From 32176338a6c809095621fd8a9e1db7168d6f1d81 Mon Sep 17 00:00:00 2001 From: Harkirat Gill Date: Thu, 27 Aug 2026 16:18:34 -0400 Subject: [PATCH 001/109] ci : build only the ggml-hip backend for windows-rocm release (#27753) --- .github/workflows/release.yml | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ab5d3389bd58..98250c860650 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -714,10 +714,10 @@ jobs: with: key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu - # TODO: build only the ggml-hip backend like the other windows backend jobs - # (windows-cuda, windows-sycl), then drop the ui-build dependency + # note: builds only the ggml-hip backend - llama-server is injected from the + # windows-cpu zip during the release "Merge artifacts" step windows-rocm: - needs: [check-release, ui-build] + needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} runs-on: windows-2022 @@ -736,12 +736,6 @@ jobs: with: fetch-depth: 0 - - name: Download UI build - uses: actions/download-artifact@v7 - with: - name: llama-ui.zip - path: tools/ui/dist - - name: Install Ninja run: | choco install ninja @@ -804,17 +798,15 @@ jobs: -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" ` -DGGML_BACKEND_DL=ON ` -DGGML_NATIVE=OFF ` - -DGGML_CPU=ON ` - -DGGML_CPU_ALL_VARIANTS=ON ` + -DGGML_CPU=OFF ` -DGGML_HIP=ON ` -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" ` -DCMAKE_C_FLAGS="-Wno-error=incompatible-pointer-types" ` -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DHIP_PATH="${env:HIP_PATH}" ` - -DGGML_HIP_ROCWMMA_FATTN=ON ` -DAMDGPU_TARGETS="${{ matrix.gpu_targets }}" - cmake --build build --config Release --parallel ${env:NUMBER_OF_PROCESSORS} + cmake --build build --config Release --parallel ${env:NUMBER_OF_PROCESSORS} --target ggml-hip - name: Verify HIP backend was built run: | @@ -866,8 +858,11 @@ jobs: - name: Pack artifacts run: | - cp "LICENSE" "build\bin\Release\" - 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip .\build\bin\Release\* + 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip ` + .\build\bin\Release\ggml-hip.dll ` + .\build\bin\Release\amdhip64_7.dll ` + .\build\bin\Release\rocm_kpack.dll ` + .\build\bin\Release\amd_comgr.dll - name: Upload artifacts uses: actions/upload-artifact@v6 From 18443257a30c884d5332abb8e7dc43c7ffe42fda Mon Sep 17 00:00:00 2001 From: Bartowski <3266127+bartowski1182@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:39:14 -0400 Subject: [PATCH 002/109] server: add ctx-per-slot (--kv-unified-per-slot) (#24124) * Add ctx-per-slot argument for unifid KV cache * Swap out ctx fractions for ctx pool slots * Formatting cleanup * Remove ctx-pool-slots, make ctx-per-slot an int * refactor it --------- Co-authored-by: Xuan Son Nguyen --- common/arg.cpp | 8 ++++++ common/common.h | 1 + tools/server/README.md | 1 + tools/server/server-context.cpp | 46 ++++++++++++++++++++++++++------- tools/server/server.cpp | 12 +++++++++ 5 files changed, 59 insertions(+), 9 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index f1e2bf690834..e346863e51fd 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1643,6 +1643,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_env("LLAMA_ARG_CTX_SIZE")); + add_opt(common_arg( + { "--kv-unified-per-slot" }, "N", + "context limit per parallel slot (default: unset, behavior unchanged).\n" + "when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N", + [](common_params & params, int value) { + params.kv_unified_per_slot = value; + } + ).set_env("LLAMA_ARG_KV_UNIFIED_PER_SLOT").set_examples({ LLAMA_EXAMPLE_SERVER })); add_opt(common_arg( {"-n", "--predict", "--n-predict"}, "N", string_format( diff --git a/common/common.h b/common/common.h index 82fed22092f6..a333f702ac1d 100644 --- a/common/common.h +++ b/common/common.h @@ -627,6 +627,7 @@ struct common_params { bool cache_prompt = true; // whether to enable prompt caching bool cache_idle_slots = true; // save and clear idle slots upon starting a new task int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot + int32_t kv_unified_per_slot = 0; // max context per parallel slot; 0 = unset int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. diff --git a/tools/server/README.md b/tools/server/README.md index 07e58fe91682..3c2228f34322 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -163,6 +163,7 @@ For the full list of features, please refer to [server's changelog](https://gith | -------- | ----------- | | `-lcs, --lookup-cache-static FNAME` | path to static lookup cache to use for lookup decoding (not updated by generation) | | `-lcd, --lookup-cache-dynamic FNAME` | path to dynamic lookup cache to use for lookup decoding (updated by generation) | +| `--kv-unified-per-slot N` | context limit per parallel slot (default: unset, behavior unchanged).
when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N
(env: LLAMA_ARG_KV_UNIFIED_PER_SLOT) | | `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) | | `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e6c991f7cfba..f5477356d61d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1208,10 +1208,31 @@ struct server_context_impl { const int n_ctx_train = llama_model_n_ctx_train(model_tgt); - int n_ctx_slot = llama_n_ctx_seq(ctx_tgt); - if (n_ctx_slot > n_ctx_train) { - SRV_WRN("the slot context (%d) exceeds the training context of the model (%d) - capping\n", n_ctx_slot, n_ctx_train); - n_ctx_slot = n_ctx_train; + { + // note: the capping itself is done in n_ctx_slot(), here we only report it + const int n_ctx_seq = llama_n_ctx_seq(ctx_tgt); + + if (params_base.kv_unified_per_slot > 0) { + if (n_ctx_seq > params_base.kv_unified_per_slot) { + SRV_INF("capping per-slot context (%d) to --kv-unified-per-slot (%d)\n", + n_ctx_seq, params_base.kv_unified_per_slot); + } else if (params_base.kv_unified_per_slot > n_ctx_seq) { + // cap is above the per-slot pool capacity, so it can never bind + SRV_WRN( + "--kv-unified-per-slot (%d) exceeds the per-slot pool capacity (%d) - cap has no effect, " + "slots are limited to %d (raise the KV pool with -c, or unset -c to size it to " + "n_parallel * kv_unified_per_slot)\n", + params_base.kv_unified_per_slot, n_ctx_seq, n_ctx_seq); + } + } + + const int n_ctx_capped = params_base.kv_unified_per_slot > 0 ? + std::min(n_ctx_seq, params_base.kv_unified_per_slot) : n_ctx_seq; + + if (n_ctx_capped > n_ctx_train) { + SRV_WRN("the slot context (%d) exceeds the training context of the model (%d) - capping\n", + n_ctx_capped, n_ctx_train); + } } slots.clear(); @@ -1227,7 +1248,7 @@ struct server_context_impl { // setup slots SRV_INF("initializing, n_slots = %d, n_ctx_slot = %d, kv_unified = '%s'\n", - params_base.n_parallel, n_ctx_slot, params_base.kv_unified ? "true" : "false"); + params_base.n_parallel, n_ctx_slot(), params_base.kv_unified ? "true" : "false"); // initialize slots for (int i = 0; i < params_base.n_parallel; i++) { @@ -1271,7 +1292,7 @@ struct server_context_impl { slot.ctx_dft = ctx_dft; slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); - slot.n_ctx = n_ctx_slot; + slot.n_ctx = n_ctx_slot(); slot.mctx = mctx; slot.prompt.tokens.has_mtmd = mctx != nullptr; @@ -3975,8 +3996,15 @@ struct server_context_impl { }); } - int get_slot_n_ctx() { - return slots.back().n_ctx; + // context size of a single slot, capped by --kv-unified-per-slot and by the training context of the model + int n_ctx_slot() const { + int res = llama_n_ctx_seq(ctx_tgt); + + if (params_base.kv_unified_per_slot > 0) { + res = std::min(res, params_base.kv_unified_per_slot); + } + + return std::min(res, llama_model_n_ctx_train(model_tgt)); } server_response_reader get_response_reader() { @@ -4142,7 +4170,7 @@ server_context_meta server_context::get_meta() const { /* has_inp_audio */ impl->chat_params.allow_audio, /* has_inp_video */ impl->chat_params.allow_video, /* json_ui_settings */ impl->json_ui_settings, - /* slot_n_ctx */ impl->get_slot_n_ctx(), + /* slot_n_ctx */ impl->n_ctx_slot(), /* pooling_type */ llama_pooling_type(impl->ctx_tgt), /* chat_params */ impl->chat_params, diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 5fe2729ba1b2..22378b38c5ef 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -157,6 +157,18 @@ int llama_server(common_params & params, int argc, char ** argv) { } } + // size the KV pool from --kv-unified-per-slot, unless the user pinned it with -c + // or with -c 0 for max context + const bool ctx_pool_auto_sized = params.kv_unified_per_slot > 0 && + params.n_ctx == 0 && + (uint32_t) params.fit_params_min_ctx != UINT32_MAX; + + if (ctx_pool_auto_sized) { + params.n_ctx = params.n_parallel * params.kv_unified_per_slot; + SRV_INF("--kv-unified-per-slot: sizing KV pool to n_parallel * kv_unified_per_slot = %d * %d = %d\n", params.n_parallel, + params.kv_unified_per_slot, params.n_ctx); + } + // for consistency between server router mode and single-model mode, we set the same model name as alias auto model_name = params.model.get_name(); if (params.model_alias.empty() && !model_name.empty()) { From 83d855c5a6d70487121edbf4020b25c96b7a04e7 Mon Sep 17 00:00:00 2001 From: Aparna M P Date: Fri, 28 Aug 2026 03:08:02 +0530 Subject: [PATCH 003/109] hex-unary: fix RMS_NORM_MUL weight-offset bugs for grouped/broadcast norms (#27798) --- ggml/src/ggml-hexagon/htp/unary-ops.c | 32 +++++++++++++++++---------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c index b21415a67d64..971f3b537778 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.c +++ b/ggml/src/ggml-hexagon/htp/unary-ops.c @@ -478,6 +478,9 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat const uint32_t nb11 = src1 ? src1->nb[1] : 0; \ const uint32_t nb12 = src1 ? src1->nb[2] : 0; \ const uint32_t nb13 = src1 ? src1->nb[3] : 0; \ + const uint32_t nb11_bc = (src1 && src1->ne[1] > 1) ? nb11 : 0; \ + const uint32_t nb12_bc = (src1 && src1->ne[2] > 1) ? nb12 : 0; \ + const uint32_t nb13_bc = (src1 && src1->ne[3] > 1) ? nb13 : 0; \ const bool src1_contig = src1 ? ((nb12 == (size_t)ne01 * nb11) && (nb13 == (size_t)ne02 * nb12)) : false; \ \ uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); \ @@ -497,8 +500,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; \ const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; \ \ - const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \ - const uint32_t dst_max_block = dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \ + const bool src1_needs_row_clip = (IS_RMS_NORM_MUL) && !uctx->broadcast_weight && !src1_contig; \ + const bool block_src0_contig = src0_contig && !src1_needs_row_clip; \ + const bool block_dst_contig = dst_contig && !src1_needs_row_clip; \ + \ + const uint32_t src0_max_block = block_src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \ + const uint32_t dst_max_block = block_dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \ const uint32_t BLOCK = MIN(src0_max_block, dst_max_block); \ if (BLOCK == 0) { \ FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small, needed at least %zu\n", \ @@ -515,8 +522,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat } \ \ for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { \ - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \ - div_ne01); \ + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \ + ne01, div_ne01); \ \ dma_queue_push(dma_queue, \ dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), \ @@ -530,7 +537,7 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ const size_t src1_off = src1_contig ? (ir * nb11) : \ - unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \ + unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, nb13_bc); \ dma_queue_push(dma_queue, \ dma_make_ptr(src1_vtcm_data + (vtcm_idx * src1_vtcm_half_size), data_src1 + src1_off), \ uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, block_size); \ @@ -540,8 +547,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat } \ \ for (uint32_t ir = src0_start_row; ir < src0_end_row; ) { \ - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \ - div_ne01); \ + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \ + ne01, div_ne01); \ \ float * dst_vtcm = (float *) dma_queue_pop(dma_queue).src; \ float * src0_vtcm = (float *) dma_queue_pop(dma_queue).dst; \ @@ -562,12 +569,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ const uint32_t next_ir = ir + block_size; \ if (next_ir < src0_end_row) { \ - const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig,\ - ne01, div_ne01); \ + const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, block_src0_contig, \ + block_dst_contig, ne01, div_ne01); \ const uint32_t pref_ir = next_ir + next_block_size; \ if (pref_ir < src0_end_row) { \ - const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig, \ - dst_contig, ne01, div_ne01); \ + const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, block_src0_contig, \ + block_dst_contig, ne01, div_ne01); \ const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : \ unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); \ dma_queue_push(dma_queue, \ @@ -576,7 +583,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ const size_t src1_pref_off = src1_contig ? (pref_ir * nb11) : \ - unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \ + unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, \ + nb13_bc); \ dma_queue_push(dma_queue, \ dma_make_ptr(src1_vtcm, data_src1 + src1_pref_off), \ uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size); \ From e70802a01f03f0ed31a26338a5664796f3824371 Mon Sep 17 00:00:00 2001 From: cqderek Date: Fri, 28 Aug 2026 06:05:57 +0800 Subject: [PATCH 004/109] ggml-hexagon: add HTP unary ops for ABS and LOG (#27786) Add HVX-accelerated implementations for GGML_OP_LOG and GGML_UNARY_OP_ABS on the HTP backend. - Register HTP_OP_UNARY_ABS and HTP_OP_UNARY_LOG in op_remap_to_htp() - Add ABS and LOG to ggml_backend_hexagon_device_supports_op() - Implement hvx_abs_f32_aa() in hvx-arith.h using hvx_vec_abs_f32() - Implement hvx_log_f32_aa() in hvx-log.h using hvx_vec_log_f32() - Add abs_f32() and log_f32() row-wise dispatch in unary-ops.c - Define tiled and non-tiled task functions via DEFINE_UNARY_TASK and DEFINE_UNARY_TILED_TASK macros - Route HTP_OP_UNARY_ABS and HTP_OP_UNARY_LOG through execute_op() in main.c --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 4 +++ ggml/src/ggml-hexagon/htp/htp-ops.h | 2 ++ ggml/src/ggml-hexagon/htp/hvx-arith.h | 28 +++++++++++++++++++ ggml/src/ggml-hexagon/htp/hvx-log.h | 24 ++++++++++++++++ ggml/src/ggml-hexagon/htp/main.c | 2 ++ ggml/src/ggml-hexagon/htp/unary-ops.c | 38 ++++++++++++++++++++++++++ ggml/src/ggml-hexagon/htp/unary-ops.h | 2 ++ 7 files changed, 100 insertions(+) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index c1e9f919d92a..53e860755910 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4643,6 +4643,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_OP_CLAMP: return HTP_OP_CLAMP; case GGML_OP_SQR: return HTP_OP_SQR; case GGML_OP_SQRT: return HTP_OP_SQRT; + case GGML_OP_LOG: return HTP_OP_UNARY_LOG; case GGML_OP_SOFT_MAX: return HTP_OP_SOFTMAX; case GGML_OP_SSM_CONV: return HTP_OP_SSM_CONV; case GGML_OP_GATED_DELTA_NET: return HTP_OP_GATED_DELTA_NET; @@ -4666,6 +4667,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_UNARY_OP_EXP: return HTP_OP_UNARY_EXP; case GGML_UNARY_OP_SOFTPLUS: return HTP_OP_UNARY_SOFTPLUS; case GGML_UNARY_OP_TANH: return HTP_OP_UNARY_TANH; + case GGML_UNARY_OP_ABS: return HTP_OP_UNARY_ABS; default: break; } @@ -5463,6 +5465,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_SQR: case GGML_OP_SQRT: + case GGML_OP_LOG: supp = ggml_hexagon_supported_unary(sess, op); break; @@ -5481,6 +5484,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons case GGML_UNARY_OP_SIGMOID: case GGML_UNARY_OP_SOFTPLUS: case GGML_UNARY_OP_TANH: + case GGML_UNARY_OP_ABS: case GGML_UNARY_OP_SILU: case GGML_UNARY_OP_GELU: case GGML_UNARY_OP_GELU_QUICK: diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index b4023b34d388..e804844d5996 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -62,6 +62,8 @@ enum htp_op_code { HTP_OP_UNARY_NEG, HTP_OP_UNARY_SOFTPLUS, HTP_OP_UNARY_TANH, + HTP_OP_UNARY_ABS, + HTP_OP_UNARY_LOG, HTP_OP_GLU_SWIGLU, HTP_OP_GLU_SWIGLU_OAI, HTP_OP_GLU_GEGLU, diff --git a/ggml/src/ggml-hexagon/htp/hvx-arith.h b/ggml/src/ggml-hexagon/htp/hvx-arith.h index c8d0003ab5c0..765c3577668e 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-arith.h +++ b/ggml/src/ggml-hexagon/htp/hvx-arith.h @@ -358,6 +358,34 @@ static inline void hvx_clamp_scalar_f32(uint8_t * restrict dst, const uint8_t * } } +// +// Abs +// + +static inline void hvx_abs_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + + HVX_Vector * restrict vdst = (HVX_Vector *) dst; + HVX_Vector * restrict vsrc = (HVX_Vector *) src; + + const uint32_t elem_size = sizeof(float); + const uint32_t epv = 128 / elem_size; + const uint32_t nvec = n / epv; + const uint32_t nloe = n % epv; + + uint32_t i = 0; + + _Pragma("unroll(4)") + for (; i < nvec; i++) { + vdst[i] = hvx_vec_abs_f32(vsrc[i]); + } + if (nloe) { + HVX_Vector v = hvx_vec_abs_f32(vsrc[i]); + hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v); + } +} + // // Square // diff --git a/ggml/src/ggml-hexagon/htp/hvx-log.h b/ggml/src/ggml-hexagon/htp/hvx-log.h index 7013dae785ac..a209f88d555c 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-log.h +++ b/ggml/src/ggml-hexagon/htp/hvx-log.h @@ -62,4 +62,28 @@ static inline HVX_Vector hvx_vec_log_f32(HVX_Vector x) { return hvx_vec_add_f32_f32(term_e, res); } +static inline void hvx_log_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + + HVX_Vector * restrict vdst = (HVX_Vector *) dst; + HVX_Vector * restrict vsrc = (HVX_Vector *) src; + + const uint32_t elem_size = sizeof(float); + const uint32_t epv = 128 / elem_size; + const uint32_t nvec = n / epv; + const uint32_t nloe = n % epv; + + uint32_t i = 0; + + _Pragma("unroll(4)") + for (; i < nvec; i++) { + vdst[i] = hvx_vec_log_f32(vsrc[i]); + } + if (nloe) { + HVX_Vector v = hvx_vec_log_f32(vsrc[i]); + hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v); + } +} + #endif /* HVX_LOG_H */ diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index 975ba0c7af51..fe7d093a81c1 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -777,6 +777,8 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_UNARY_NEG: case HTP_OP_UNARY_EXP: case HTP_OP_UNARY_TANH: + case HTP_OP_UNARY_ABS: + case HTP_OP_UNARY_LOG: case HTP_OP_L2_NORM: return op_unary(octx); diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c index 971f3b537778..1a632bf5631e 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.c +++ b/ggml/src/ggml-hexagon/htp/unary-ops.c @@ -443,6 +443,34 @@ static void tanh_f32(const float * restrict src, } } +static void abs_f32(const float * restrict src, + float * restrict dst, + const uint32_t num_rows, + const struct htp_unary_context * uctx) { + htp_unary_op_preamble; + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned); + + hvx_abs_f32_aa(dst_local, src_local, ne0); + } +} + +static void log_f32(const float * restrict src, + float * restrict dst, + const uint32_t num_rows, + const struct htp_unary_context * uctx) { + htp_unary_op_preamble; + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned); + + hvx_log_f32_aa(dst_local, src_local, ne0); + } +} + #define DEFINE_UNARY_TASK(NAME, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR) \ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * data) { \ const struct htp_unary_context * uctx = (const struct htp_unary_context *) data; \ @@ -611,6 +639,8 @@ DEFINE_UNARY_TASK(unary_silu, false, false, silu_f32(src0_vtcm, dst_vtcm, bl DEFINE_UNARY_TASK(unary_gelu, false, false, gelu_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(unary_softplus, false, false, softplus_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(unary_tanh, false, false, tanh_f32(src0_vtcm, dst_vtcm, block_size, uctx)) +DEFINE_UNARY_TASK(unary_abs, false, false, abs_f32(src0_vtcm, dst_vtcm, block_size, uctx)) +DEFINE_UNARY_TASK(unary_log, false, false, log_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(l2_norm, false, false, l2_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(tri, false, true, tri_f32(src0_vtcm, dst_vtcm, block_size, ir, uctx)) @@ -858,6 +888,8 @@ DEFINE_UNARY_TILED_TASK(unary_silu, false, tile_silu_f32(dst_vtcm, src_vtcm, DEFINE_UNARY_TILED_TASK(unary_gelu, false, tile_gelu_f32(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(unary_softplus, false, tile_unary_softplus_f32(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(unary_tanh, false, hvx_tanh_f32_aa(dst_vtcm, src_vtcm, tw)) +DEFINE_UNARY_TILED_TASK(unary_abs, false, hvx_abs_f32_aa(dst_vtcm, src_vtcm, tw)) +DEFINE_UNARY_TILED_TASK(unary_log, false, hvx_log_f32_aa(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(tri, true, tri_apply_tile_f32(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype)) static int execute_op_unary_f32(struct htp_ops_context * octx) { @@ -883,6 +915,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: op_type = "gelu-f32"; break; case HTP_OP_UNARY_SOFTPLUS: op_type = "softplus-f32"; break; case HTP_OP_UNARY_TANH: op_type = "tanh-f32"; break; + case HTP_OP_UNARY_ABS: op_type = "abs-f32"; break; + case HTP_OP_UNARY_LOG: op_type = "log-f32"; break; case HTP_OP_L2_NORM: op_type = "l2norm-f32"; break; case HTP_OP_TRI: op_type = "tri-f32"; break; @@ -981,6 +1015,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: task_func = unary_task_f32_tiled_unary_gelu; break; case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_tiled_unary_softplus; break; case HTP_OP_UNARY_TANH: task_func = unary_task_f32_tiled_unary_tanh; break; + case HTP_OP_UNARY_ABS: task_func = unary_task_f32_tiled_unary_abs; break; + case HTP_OP_UNARY_LOG: task_func = unary_task_f32_tiled_unary_log; break; case HTP_OP_TRI: task_func = unary_task_f32_tiled_tri; break; default: break; } @@ -1000,6 +1036,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: task_func = unary_task_f32_unary_gelu; break; case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_unary_softplus; break; case HTP_OP_UNARY_TANH: task_func = unary_task_f32_unary_tanh; break; + case HTP_OP_UNARY_ABS: task_func = unary_task_f32_unary_abs; break; + case HTP_OP_UNARY_LOG: task_func = unary_task_f32_unary_log; break; case HTP_OP_L2_NORM: task_func = unary_task_f32_l2_norm; break; case HTP_OP_TRI: task_func = unary_task_f32_tri; break; default: break; diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.h b/ggml/src/ggml-hexagon/htp/unary-ops.h index 1f4c3a5c4d96..458218ff4431 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.h +++ b/ggml/src/ggml-hexagon/htp/unary-ops.h @@ -55,6 +55,8 @@ static inline bool htp_op_is_unary(uint32_t opcode) { case HTP_OP_UNARY_GELU: case HTP_OP_UNARY_SOFTPLUS: case HTP_OP_UNARY_TANH: + case HTP_OP_UNARY_ABS: + case HTP_OP_UNARY_LOG: case HTP_OP_L2_NORM: case HTP_OP_TRI: return true; From ca3d5a3e10d53f7ea672cb9b6178faca3e2807bc Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Fri, 28 Aug 2026 01:49:27 +0200 Subject: [PATCH 005/109] model: add DSpark support for Nemotron3.5 (#27804) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * model: add DSpark support for Nemotron3.5 * Update src/models/dflash.cpp Co-authored-by: Sigbjørn Skjæret --------- Co-authored-by: Sigbjørn Skjæret Co-authored-by: Xuan Son Nguyen --- common/speculative.cpp | 18 +++++++++++- conversion/qwen.py | 18 ++++++++++-- gguf-py/gguf/constants.py | 1 + gguf-py/gguf/gguf_writer.py | 3 ++ src/llama-model.h | 1 + src/models/dflash.cpp | 58 ++++++++++++++++++++++--------------- 6 files changed, 71 insertions(+), 28 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index b5348ab6f3d5..d34d1c9c5950 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -935,6 +935,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // dspark speculators bool sample_from_anchor = true; + // block-internal attention + bool causal_attn = false; + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; @@ -972,12 +975,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) { sample_from_anchor = std::strcmp(buf, "true") == 0; } + if (llama_model_meta_val_str(model_dft, "dflash.attention.causal", buf, sizeof(buf)) >= 0) { + causal_attn = std::strcmp(buf, "true") == 0; + } } selector_top_k = llama_model_dflash_selector_top_k(model_dft); is_dflash2 = selector_top_k > 0; mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); + if (is_dspark && this->params.p_min > 0.0f) { + char buf[16] = {}; + const bool has_conf = + llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 || + std::strcmp(buf, "true") == 0; + if (!has_conf) { + throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0"); + } + } + LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__, @@ -1036,7 +1052,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits. llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2); - llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention + llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise } ~common_speculative_impl_draft_dflash() override { diff --git a/conversion/qwen.py b/conversion/qwen.py index c5297418c99b..419611896fc3 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -709,14 +709,20 @@ def set_gguf_parameters(self): extract_layer_ids = [i + 1 for i in target_layer_ids] self.gguf_writer.add_target_layers(extract_layer_ids) - use_sliding_window = self.hparams.get("use_sliding_window", False) - sliding_window = self.hparams.get("sliding_window") + use_sliding_window = self.hparams.get("use_sliding_window", False) or dflash_config.get("use_swa", False) + sliding_window = dflash_config.get("swa_window_size") or self.hparams.get("sliding_window") layer_types = self.hparams.get("layer_types") if use_sliding_window and sliding_window and layer_types: is_swa = [lt == "sliding_attention" for lt in layer_types] self.gguf_writer.add_sliding_window(sliding_window) self.gguf_writer.add_sliding_window_pattern(is_swa) + causal = self.hparams.get("is_causal") + if causal is None: + causal = dflash_config.get("causal") + if causal is not None: + self.gguf_writer.add_causal_attention(bool(causal)) + # M-RoPE target: the draft ropes on the temporal dim only, so write # degenerate sections [n_rot/2, 0, 0, 0] if self._target_uses_mrope(): @@ -737,6 +743,8 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca name, gen = item if not name.startswith("model."): name = "model." + name + if "sink" in name and not name.endswith(".weight"): + name += ".weight" return super().filter_tensors((name, gen)) _ROPE_PERMUTE_SUFFIXES = ( @@ -815,6 +823,10 @@ def set_gguf_parameters(self): super().set_gguf_parameters() self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor) + # confidence head is optional: vanilla-markov exports ship without it + has_conf = any("confidence_head.proj" in name for name in self.model_tensors) + self.gguf_writer.add_has_confidence_head(has_conf) + @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: if item[0] == "t2d": # not used at runtime @@ -833,7 +845,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter self._d2t = data_torch return - if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")): + if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith("lm_head.weight"): return # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index fffbd6745397..c99feb3c795c 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -167,6 +167,7 @@ class LLM: SELECTOR_RANK = "{arch}.selector_rank" SELECTOR_TOP_K = "{arch}.selector_top_k" SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" + HAS_CONFIDENCE_HEAD = "{arch}.has_confidence_head" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index b1d161bcb37e..1f309ad2eafd 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1008,6 +1008,9 @@ def add_selector_top_k(self, value: int) -> None: def add_sample_from_anchor(self, value: bool) -> None: self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value) + def add_has_confidence_head(self, value: bool) -> None: + self.add_bool(Keys.LLM.HAS_CONFIDENCE_HEAD.format(arch=self.arch), value) + def add_target_layers(self, value: Sequence[int]) -> None: self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value) diff --git a/src/llama-model.h b/src/llama-model.h index 0d7352ac500b..38066538ed10 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -672,6 +672,7 @@ struct llama_model { // dspark struct ggml_tensor * dspark_markov_w1 = nullptr; struct ggml_tensor * dspark_markov_w2 = nullptr; + struct ggml_tensor * dspark_markov_w2_s = nullptr; struct ggml_tensor * dspark_conf_proj = nullptr; struct ggml_tensor * dspark_conf_proj_b = nullptr; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 516651ce4074..f2c7d1d2462e 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -115,10 +115,11 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { if (markov_meta) { const int64_t dspark_markov_rank = markov_meta->ne[0]; - dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); - dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); + dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); + dspark_markov_w2_s = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "scale"), { 1 }, TENSOR_NOT_REQUIRED); - dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0); + dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED); dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank); @@ -219,6 +220,9 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0); + // optional per-head attention sinks (e.g. Nemotron DSpark) + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), { n_head }, TENSOR_NOT_REQUIRED); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0); layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); @@ -290,7 +294,10 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * w1 = model.dspark_markov_w1; ggml_tensor * w2 = model.dspark_markov_w2; - GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded"); + GGML_ASSERT(w1 && w2 && "DSpark markov weights not loaded"); + + // confidence head is optional + const bool has_conf = model.dspark_conf_proj != nullptr; ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens] const int64_t n_vocab = base->ne[0]; @@ -321,23 +328,22 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0); prev = ggml_cont_1d(ctx0, prev, n_blocks); - // confidence head input: predicts per-position acceptance - ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok] - ggml_tensor * cat = nullptr; ggml_tensor * cat_conf = nullptr; if (!sample_from_anchor) { // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column - cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); - cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); + if (has_conf) { + cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + } } // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final // token pick, not the Markov conditioning path for (int64_t i = i_draft_beg; i < block_drafts; ++i) { - ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] - ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks] + ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] + ggml_tensor * bias = g.build_lora_mm(w2, w1_prev, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks] if (model.d2t) { // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others) const int64_t n_draft_vocab = bias->ne[0]; @@ -354,17 +360,21 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; - // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] - ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks, - (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]); - ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0); - ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); - if (model.dspark_conf_proj_b) { - conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b); - } - conf = ggml_sigmoid(ctx0, conf); + if (has_conf) { + // confidence head input: predicts per-position acceptance + ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok] + // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] + ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks, + (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]); + ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0); + ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); + if (model.dspark_conf_proj_b) { + conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b); + } + conf = ggml_sigmoid(ctx0, conf); - cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; + cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; + } if (i + 1 < block_drafts) { prev = ggml_argmax(ctx0, col); @@ -376,7 +386,7 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks] out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); - { + if (has_conf) { ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts); conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3)); conf = ggml_reshape_2d(ctx0, conf, 1, n_tok); @@ -707,8 +717,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra // cache-aware, non-causal attention ggml_tensor * cur = use_iswa - ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) - : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il) + : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il); if (attn_dynamic) { cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1); From 4e97ac86ebe2c4cb8212d98d2641ad6768810896 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 28 Aug 2026 09:45:19 +0300 Subject: [PATCH 006/109] tests : run test-save-load-state across all architectures (#27755) * tests : run test-save-load-state across all architectures test-save-load-state previously only ran in ctest against a single downloaded model (tinyllamas/stories15M), i.e. only the llama arch. Add a --models DIR mode to test-save-load-state that runs the full save/load suite over every *.gguf in a directory, reporting a per-model PASS/FAIL and exiting non-zero if any model fails, and wire a ctest to run it over all architectures using the existing generate-models fixture (test-llama-archs). The single-model -m mode is preserved (still used by ci/run.sh). Also bump the dummy-model training context in test-llama-archs from 128 to 256 so that the per-sequence context (which is padded up to a multiple of 256) no longer exceeds n_ctx_train and emits the "possible training context overflow" warning. The test is expected to fail until the affected arches are fixed: deepseek4 (host seq-copy), gemma2/gpt-oss/lfm2 (device seq-copy), minimax-01 (state load). It aborts at the first arch that crashes. Assisted-by: pi:llama.cpp/Qwen3.8-27B * tests : match dummy DSA indexer to fused Lightning Indexer kernel The dummy DSA indexer (deepseek32, glm-dsa, ...) used key_length=64 and head_count=1, so the fused Lightning Indexer op's q tensor was shaped [64, 1, ...]. The Metal fused kernel is fixed to DK=128, NH=64, so it rejected the op and the scheduler fell back to CPU, emitting a 'layer assigned to MTL but Lightning Indexer on CPU' warning. Bump key_length to 128 and the DSA head_count to 64 so the fused op runs on the GPU. Assisted-by: pi:llama.cpp/Qwen3.8-27B * tests : add --help and document -o in test-llama-archs Add a --help/-h flag to test-llama-archs and list the existing -o/--out option in the usage text, which was previously missing. Assisted-by: pi:llama.cpp/Qwen3.8-27B * tests : use 64 indexer heads for deepseek4 deepseek4's indexer head count was set to n_head (8), which does not match the fused Lightning Indexer kernel's fixed NH=64, so the fused op fell back to the CPU backend and emitted a device-mismatch warning. Give it the same fixed 64 as the other indexer archs by dropping it from the n_head ternary (only minimax-m3 keeps n_head, since it does not use the fused Lightning Indexer op). Assisted-by: pi:llama.cpp/Qwen3.8-27B * tests : fix dsv4 save-load n_stream mismatch The dsv4 KV cache keeps per-sequence KV/state streams even in unified mode, so its n_stream equals n_seq_max. The test saved the state in the baseline with n_seq_max=1 but loaded it in the seq-copy tests with n_seq_max=2, so state_read threw an n_stream mismatch. Use n_seq_max=2 in the baseline and state-load tests so the save and load agree. Assisted-by: pi:llama.cpp/Qwen3.8-27B * context : relax on-device seq-copy chunk alignment The on-device state seq copy (llama_state_seq_set_data with LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) copied the write-side cpy tensors to the read-side targets 1:1 by index, requiring the writer and reader to emit the same number of chunks in the same order with the same per-chunk sizes. state_write_data chunks per cell-range while state_read_data chunks contiguous-or-per-cell, so the counts diverged for non-contiguous sources (dsv4, SWA) and the copy aborted with "memory buffer mismatch". All state writers and readers enumerate the same logical data in the same order, differing only in chunking. Copy the flat write-side data into the read-side targets with a byte cursor that walks both tensor lists across their boundaries, so the chunking no longer needs to match. Keep the total-size guard; drop the n_tensors equality check. Assisted-by: pi:llama.cpp/Qwen3.8-27B * model : fix dangling hparams ref in minimax-01 LA graph input llm_graph_input_la stored const llama_hparams & hparams, bound to the llm_graph_params temporary in llama_context::process_ubatch. The input object outlives that temporary (it is kept in llm_graph_result::inputs for graph reuse), so set_input() read destroyed stack memory on every graph reuse - test-save-load-state crashed for minimax-01 when the stack region was overwritten (n_layer_all read as 0, abort in llama_hparams::n_head). Store a copy like every other graph input class. Assisted-by: pi:llama.cpp/Qwen3.8-27B * context : handle "worst case" graph and add TODO --- src/llama-context.cpp | 94 ++++++++++++++++++-- src/models/minimax-01.cpp | 2 +- tests/CMakeLists.txt | 13 ++- tests/test-llama-archs.cpp | 14 ++- tests/test-save-load-state.cpp | 155 +++++++++++++++++++++++++-------- 5 files changed, 226 insertions(+), 52 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index fb88919f9d67..179c526c2940 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -661,11 +661,19 @@ void llama_context::sched_reserve() { // reserve again with pp graph to avoid ggml-alloc reallocations during inference { - // TODO: not sure if the following graph would be worst case for multi-stream KV caches: - // - // auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get()); - // - auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + // TODO: the worst case graph is not always reached for `n_seqs > 1` + // need to implement a more robust mechanism that tries a few different inputs and analyzes the results + ggml_cgraph * gf = nullptr; + switch (model.arch) { + case LLM_ARCH_MINIMAX_01: + // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which + // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1` + gf = graph_reserve(n_tokens, 1, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + break; + default: + gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + }; + if (!gf) { throw std::runtime_error("failed to allocate compute pp buffers"); } @@ -2892,13 +2900,83 @@ class llama_io_read_device : public llama_io_read_i { for (auto & [buft, mbuf] : mbufs_new) { const auto & mbuf_cur = mbufs.at(buft); - if (!mbuf_cur.buf || mbuf_cur.n_tensors != mbuf.n_tensors || mbuf_cur.total_size != mbuf.total_size) { + if (!mbuf_cur.buf || mbuf_cur.total_size != mbuf.total_size) { GGML_ABORT("%s: memory buffer mismatch\n", __func__); } - for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { - ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + if (mbuf_cur.n_tensors == mbuf.n_tensors) { + // same chunking: copy 1:1 by index + for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i])); + ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + } + continue; } + + // different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org) + // with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk + // it differently, so copy across tensor boundaries rather than 1:1 by index. + const size_t total = mbuf_cur.total_size; + + ggml_init_params params_scratch = { + /*.mem_size =*/ 2*(mbuf_cur.cpy.size() + mbuf.org.size())*ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context * ctx_scratch = ggml_init(params_scratch); + + size_t src_pos = 0; + size_t dst_pos = 0; + size_t src_j = 0; + size_t dst_i = 0; + size_t src_base = 0; + size_t dst_base = 0; + + while (src_pos < total) { + const auto & src_t = mbuf_cur.cpy[src_j]; + const auto & dst_t = mbuf.org[dst_i]; + + const size_t src_size = ggml_nbytes(src_t); + const size_t dst_size = ggml_nbytes(dst_t); + + const size_t src_off = src_pos - src_base; + const size_t dst_off = dst_pos - dst_base; + + const size_t n_copy = std::min(src_size - src_off, dst_size - dst_off); + + const size_t el = ggml_element_size(src_t); + const int64_t n_el = (int64_t) (n_copy / el); + + auto * src_v = ggml_view_1d(ctx_scratch, src_t, n_el, src_off); + ggml_backend_view_init(src_v); + auto * dst_v = ggml_view_1d(ctx_scratch, dst_t, n_el, dst_off); + ggml_backend_view_init(dst_v); + + ggml_backend_tensor_copy(src_v, dst_v); + + src_pos += n_copy; + dst_pos += n_copy; + + if (src_pos - src_base == src_size) { + src_base = src_pos; + ++src_j; + } + if (dst_pos - dst_base == dst_size) { + dst_base = dst_pos; + ++dst_i; + } + } + + GGML_ASSERT(src_pos == total && dst_pos == total); + // any tensors left unvisited hold no data + for (size_t i = src_j; i < mbuf_cur.cpy.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == 0); + } + for (size_t i = dst_i; i < mbuf.org.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf.org[i]) == 0); + } + + ggml_free(ctx_scratch); } GGML_ASSERT(buf_size == 0); diff --git a/src/models/minimax-01.cpp b/src/models/minimax-01.cpp index f14626b2c74f..361114acc327 100644 --- a/src/models/minimax-01.cpp +++ b/src/models/minimax-01.cpp @@ -181,7 +181,7 @@ class llm_graph_input_la : public llm_graph_input_i { return res; } - const llama_hparams & hparams; + const llama_hparams hparams; ggml_tensor * inp_slopes = nullptr; // F32 [n_head] ggml_tensor * inp_q_decay = nullptr; // F32 [1, n_head, n_batch] diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b9f9d4b78af2..fe3d14ffc552 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -149,6 +149,7 @@ if (LLAMA_LLGUIDANCE) endif () llama_build(test-recurrent-state-rollback.cpp) +llama_build(test-save-load-state.cpp) if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) @@ -237,6 +238,14 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) set_tests_properties(test-recurrent-state-rollback-dsv4 PROPERTIES FIXTURES_REQUIRED generate-models ) + + # Test state save/load functionality across all architectures, using the generated dummy models + llama_test( + test-save-load-state + LABEL main + ARGS --models "${MODEL_DIR}" + ) + set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED generate-models) endif() llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) @@ -299,10 +308,6 @@ llama_build_and_test(test-backend-sampler.cpp LABEL "model") llama_build_and_test(test-state-restore-fragmented.cpp LABEL "model" ARGS -m "${MODEL_DEST}") set_tests_properties(test-state-restore-fragmented PROPERTIES FIXTURES_REQUIRED test-download-model) -# Test state save/load functionality -llama_build_and_test(test-save-load-state.cpp LABEL "model" ARGS -m "${MODEL_DEST}") -set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED test-download-model) - if (APPLE) llama_build(test-rset-release.cpp) endif() diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d58d90952eb0..35a3286e4a1b 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -65,7 +65,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } static void usage(char ** argv) { - printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-v/--verbose]\n", argv[0]); + printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v/--verbose] [-h/--help]\n", argv[0]); } static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ @@ -82,7 +82,7 @@ static std::vector get_tokens(const uint32_t n_tokens, const uint32 static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { gguf_context_ptr ret(gguf_init_empty()); llama_model_saver ms(arch, ret.get()); - const uint32_t n_ctx = 128; + const uint32_t n_ctx = 256; uint32_t n_vocab = 128; uint32_t n_embd = 256; @@ -256,10 +256,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector(n_layer, 4)); } - ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_DEEPSEEK4 ? n_head : uint32_t(1)); + // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused + ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(64)); // qwen4exp ropes indexer keys with the main rotary width, so its head can't be < n_rot ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, - arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(64)); + arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(128)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); @@ -762,6 +764,10 @@ int main(int argc, char ** argv) { std::string out; for (int i = 1; i < argc; i++) { + if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) { + usage(argv); + return 0; + } if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) { if (i + 1 < argc) { const std::string arch_name = argv[++i]; diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 6e93ce6fb8da..0ceab7c5452b 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -3,8 +3,12 @@ #include "log.h" #include "llama-cpp.h" +#include #include +#include +#include #include +#include #include struct llama_batch_ptr { @@ -53,7 +57,9 @@ static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, i // - decode the last token // - generate n_predict tokens static llama_tokens test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) { - auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))}; + auto params_ctx = common_context_params_to_llama(params); + params_ctx.n_seq_max = 2; + auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; auto sparams = llama_sampler_chain_default_params(); auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)}; @@ -161,7 +167,9 @@ static bool test_seq_rm_isolated( // - replay the last prompt token // - generate n_predict tokens and compare against expected result static bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) { - auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))}; + auto params_ctx = common_context_params_to_llama(params); + params_ctx.n_seq_max = 2; + auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; auto sparams = llama_sampler_chain_default_params(); auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)}; @@ -347,38 +355,18 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p } -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - params.prompt = ""; - params.n_batch = 100; - params.out_file = "dump_state.bin"; - params.sampling.seed = 1234; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { - return 1; - } - - if (params.n_parallel == 1) { - LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__); - params.kv_unified = true; - } - - if (params.n_predict < 0) { - params.n_predict = 16; - } - - ggml_backend_load_all(); +// Run the full save/load test suite (tests 1-5) for a single model. +// Returns true if all tests pass, false otherwise. +static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) { + struct common_params params = base_params; + params.model.path = model_path; auto llama_init = common_init_from_params(params, true); auto * model = llama_init->model(); if (model == nullptr) { - LOG_ERR("%s: failed to init\n", __func__); - return 1; + LOG_ERR("%s: failed to init model '%s'\n", __func__, model_path.c_str()); + return false; } GGML_ASSERT(llama_init->context() == nullptr); @@ -411,30 +399,127 @@ int main(int argc, char ** argv) { // Test 1: baseline (saves state to disk) auto result_baseline = test_baseline(model, params, tokens); if (result_baseline.empty()) { - return 1; + return false; } // Test 2: sequence removal isolation if (!test_seq_rm_isolated(model, params, tokens)) { - return 1; + return false; } // Test 3: state load if (!test_state_load(model, params, tokens, result_baseline)) { - return 1; + return false; } // Test 4: seq copy (host) if (!test_seq_cp_host(model, params, tokens, result_baseline)) { - return 1; + return false; } // Test 5: seq copy (device) if (!test_seq_cp_device(model, params, tokens, result_baseline)) { - return 1; + return false; } LOG("\nAll tests passed.\n"); - return 0; + return true; +} + + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + + common_params params; + params.prompt = ""; + params.n_batch = 100; + params.out_file = "dump_state.bin"; + params.sampling.seed = 1234; + + common_init(); + + // extract our own --models DIR option before handing the rest to the common arg parser + std::string models_dir; + std::vector filtered_argv; + filtered_argv.push_back(argv[0]); + for (int i = 1; i < argc; i++) { + if (strcmp(argv[i], "--models") == 0) { + if (i + 1 >= argc) { + LOG_ERR("%s: --models requires a directory argument\n", __func__); + return 1; + } + models_dir = argv[i + 1]; + i++; + } else { + filtered_argv.push_back(argv[i]); + } + } + filtered_argv.push_back(nullptr); + const int fargc = (int)filtered_argv.size() - 1; + + // in --models mode there is no single model; set a placeholder so the common parser's + // "--model is required" check passes (each model is set individually inside the loop) + if (!models_dir.empty()) { + params.model.path = models_dir; + } + + if (!common_params_parse(fargc, filtered_argv.data(), params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + if (params.n_parallel == 1) { + LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__); + params.kv_unified = true; + } + + if (params.n_predict < 0) { + params.n_predict = 16; + } + + ggml_backend_load_all(); + + if (!models_dir.empty()) { + // run the suite over every dummy model in the directory + if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) { + LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str()); + return 1; + } + + std::vector models; + for (const auto & entry : std::filesystem::directory_iterator(models_dir)) { + if (entry.is_regular_file() && entry.path().extension() == ".gguf") { + models.push_back(entry.path().string()); + } + } + std::sort(models.begin(), models.end()); + + if (models.empty()) { + LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str()); + return 1; + } + + LOG_INF("%s: running save/load tests over %zu models in '%s'\n", __func__, models.size(), models_dir.c_str()); + + size_t n_pass = 0; + size_t n_fail = 0; + for (const auto & model_path : models) { + LOG("\n================================================================\n"); + LOG_INF("%s: model %s\n", __func__, model_path.c_str()); + + if (run_save_load_tests_for_model(model_path, params)) { + n_pass++; + } else { + n_fail++; + } + } + + LOG("\n================================================================\n"); + LOG_INF("%s: summary: %zu passed, %zu failed (of %zu)\n", __func__, n_pass, n_fail, models.size()); + + return n_fail == 0 ? 0 : 1; + } + + // single-model mode + return run_save_load_tests_for_model(params.model.path, params) ? 0 : 1; } From 6d6b697cd53885de8f6f4e1b80902e3559817d7b Mon Sep 17 00:00:00 2001 From: Brad Smith <1472326+infinitewarp@users.noreply.github.com> Date: Fri, 28 Aug 2026 04:37:43 -0400 Subject: [PATCH 007/109] metal : add fa-vec tunings for M4 Pro (#27824) This is a followup contribution to efeda76b948f59ee52ea20db640bc4cf3dfe8ac1 as requested in https://github.com/ggml-org/llama.cpp/discussions/27668 to add support for additional Apple GPUs. I generated this output using the provided instructions: ```sh git clone https://github.com/ggml-org/llama.cpp cd llama.cpp cmake -B build -DGGML_METAL=ON cmake --build build --target ggml-metal-tuning -j ./build/bin/ggml-metal-tuning fa-vec --dtype f16,q8_0 > fa_vec_rows.txt 2> fa_vec_sweep.log ``` This ran on a MacBook Pro (14-inch, Nov 2024) with Apple M4 Pro. The `ggml-metal-tuning` command completed successfully in 1h 13m 1s with no other notable load on the system. --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 61 ++++++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 6d8c18e6a6a2..7e99c1dd6612 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -448,7 +448,66 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, - + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 128, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, From 8963a9bdcdf312abc9aab2e662e525c06c9964ec Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 28 Aug 2026 11:52:03 +0300 Subject: [PATCH 008/109] metal : add fa-vec tunings for M3 Max, M5 and M5 Pro (#27863) * metal : add fa-vec tunings for M5 This is a followup contribution to efeda76b948f59ee52ea20db640bc4cf3dfe8ac1 as requested in https://github.com/ggml-org/llama.cpp/discussions/27668 to add support for additional Apple GPUs. I generated this output using the provided instructions: ```sh git clone https://github.com/ggml-org/llama.cpp cd llama.cpp cmake -B build -DGGML_METAL=ON cmake --build build --target ggml-metal-tuning -j ./build/bin/ggml-metal-tuning fa-vec --dtype f16,q8_0 > fa_vec_rows.txt 2> fa_vec_sweep.log ``` This ran on a machine with Apple M5. Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : add fa-vec tunings for M5 Pro This adds fa_vec_tuned_table records for Apple M5 Pro to ggml-metal-tuning.cpp. Contributed by SerayaEryn in https://github.com/ggml-org/llama.cpp/discussions/27668#discussioncomment-18157544 (F16, Q4_0, Q8_0; M5 Pro, 20 GPU cores). Assisted-by: pi:llama.cpp/Qwen3.8-27B * metal : add fa-vec tunings for M3 Max This adds fa_vec_tuned_table records for Apple M3 Max to ggml-metal-tuning.cpp. Contributed by TeeAaTeeUu in https://github.com/ggml-org/llama.cpp/discussions/27668#discussioncomment-18175220 (F16, Q8_0; M3 Max, MacBook Pro 64GB, low power mode). Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : whitespaces --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 305 +++++++++++++++++++++- 1 file changed, 304 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 7e99c1dd6612..c2139fe20b00 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -66,6 +66,7 @@ fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv) { // One row per kept bucket, plus per-(dtype,dk,dv) ne11-collapsed domain defaults // (ne11_b = FA_VEC_NE11_DEFAULT, ne01_b = domain). To retune or add a device, re-run the // sweep and paste its output. See ggml-metal-tuning.h for the row/lookup semantics. +// ref: https://github.com/ggml-org/llama.cpp/pull/27824 constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, @@ -448,6 +449,99 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, 2, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, 2, 3 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 0 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 0 }, { 4, 1 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 1 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 0 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 3, 1 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, @@ -508,6 +602,7 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 3 }, { 1, 2 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, @@ -699,7 +794,215 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, - { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 3, 2 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 3, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 2 }, { 4, 1 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 4 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 2, 0 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 3, 0 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 2, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 4 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, 3, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 2, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 2, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 3, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + + { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 1, 2 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 1, 4 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 2, 2 }, { 4, 4 } }, From be876204aa23bc0d8f890981fde2853555b132f9 Mon Sep 17 00:00:00 2001 From: Titaniumtown Date: Fri, 28 Aug 2026 01:53:31 -0700 Subject: [PATCH 009/109] sycl: bind the f16 KV cache in place for the oneDNN SDPA path (#27468) Measured at a live KV length of 34816 (32768 depth plus one 2048 ubatch), on Qwen3.8 27B Q4_K_S: per tensor 4 * 34816 * 256 * 2 B = 71.3 MB staged per call K and V, so 2x = 142.6 MB traffic per call read once, write once = 285.2 MB traffic per ubatch 285.2 MB * 16 calls = 4.56 GB One ubatch is one ggml_cgraph submission (llama_context::process_ubatch -> graph_compute), so that 4.56 GB is the cost of a single 2048-token prefill chunk, and it scales with the live KV length: the first ubatch of the same run, at seq = 2048, moves 0.27 GB. Reproduce the two measured inputs with: GGML_SCHED_DEBUG=2 llama-bench -m MODEL -p 8 -n 0 -r 1 -ngl 0 \ -fa on -ctk f16 -ctv f16 -v > nd.txt 2>&1 grep -E 'n_layer|n_head_kv|n_embd_head_k' nd.txt awk '/node # 0 /{g++} g==1 && /\(FLASH_ATTN\)/{n++} END{print n+0}' nd.txt --- ggml/src/ggml-sycl/fattn-onednn.cpp | 62 +++++++++++++++++++++++------ 1 file changed, 49 insertions(+), 13 deletions(-) diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index a501295192fb..d41c2ddce345 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -150,7 +151,8 @@ struct sdpa_partition { // Build + compile the contiguous-input GQA SDPA graph (MatMul->Divide->Add->SoftMax->MatMul), f32 out. // Mirrors the hardware-verified scratch/onednn_sdpa_probe.cpp build_gqa (partitions=1, sdp_primitive_kernel_t). -static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d) { +static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d, + const std::array & k_str, const std::array & v_str) try { using ltype = logical_tensor::layout_type; using dt = logical_tensor::data_type; using ldims = logical_tensor::dims; @@ -158,11 +160,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int const int rep = H / Hkv; const ldims q_sz = {1, Hkv, rep, q, d}, kv_sz = {1, Hkv, 1, seq, d}, s_sz = {1, Hkv, rep, q, seq}, sc = {1, 1, 1, 1, 1}, msk = {1, 1, 1, q, seq}, o_sz = {1, Hkv, rep, q, d}; + const ldims k_st(k_str.begin(), k_str.end()), v_st(v_str.begin(), v_str.end()); int64_t id = 0; sdpa_partition E; auto query = logical_tensor(id++, t, q_sz, ltype::strided); - auto key = logical_tensor(id++, t, kv_sz, ltype::strided); + auto key = logical_tensor(id++, t, kv_sz, k_st); auto score = logical_tensor(id++, fi, s_sz, ltype::strided); auto bmm1 = op(id++, op::kind::MatMul, "bmm1"); bmm1.set_attr(op::attr::transpose_b, true); // key is [.., seq, d] @@ -184,7 +187,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int smax.set_attr(op::attr::mode, "inf_as_zero"); smax.add_inputs({masked}); smax.add_outputs({probs}); - auto value = logical_tensor(id++, t, kv_sz, ltype::strided); + auto value = logical_tensor(id++, t, kv_sz, v_st); // f16 output is REQUIRED to hit sdp_primitive_kernel_t (the systolic micro-kernel); an f32 output // falls to larger_partition_kernel_t which materializes N^2 (confirmed: scratch/onednn_sdpa_kernel_probe.cpp). // converted to the f32 ggml dst in the permute below. @@ -198,6 +201,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int auto parts = g.get_partitions(); if (parts.size() != 1 || !parts[0].is_supported()) { + GGML_LOG_WARN("%s: oneDNN did not fuse the SDPA graph; falling back to TILE kernel\n", __func__); return E; // ok stays false -> caller falls back to TILE } E.ins = parts[0].get_input_ports(); @@ -209,6 +213,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int E.ok = true; return E; } +catch (const std::exception & e) { + // compile() can reject a stride set the partitioner never inspects; memoise the failure so the + // fallback costs one build rather than one per call. + GGML_LOG_WARN("%s: oneDNN SDPA partition build failed (%s); falling back to TILE kernel\n", __func__, e.what()); + return {}; +} void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) try { const ggml_tensor * Q = dst->src[0]; @@ -234,13 +244,34 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso ggml_sycl_pool_alloc Qf(ctx.pool(), (size_t) H * q * d); cont_to_f16_sycl((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream); - // K/V: use pool-alloc for both F16 and dequant paths. + // K/V: bind the f16 cache in place. llama.cpp permutes it to [token][head][dim], so its head + // plane is strided rather than dense, which is what an explicit stride vector expresses. + // Quantized and f32 KV still stage a dense copy -- the layout the k_str/v_str defaults describe. sycl::half * K_ptr = nullptr; sycl::half * V_ptr = nullptr; + std::array k_str{ Hkv * seq * d, seq * d, seq * d, d, 1 }; + std::array v_str = k_str; std::optional> Kf_pool; std::optional> Vf_pool; - if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) { + auto bindable = [](const ggml_tensor * t) { + return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 && + t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0; + }; + auto elem_strides = [](const ggml_tensor * t) { + const int64_t s1 = (int64_t) (t->nb[1] / t->nb[0]); + const int64_t s2 = (int64_t) (t->nb[2] / t->nb[0]); + const int64_t s3 = (int64_t) (t->nb[3] / t->nb[0]); + // dims are {mb=1, Hkv, rep=1, seq, d}; the size-1 dims at 0 and 2 never advance an address. + return std::array{ s3, s2, s2, s1, 1 }; + }; + + if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && bindable(K) && bindable(V)) { + K_ptr = (sycl::half *) K->data; + V_ptr = (sycl::half *) V->data; + k_str = elem_strides(K); + v_str = elem_strides(V); + } else if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) { Kf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d); Vf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d); cont_to_f16_sycl((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream); @@ -341,19 +372,24 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso ggml_sycl_pool_alloc outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d] - // compile once per (device, shape), reuse across layers/calls. + // compile once per (device, shape, KV strides), reuse across layers/calls. Stride 2 always + // repeats stride 1 and stride 4 is always 1, so the key covers every entry that can differ. static std::unordered_map cache; - char keyb[96]; - snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(), - (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d); + char keyb[256]; + snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(), + (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d, + (long long) k_str[0], (long long) k_str[1], (long long) k_str[3], + (long long) v_str[0], (long long) v_str[1], (long long) v_str[3]); auto it = cache.find(keyb); if (it == cache.end()) { - it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d)).first; + it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d, k_str, v_str)).first; } sdpa_partition & E = it->second; - // _supported() is authoritative: if it accepted this op the partition must build. - // A failure here is a gap in _supported() -- surface it, don't mask it with a fallback. - GGML_ASSERT(E.ok && "oneDNN SDPA partition failed to build for a _supported() shape"); + if (!E.ok) { + // oneDNN can decline a shape or a stride set that _supported() never sees; build_sdpa warns per key. + ggml_sycl_flash_attn_ext_tile(ctx, dst); + return; + } auto id2ptr = [&](size_t r) -> void * { if (r == E.id_q) return Qf.get(); From d077b4c21466cfad678b07b05b557599f4db3974 Mon Sep 17 00:00:00 2001 From: Ozymandias_EBON <112784549+johnkarlhill@users.noreply.github.com> Date: Fri, 28 Aug 2026 03:58:58 -0500 Subject: [PATCH 010/109] sycl: use TILE for quantized KV decode on BMG (#26689) Route quantized KV decode to TILE on Xe2 (BMG) only, keep VEC on other archs until validated there. --- ggml/src/ggml-sycl/fattn.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.cpp index a85eb721f6af..a85bca7cb3cd 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.cpp @@ -104,7 +104,6 @@ enum best_fattn_kernel { static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const ggml_tensor * dst) { - GGML_UNUSED(device); #ifndef SYCL_FLASH_ATTN GGML_UNUSED(dst); return BEST_FATTN_KERNEL_NONE; @@ -263,6 +262,11 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const } } else { if (Q->ne[1] <= 2) { + // TILE is faster for quantized KV decode on Xe2 (BMG); keep VEC on untested archs + const gpu_arch arch = ggml_sycl_info().devices[device].hw_info.arch; + if (arch == gpu_arch::intel_gpu_bmg_g21 || arch == gpu_arch::intel_gpu_bmg_g31) { + return BEST_FATTN_KERNEL_TILE; + } return BEST_FATTN_KERNEL_VEC; } } From b19cbe925be361d229f0fe03435affe4a2717f37 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Fri, 28 Aug 2026 11:46:30 +0200 Subject: [PATCH 011/109] convert: prevent ndarray conversion in LazyChunkedTensor (#27869) --- gguf-py/gguf/gguf_writer.py | 7 ++++++- gguf-py/gguf/lazy.py | 4 ++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 1f309ad2eafd..d95fe9b1ac3c 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -467,10 +467,15 @@ def write_tensors_to_file(self, *, progress: bool = False) -> None: shard_bar.reset(total=(total if total > 0 else None)) # relying on the fact that Python dicts preserve insertion order (since 3.7) - for ti in tensors.values(): + for name, ti in tensors.items(): assert ti.tensor is not None # can only iterate once over the tensors assert ti.tensor.nbytes == ti.nbytes + start = fout.tell() ti.tensor.tofile(fout) + # a short write here would only surface as a corrupt file at load time + if fout.tell() - start != ti.nbytes: + raise ValueError( + f"tensor {name!r} wrote {fout.tell() - start} bytes, expected {ti.nbytes}") if shard_bar is not None: shard_bar.update(ti.nbytes) if bar is not None: diff --git a/gguf-py/gguf/lazy.py b/gguf-py/gguf/lazy.py index 6a0aee881107..a39f22321597 100644 --- a/gguf-py/gguf/lazy.py +++ b/gguf-py/gguf/lazy.py @@ -251,6 +251,10 @@ def nbytes(self) -> int: def numpy(self) -> LazyChunkedTensor: return self + def __array__(self, *args, **kwargs): + # numpy would otherwise make a 1-element object array of self, and write 8 bytes + raise TypeError("LazyChunkedTensor cannot become an ndarray, it is written in chunks") + def quantize(self, qtype: Any) -> LazyChunkedTensor: from .constants import GGMLQuantizationType from .quants import QuantError, quant_shape_to_byte_shape From 511f9c1379a52516f328859af86daa124eddc717 Mon Sep 17 00:00:00 2001 From: Zijun Yu Date: Fri, 28 Aug 2026 19:42:07 +0800 Subject: [PATCH 012/109] OpenVINO: Update OV to 2026.3.1, whisper.cpp support, Qwen3.5 on NPU, and new ops (#27843) * OpenVINO Backend: Fuse IM2COL + MatMul convolution into OpenVINO convolution * ci:ggml-ov: Skip recurrent state rollback tests * ci:ggml-ov: Skip recurrent state rollback tests * Update OPENVINO.md * ggml-openvino : add env-var gated op support debugging * Fix ggml_rope_set_offset case * OpenVINO backend: Support Whisper.cpp * Fix code style * openvino : enable qwen35 on NPU Static shapes: - get_graph_input_shape() left the s_copy / s_copy-leaf inputs dynamic ([1,1,1,-1]) even in static mode, which propagated a dynamic slot dim through GET_ROWS into the conv/GDN state, the state reshapes and the GDN output. - With -np 1 the s_copy defrag remainder gathers zero rows; short-circuit that CPY to the untouched cache instead of emitting a degenerate Slice/Concat, and skip binding its zero-byte ggml tensor as an output (the dynamic path already did the latter, the static path wrote the full cache over a 0-byte buffer). Token-count independence: - In static mode the compiled model's token count is the prefill chunk size or 1, not the captured cgraph's. Offsets derived from the captured count were therefore wrong. Anchor the GDN state slice at the end of the packed [attn | state] output and drop the rs_src_begin runtime inputs, and make VIEWs over the GDN output / conv_input pass through so the consumer does the slicing. - CONT could not identify its token axis when the graph was captured with a single token (every trailing dim has the same stride and size 1) and baked the captured shape into the prefill model. Chunked prefill: - The last chunk is padded with fabricated tokens. Attention masks them, but the recurrent path folded them into cache_r/cache_s permanently. Add a chunk_valid_len runtime input, use it to zero g and beta for padded steps (making the recurrence an exact identity) and to end the conv snapshot window at the last valid token, and disable the recurrent-cache reset after the first chunk so earlier chunks are not wiped. - get_is_prefill() and the chunk loop bound read inp_pos->ne[0] directly, but IMROPE stacks 4 position planes, so every decode step was run through the padded prefill model and the loop ran extra out-of-bounds chunks. cache_rs_reset_idx/len now stay runtime Parameters in static mode, since can_reuse_statically() does not invalidate the cached model on ComputeParams changes. Add GGML_OPENVINO_FORCE_STATIC to exercise the static path on CPU. * Update to OpenVINO 2026.3.1 * ggml-openvino: forward NPU compilation mode parameters Add GGML_OPENVINO_NPU_COMPILE_CONFIG to the backend's cached environment so callers can configure the NPU compiler without using the generic property escape hatch. When the value is non-empty, pass it to OpenVINO as NPU_COMPILATION_MODE_PARAMS. This enables settings such as optimization-level=3 for NPU compilation while preserving the existing behavior when the variable is unset and leaving CPU and GPU configuration unchanged. Document the variable, its NPU-only scope, and the optimization-level=3 example in the OpenVINO backend runtime configuration table. * ggml-openvino : support RELU, POOL_2D, QUICK_GEGLU, and ROLL ops * reorder op table * exclude GPU/NPU failing POOL_2D case * move op type detection to compute_op_case * Relax rope supported cases * Fix pool case * Update openvino doc, gpu driver in ov docker * openvino: remove unused static remote context branch * openvino: parallelize static model build * Apply editorconfig --------- Co-authored-by: Mostafa Faheem Co-authored-by: Ravi Panchumarthy Co-authored-by: zhaixuejun1993 --- .devops/openvino.Dockerfile | 12 +- .github/workflows/build-cache.yml | 8 +- .github/workflows/build-openvino.yml | 19 +- .github/workflows/build-self-hosted.yml | 4 +- .github/workflows/release.yml | 8 +- ci/run.sh | 4 +- docs/backend/OPENVINO.md | 86 ++++--- ggml/src/ggml-openvino/CMakeLists.txt | 2 + ggml/src/ggml-openvino/ggml-decoder.cpp | 153 ++++++++++-- ggml/src/ggml-openvino/ggml-decoder.h | 12 +- .../src/ggml-openvino/ggml-openvino-extra.cpp | 12 +- ggml/src/ggml-openvino/ggml-openvino.cpp | 231 ++++++++++-------- ggml/src/ggml-openvino/openvino/op/cpy.cpp | 140 +++++++++-- .../openvino/op/flash_attn_ext.cpp | 104 ++++++-- .../openvino/op/gated_delta_net.cpp | 25 ++ .../openvino/op/glu_geglu_quick.cpp | 64 +++++ .../src/ggml-openvino/openvino/op/pool_2d.cpp | 53 ++++ ggml/src/ggml-openvino/openvino/op/roll.cpp | 36 +++ ggml/src/ggml-openvino/openvino/op/view.cpp | 7 + ggml/src/ggml-openvino/openvino/op_table.cpp | 5 + ggml/src/ggml-openvino/openvino/op_table.h | 3 + .../openvino/pass/fuse_to_conv.cpp | 212 ++++++++++++++++ .../openvino/pass/fuse_to_conv.h | 17 ++ .../openvino/translate_session.cpp | 6 +- ggml/src/ggml-openvino/openvino/utils.cpp | 1 + ggml/src/ggml-openvino/utils.cpp | 159 ++++++++---- ggml/src/ggml-openvino/utils.h | 4 +- 27 files changed, 1121 insertions(+), 266 deletions(-) create mode 100644 ggml/src/ggml-openvino/openvino/op/glu_geglu_quick.cpp create mode 100644 ggml/src/ggml-openvino/openvino/op/pool_2d.cpp create mode 100644 ggml/src/ggml-openvino/openvino/op/roll.cpp create mode 100644 ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.cpp create mode 100644 ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.h diff --git a/.devops/openvino.Dockerfile b/.devops/openvino.Dockerfile index a43e5c4993f8..13301ba287dd 100644 --- a/.devops/openvino.Dockerfile +++ b/.devops/openvino.Dockerfile @@ -1,12 +1,12 @@ -ARG OPENVINO_VERSION_MAJOR=2026.3 -ARG OPENVINO_VERSION_FULL=2026.3.0.22451.bd8d6542e3c +ARG OPENVINO_VERSION_MAJOR=2026.3.1 +ARG OPENVINO_VERSION_FULL=2026.3.1.22476.56d9685302d ARG UBUNTU_VERSION=24.04 # Intel GPU driver versions. https://github.com/intel/compute-runtime/releases -ARG IGC_VERSION=v2.38.2 -ARG IGC_VERSION_FULL=2_2.38.2+22051 -ARG COMPUTE_RUNTIME_VERSION=26.27.39122.11 -ARG COMPUTE_RUNTIME_VERSION_FULL=26.27.39122.11-0 +ARG IGC_VERSION=v2.40.13 +ARG IGC_VERSION_FULL=2_2.40.13+22418 +ARG COMPUTE_RUNTIME_VERSION=26.31.39395.13 +ARG COMPUTE_RUNTIME_VERSION_FULL=26.31.39395.13-0 ARG IGDGMM_VERSION=22.10.0 # Intel NPU driver versions. https://github.com/intel/linux-npu-driver/releases diff --git a/.github/workflows/build-cache.yml b/.github/workflows/build-cache.yml index 187427a8d4b0..4a23ec2d4d36 100644 --- a/.github/workflows/build-cache.yml +++ b/.github/workflows/build-cache.yml @@ -41,8 +41,8 @@ jobs: env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone @@ -69,8 +69,8 @@ jobs: env: # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone diff --git a/.github/workflows/build-openvino.yml b/.github/workflows/build-openvino.yml index 0316e7ad97e3..8e0326f4a44f 100644 --- a/.github/workflows/build-openvino.yml +++ b/.github/workflows/build-openvino.yml @@ -32,6 +32,8 @@ env: LLAMA_ARG_LOG_COLORS: 1 LLAMA_ARG_LOG_PREFIX: 1 LLAMA_ARG_LOG_TIMESTAMPS: 1 + # TODO: fix and re-enable the `test-llama-archs` and `test-recurrent-state-rollback` + CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-rollback" jobs: ubuntu-24-openvino: @@ -39,8 +41,8 @@ jobs: env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone @@ -78,26 +80,24 @@ jobs: - name: Test (CPU) id: cmake_test_cpu - # TODO: fix and re-enable the `test-llama-archs` test below run: | cd ${{ github.workspace }} - ctest --test-dir build/ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" --verbose --timeout 2000 + ctest --test-dir build/ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" --verbose --timeout 3000 - name: Test (GPU) id: cmake_test_gpu - # TODO: fix and re-enable the `test-llama-archs` test below run: | cd ${{ github.workspace }} export GGML_OPENVINO_DEVICE=GPU - ctest --test-dir build/ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" --verbose --timeout 3000 + ctest --test-dir build/ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" --verbose --timeout 3000 openvino-windows-2022: runs-on: windows-2022 env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone @@ -159,14 +159,13 @@ jobs: - name: Test (CPU) id: cmake_test_cpu shell: cmd - # TODO: fix and re-enable the `test-llama-archs` test below run: | REM Find extracted OpenVINO folder dynamically for /d %%i in (openvino_toolkit\*) do set OPENVINO_ROOT=%%i call "%OPENVINO_ROOT%\setupvars.bat" cd build - ctest --test-dir ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" -C Release --verbose --timeout 3000 + ctest --test-dir ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" -C Release --verbose --timeout 3000 - name: ccache-clear uses: ./.github/actions/ccache-clear diff --git a/.github/workflows/build-self-hosted.yml b/.github/workflows/build-self-hosted.yml index fe2ab815473c..ccfe2a604645 100644 --- a/.github/workflows/build-self-hosted.yml +++ b/.github/workflows/build-self-hosted.yml @@ -288,8 +288,8 @@ jobs: env: # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 98250c860650..76717d064bd0 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -415,8 +415,8 @@ jobs: env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Set OpenVINO version output @@ -529,8 +529,8 @@ jobs: env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Set OpenVINO version output diff --git a/ci/run.sh b/ci/run.sh index 1f1e4bc033c9..1701bc7ed058 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -189,8 +189,8 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then fi CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON" - # TODO: fix and re-enable the `test-llama-archs` test below - CTEST_EXTRA="-E test-llama-archs|test-recurrent-state-rollback-nemotron-h" + # TODO: fix and re-enable the `test-llama-archs` and `test-recurrent-state-rollback*` + CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-rollback" fi ## helpers diff --git a/docs/backend/OPENVINO.md b/docs/backend/OPENVINO.md index 3cdf631cebc2..9b43807d36b3 100644 --- a/docs/backend/OPENVINO.md +++ b/docs/backend/OPENVINO.md @@ -22,8 +22,8 @@ The OpenVINO backend is implemented in `ggml/src/ggml-openvino` and provides a t - [0. Prerequisites](#0-prerequisites) - [1. Install OpenVINO Runtime](#1-install-openvino-runtime) - [2. Build llama.cpp with OpenVINO Backend](#2-build-llamacpp-with-openvino-backend) - - [Automated Ubuntu Build Script](#automated-ubuntu-build-script) - - [Automated Windows Build Script](#automated-windows-build-script) + - [Ubuntu Build Script](#ubuntu-build-script) + - [Windows Build Script](#windows-build-script) - [3. Download Sample Model](#3-download-sample-model) - [4. Run Inference with OpenVINO Backend](#4-run-inference-with-openvino-backend) - [5. Docker Build](#5-docker-build) @@ -96,7 +96,7 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_ - **SL** = Stateless (`GGML_OPENVINO_STATEFUL_EXECUTION=0`) - **SF** = Stateful (`GGML_OPENVINO_STATEFUL_EXECUTION=1`) - Note: The NPU operates in stateless mode only. -- **Validation system:** Intel® Core™ Ultra 5 238V (Lunar Lake) | 32 GB RAM | Ubuntu 24.04 | Intel OpenCL GPU Driver 26.18.38308.1 | Intel NPU Driver 1.33.0. +- **Validation system:** Intel® Core™ Ultra 5 238V (Lunar Lake) | 32 GB RAM | Ubuntu 24.04 | Intel OpenCL GPU Driver 26.31.39395.13-0 | Intel NPU Driver 1.35.0. - See [Known Limitations](#known-limitations) for context on observed failures. | Model | CPU (SL / SF) | GPU (SL / SF) | NPU (SL) | @@ -105,27 +105,32 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_ | [bartowski/Llama-3.2-3B-Instruct-Q4_K_M](https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | [bartowski/Meta-Llama-3.1-8B-Instruct-Q4_K_M](https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | -| [Qwen/qwen2.5-1.5b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [Qwen/qwen2.5-coder-7b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/Qwen_Qwen3-0.6B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-0.6B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/Qwen_Qwen3-1.7B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-1.7B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [Qwen/Qwen3-4B-Q4_K_M](https://huggingface.co/Qwen/Qwen3-4B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [lm-kit/Qwen3-8B-Q4_K_M](https://huggingface.co/lm-kit/qwen-3-8b-instruct-gguf) | ✓ / ✓ | ✓ / ✗ | ✓ | +| [Qwen/qwen2.5-1.5b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [Qwen/qwen2.5-coder-7b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/Qwen_Qwen3-0.6B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-0.6B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/Qwen_Qwen3-1.7B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-1.7B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [Qwen/Qwen3-4B-Q4_K_M](https://huggingface.co/Qwen/Qwen3-4B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [lm-kit/Qwen3-8B-Q4_K_M](https://huggingface.co/lm-kit/qwen-3-8b-instruct-gguf) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/Qwen_Qwen3.5-0.8B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-0.8B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | +| [bartowski/Qwen_Qwen3.5-2B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-2B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | +| [bartowski/Qwen_Qwen3.5-4B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-4B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | +| [lmstudio-community/Qwen3.5-9B-Q4_K_M](https://huggingface.co/lmstudio-community/Qwen3.5-9B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | | | | | | -| [unsloth/gemma-3-4b-it-Q4_K_M](https://huggingface.co/unsloth/gemma-3-4b-it-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/google_gemma-4-E2B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E2B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ | +| [unsloth/gemma-3-4b-it-Q4_K_M](https://huggingface.co/unsloth/gemma-3-4b-it-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/google_gemma-4-E2B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E2B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | | [bartowski/google_gemma-4-E4B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E4B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ | -| [bartowski/gemma-4-12B-it-Q4_K_M](https://huggingface.co/bartowski/gemma-4-12B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | +| [bartowski/gemma-4-12B-it-Q4_K_M](https://huggingface.co/bartowski/gemma-4-12B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ | | | | | | -| [bartowski/Phi-3-mini-4k-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3-mini-4k-instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/Phi-3.5-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3.5-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | +| [bartowski/Phi-3-mini-4k-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3-mini-4k-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/Phi-3.5-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3.5-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/microsoft_Phi-4-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/microsoft_Phi-4-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [bartowski/Mistral-7B-Instruct-v0.3-Q4_K_M](https://huggingface.co/bartowski/Mistral-7B-Instruct-v0.3-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | [QuantFactory/Ministral-3b-instruct.Q4_K_M](https://huggingface.co/QuantFactory/Ministral-3b-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | [bartowski/Ministral-8B-Instruct-2410-Q4_K_M](https://huggingface.co/bartowski/Ministral-8B-Instruct-2410-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [bartowski/DeepSeek-R1-Distill-Llama-8B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Llama-8B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | -| [bartowski/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | +| [bartowski/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [ibm-granite/granite-4.0-350m-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-350m-GGUF) | ✓ / ✓ | ✗ / ✗ | ✓ | | [ibm-granite/granite-4.0-micro-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-micro-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | @@ -133,10 +138,10 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_ | [ibm-research/granite-3.2-8b-instruct-Q4_K_M](https://huggingface.co/ibm-research/granite-3.2-8b-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [HuggingFaceTB/smollm2-1.7b-instruct-q4_k_m](https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | -| [openbmb/MiniCPM-V-2_6-Q4_K_M](https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/tencent_Hunyuan-7B-Instruct-Q4_K_M](https://huggingface.co/bartowski/tencent_Hunyuan-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-Q4_K_M](https://huggingface.co/LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/prism-ml_Bonsai-8B-unpacked-Q4_K_M](https://huggingface.co/bartowski/prism-ml_Bonsai-8B-unpacked-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | +| [openbmb/MiniCPM-V-2_6-Q4_K_M](https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/tencent_Hunyuan-7B-Instruct-Q4_K_M](https://huggingface.co/bartowski/tencent_Hunyuan-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-Q4_K_M](https://huggingface.co/LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/prism-ml_Bonsai-8B-unpacked-Q4_K_M](https://huggingface.co/bartowski/prism-ml_Bonsai-8B-unpacked-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [gpustack/bge-m3-Q4_K_M.gguf](https://huggingface.co/gpustack/bge-m3-GGUF) | ✓ | ✗ | ✗ | @@ -217,18 +222,18 @@ cmake --build build\ReleaseOV --parallel > [!NOTE] > The Windows install path is `C:\Intel\openvino` (no spaces) to avoid quoting problems some CMake/Ninja toolchains have with `C:\Program Files (x86)\...`. Adjust to wherever you installed OpenVINO Runtime. From `cmd`, run `C:\Intel\openvino\setupvars.bat`; from PowerShell, run `& "C:\Intel\openvino\setupvars.ps1"` instead. Once the build is finished you can launch the binaries from any `cmd` or `PowerShell` window after sourcing the matching `setupvars` script for that shell. -#### Automated Ubuntu Build Script +#### Ubuntu Build Script For Ubuntu24 users, the following shell script automates the prerequisite installs (build tools, OpenCL ICD), the OpenVINO Runtime download/extract/setup, and the Ninja-based llama.cpp build. -Save the following as `ubuntu-llamacpp-ov-install.sh` next to where you want the `llama.cpp` folder to land, then run it: +Save the following as `build-llamacpp-ov.sh` next to where you want the `llama.cpp` folder to land, then run it: ```bash -chmod +x ubuntu-llamacpp-ov-install.sh -./ubuntu-llamacpp-ov-install.sh +chmod +x build-llamacpp-ov.sh +./build-llamacpp-ov.sh ```
-Click to expand ubuntu-llamacpp-ov-install.sh +Click to expand build-llamacpp-ov.sh ```bash #!/usr/bin/env bash @@ -237,8 +242,8 @@ chmod +x ubuntu-llamacpp-ov-install.sh # ============================================ set -euo pipefail -OPENVINO_VERSION_MAJOR="2026.3" -OPENVINO_VERSION_FULL="2026.3.0.22451.bd8d6542e3c" +OPENVINO_VERSION_MAJOR="2026.3.1" +OPENVINO_VERSION_FULL="2026.3.1.22476.56d9685302d" SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" OPENVINO_INSTALL_DIR="/opt/intel/openvino_${OPENVINO_VERSION_MAJOR}" @@ -313,8 +318,9 @@ fi echo "============================================" echo "Configuring with CMake..." echo "============================================" -# shellcheck disable=SC1091 +set +u source "${OPENVINO_ROOT}/setupvars.sh" +set -u cmake -B build/ReleaseOV -G Ninja \ -DCMAKE_BUILD_TYPE=Release \ @@ -334,27 +340,27 @@ echo " ./build/ReleaseOV/bin/llama-cli -m model.gguf" ``` > [!NOTE] -> The script pins OpenVINO `2026.3` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. +> The script pins OpenVINO `2026.3.1` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release.
-#### Automated Windows Build Script +#### Windows Build Script For Windows users, the following `.bat` script automates the prerequisite installs (Git, Ninja, CMake, Visual Studio 2022 Build Tools, vcpkg + OpenCL), the OpenVINO Runtime download/extract, and the Ninja-based llama.cpp build. -Save the following as `windows-llamacpp-ov-install.bat` next to where you want the `llama.cpp` to land, then run it from either **Command Prompt** or **PowerShell**: +Save the following as `build-llamacpp-ov.bat` next to where you want the `llama.cpp` to land, then run it from either **Command Prompt** or **PowerShell**: ```cmd :: Command Prompt -windows-llamacpp-ov-install.bat +build-llamacpp-ov.bat ``` ```powershell # PowerShell -.\windows-llamacpp-ov-install.bat +.\build-llamacpp-ov.bat ```
-Click to expand windows-llamacpp-ov-install.bat +Click to expand build-llamacpp-ov.bat ```bat @echo off @@ -364,8 +370,8 @@ REM ============================================ REM llama.cpp OpenVINO Build Script (Ninja) REM ============================================ -set "OPENVINO_VERSION_MAJOR=2026.3" -set "OPENVINO_VERSION_FULL=2026.3.0.22451.bd8d6542e3c" +set "OPENVINO_VERSION_MAJOR=2026.3.1" +set "OPENVINO_VERSION_FULL=2026.3.1.22476.56d9685302d" set "SCRIPT_DIR=%~dp0" set "VCPKG_DIR=C:\vcpkg" @@ -453,9 +459,6 @@ if exist "%OPENVINO_INSTALL_DIR%\setupvars.bat" ( ) REM Move the single top-level folder contents into the versioned install dir. - REM NOTE: delayed expansion (!VAR!) is required because the surrounding else( ... ) - REM block is parsed once up-front, so %OPENVINO_EXTRACTED% would expand to "" here - REM and xcopy would then treat "\*" as C:\* and fail with "Cannot perform a cyclic copy". set "OPENVINO_EXTRACTED=" for /d %%i in ("%OPENVINO_EXTRACT_TMP%\*") do set "OPENVINO_EXTRACTED=%%i" if not defined OPENVINO_EXTRACTED ( @@ -547,7 +550,7 @@ endlocal ``` > [!NOTE] -> The script pins OpenVINO `2026.3` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. From any new shell, source the matching `setupvars` script via the junction — `call "C:\Intel\openvino\setupvars.bat"` from `cmd`, or `& "C:\Intel\openvino\setupvars.ps1"` from PowerShell. If `winget` cannot register Visual Studio Build Tools on first run, install them once manually and re-run the script from an elevated **Developer Command Prompt for VS 2022**. +> The script pins OpenVINO `2026.3.1` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. From any new shell, source the matching `setupvars` script via the junction — `call "C:\Intel\openvino\setupvars.bat"` from `cmd`, or `& "C:\Intel\openvino\setupvars.ps1"` from PowerShell. If `winget` cannot register Visual Studio Build Tools on first run, install them once manually and re-run the script from an elevated **Developer Command Prompt for VS 2022**.
@@ -712,6 +715,7 @@ Boolean flags follow a uniform convention: set to a **positive integer** (e.g. ` | `GGML_OPENVINO_CACHE_DIR` | String | `not set` | Directory for OpenVINO model caching (recommended: `/tmp/ov_cache`). Enables model caching when set. **Not supported on NPU devices.** | | `GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR` | String | `not set` | Directory for the frontend compiled-model cache. When set, OpenVINO compiled models are exported as blobs and imported on later runs to skip weight requantization, graph conversion, and compilation for matching single-graph models. | | `GGML_OPENVINO_PREFILL_CHUNK_SIZE`| Integer | `256` | Token chunk size for **NPU** prefill (NPU-only; ignored on CPU/GPU). Must be a positive integer; otherwise the default is used. | +| `GGML_OPENVINO_NPU_COMPILE_CONFIG` | String | `not set` | NPU-only compiler mode parameters forwarded to OpenVINO as `NPU_COMPILATION_MODE_PARAMS`, for example `optimization-level=3`. | | `GGML_OPENVINO_STATEFUL_EXECUTION`| Boolean | `0` | Enable stateful KV cache for better performance. Recommended on CPU, GPU. | | `GGML_OPENVINO_DISABLE_CACHE` | Boolean | `0` | Disable the in-process compiled-model / decoder cache (cache is on by default). Set to `1` to disable. | | `GGML_OPENVINO_DISABLE_KV_SLICE` | Boolean | `0` | Disable the KV-cache input-tensor slicing optimization (slicing is on by default on CPU/GPU). Set to `1` to disable. | @@ -725,9 +729,11 @@ Boolean flags follow a uniform convention: set to a **positive integer** (e.g. ` | `GGML_OPENVINO_DEBUG_INPUT` | Boolean | `0` | Enable input debugging and print input tensor info. | | `GGML_OPENVINO_DEBUG_OUTPUT` | Boolean | `0` | Enable output debugging and print output tensor info. | | `GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS` | Boolean | `0` | Print tensor address map once. | +| `GGML_OPENVINO_LOG_UNSUPPORTED_OPS`| Boolean | `0` | Log warning messages with tensor details and rejection reasons for any ops not supported by the OpenVINO backend. Emits at `WARN` level (requires `--log-verbosity >= 2`, enabled by default). | > [!NOTE] ->`GGML_OPENVINO_STATEFUL_EXECUTION` is an **Experimental** feature to allow stateful execution for managing the KV cache internally inside the OpenVINO model, improving performance on CPUs and GPUs. Stateful execution is not effective on NPUs, and not all models currently support this feature. This feature is experimental and has been validated only with the llama-simple, llama-cli, llama-bench, and llama-run applications and is recommended to enable for the best performance. Other applications, such as llama-server and llama-perplexity, are not yet supported. +> - `GGML_OPENVINO_STATEFUL_EXECUTION` is an **Experimental** feature to allow stateful execution for managing the KV cache internally inside the OpenVINO model, improving performance on CPUs and GPUs. Stateful execution is not effective on NPUs, and not all models currently support this feature. This feature is experimental and has been validated only with the llama-simple, llama-cli, llama-bench, and llama-run applications and is recommended to enable for the best performance. Other applications, such as llama-server and llama-perplexity, are not yet supported. +> - `GGML_OPENVINO_LOG_UNSUPPORTED_OPS` emits logs at `WARN` level (`GGML_LOG_WARN`), which requires application log verbosity `--log-verbosity >= 2` (or `-lv 2`). ### Example Usage diff --git a/ggml/src/ggml-openvino/CMakeLists.txt b/ggml/src/ggml-openvino/CMakeLists.txt index cc089b721fc3..af3e0758ca2f 100644 --- a/ggml/src/ggml-openvino/CMakeLists.txt +++ b/ggml/src/ggml-openvino/CMakeLists.txt @@ -1,6 +1,8 @@ find_package(OpenVINO REQUIRED COMPONENTS Runtime Threading) find_package(OpenCL REQUIRED) +message(STATUS "Found OpenVINO: ${OpenVINO_DIR} (found version \"${OpenVINO_VERSION}\")") + file(GLOB_RECURSE GGML_HEADERS_OPENVINO "*.h" "*.hpp") file(GLOB_RECURSE GGML_SOURCES_OPENVINO "*.cpp") diff --git a/ggml/src/ggml-openvino/ggml-decoder.cpp b/ggml/src/ggml-openvino/ggml-decoder.cpp index 599f41aebbdc..006e005cb7aa 100644 --- a/ggml/src/ggml-openvino/ggml-decoder.cpp +++ b/ggml/src/ggml-openvino/ggml-decoder.cpp @@ -357,6 +357,18 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { break; } case GGML_OP_VIEW: { + if (m_is_static && node->src[0] != nullptr && + (node->src[0]->op == GGML_OP_GATED_DELTA_NET || node->src[0]->op == GGML_OP_CONCAT)) { + // VIEW slicing a GATED_DELTA_NET combined [attn|state] output, or the conv_input + // CONCAT. The consuming CPY/RMS_NORM op recovers the true window at runtime via + // ssm_state_size / the fixed conv kernel width, so this VIEW must stay an identity + // pass-through of the full source here too (it already is on the dynamic path); + // otherwise the generic static-mode Slice below would bake in the *captured* + // cgraph's token count, which is wrong once the compiled static model runs with a + // different token count (prefill chunk size or 1). + op_case = 1; + break; + } if (node->src[0]->op == GGML_OP_VIEW) { auto * src = node->src[0]; if (ggml_nelements(node) != ggml_nelements(src)) { @@ -408,6 +420,23 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { } break; } + case GGML_OP_POOL_2D: { + const ggml_op_pool pool_mode = static_cast(node->op_params[0]); + switch (pool_mode) { + case GGML_OP_POOL_MAX: { + op_case = 1; + break; + } + case GGML_OP_POOL_AVG: { + op_case = 2; + break; + } + default: + op_case = 0; + break; + } + break; + } case GGML_OP_CPY: { if (node->src[0]->op == GGML_OP_VIEW) { if (node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET) { @@ -425,6 +454,31 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { is_kvcache(node->src[1]->view_src, nullptr)) { // s_copy defrag remainder writeback: gathered extra state rows copied back into the cache op_case = 3; + } else if (node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src != nullptr) { + // op_case 5: KV write for decoder self-attention (dynamic write offset) + // op_case 6: KV write for encoder self-attn or cross-attn (static offset) + const ggml_tensor * kv_buf = node->src[1]->view_src; + if (kv_buf->ne[1] == 1 && kv_buf->ne[2] == 1 && kv_buf->ne[3] == 1) { + op_case = 6; + // Forward-scan the graph for a FLASH_ATTN_EXT that reads from + // the same buffer. Having a mask (src[3] != nullptr) implies + // decoder self-attention and the write offset is dynamic. + for (int i = 0; i < m_cgraph->n_nodes; i++) { + const ggml_tensor * n = m_cgraph->nodes[i]; + if (n->op != GGML_OP_FLASH_ATTN_EXT) { + continue; + } + // K (src[1]) and V (src[2]) are 3-D views whose view_src is + // the flat KV buffer we are writing to. + if ((n->src[1] != nullptr && n->src[1]->view_src == kv_buf) || + (n->src[2] != nullptr && n->src[2]->view_src == kv_buf)) { + if (n->src[3] != nullptr) { + op_case = 5; // decoder self-attention: mask present + } + break; + } + } + } } break; } @@ -448,6 +502,15 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { } break; } + case GGML_OP_FLASH_ATTN_EXT: { + if (node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src != nullptr) { + const ggml_tensor * kv_buf = node->src[1]->view_src; + if (kv_buf->ne[1] == 1 && kv_buf->ne[2] == 1 && kv_buf->ne[3] == 1) { + op_case = (node->src[3] != nullptr) ? 1 : 2; + } + } + break; + } default: break; } @@ -479,23 +542,35 @@ std::pair GgmlOvDecoder::compute_llm_params(ggml_cgr switch (node->op) { case GGML_OP_FLASH_ATTN_EXT: - if (node->src[0] == nullptr || node->src[1] == nullptr || node->src[3] == nullptr) { + if (node->src[0] == nullptr || node->src[1] == nullptr) { return -1; } switch (node->src[1]->op) { case GGML_OP_PERMUTE: - // case 0: node op is FLASH_ATTN_EXT, src 1 not null & op is PERMUTE & the permuted tensor src is the view of cache k - if (node->src[1]->src[0] != nullptr && node->src[1]->src[0]->op == GGML_OP_VIEW) { + // case 0: src[1] is PERMUTE of a cache VIEW, mask required + if (node->src[3] != nullptr && node->src[1]->src[0] != nullptr && + node->src[1]->src[0]->op == GGML_OP_VIEW) { return 0; } break; case GGML_OP_CPY: - // case 1: node op is FLASH_ATTN_EXT, src 1 not null & op is CPY & the copied tensor src is PERMUTE & the permuted tensor src is the view of cache k - if (node->src[1]->src[0] != nullptr && node->src[1]->src[0]->op == GGML_OP_PERMUTE && - node->src[1]->src[0]->src[0] != nullptr && node->src[1]->src[0]->src[0]->op == GGML_OP_VIEW) { + // case 1: src[1] is CPY of a PERMUTE(VIEW), mask required + if (node->src[3] != nullptr && node->src[1]->src[0] != nullptr && + node->src[1]->src[0]->op == GGML_OP_PERMUTE && node->src[1]->src[0]->src[0] != nullptr && + node->src[1]->src[0]->src[0]->op == GGML_OP_VIEW) { return 1; } break; + case GGML_OP_VIEW: + // cases 4/5/6: whisper - K is a direct non-contiguous VIEW_3D of a KV cache + if (node->src[1]->view_src != nullptr) { + if (node->src[3] != nullptr) { + return 4; // decoder self-attention + } else { + return 5; // cross-attention or encoder self-attention + }; + } + break; default: break; } @@ -548,6 +623,18 @@ std::pair GgmlOvDecoder::compute_llm_params(ggml_cgr cache_k_permute = node->src[0]->src[0]->src[0]; mask = node->src[1]; break; + case 4: + case 5: { + // whisper: K is a direct VIEW_3D of the KV buffer, no PERMUTE node + auto * cache_k_view = node->src[1]; // VIEW_3D of kv_self.k or kv_cross.k` + compute_params.token_len_per_seq = node->src[0]->ne[1]; + if (attention_pattern_case == 4) { + compute_params.attention_size = cache_k_view->ne[1]; + } else { + compute_params.attention_size_static = cache_k_view->ne[1]; + } + continue; + } default: break; } @@ -654,10 +741,8 @@ std::pair GgmlOvDecoder::compute_llm_params(ggml_cgr ComputeParams::RsWriteback writeback; writeback.slot_begin = (int) (dest_view->view_offs / row_bytes); if (is_conv) { - // conv_input column the copied window starts at writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[0]); } else if (is_gdn) { - // first row of the state part of the gated-delta-net output writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[1]); } compute_params.rs_writebacks[get_tensor_ov_name(cgraph, node)] = writeback; @@ -718,11 +803,15 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op, } else if (is_kvcache(input, op)) { // kvcache input_shape = ov::PartialShape{get_shape(input)}; - if (!m_is_static) { + // Whisper.cpp uses a fixed size 1D KV buffer [N, 1, 1, 1] (GGML) or [1, 1, 1, N] (OV). + // the token fill level is handled by token_len_per_seq + dynamic mask input. + // skip dynamic dim and stateful reshape for this layout. + const bool is_flat_kv = (input->ne[1] == 1 && input->ne[2] == 1 && input->ne[3] == 1); + if (!m_is_static && !is_flat_kv) { // do not fix ctx size to make llama-bench work across test params input_shape[2] = -1; } - if (is_stateful()) { + if (is_stateful() && !is_flat_kv) { // Convert stateless KV cache layout [1, 1, seq, n_heads_kv * head_size] // to stateful layout [1, seq, n_heads_kv, head_size]. assert(input_shape.size() == 4 && input_shape[0] == 1 && input_shape[1] == 1 && @@ -738,7 +827,9 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op, input_shape = ov::PartialShape{1, 1, 1, len}; } else if (is_inp_s_copy(input, op) || is_s_copy_leaf(input)) { - input_shape = ov::PartialShape{1, 1, 1, -1}; + // On NPU the total slot count (n_seq_max) is fixed at translation time, so the s_copy + // index list has a static length; on CPU/GPU it may change across compiles (defrag). + input_shape = m_is_static ? ov::PartialShape{get_shape(input)} : ov::PartialShape{1, 1, 1, -1}; } else { input_shape = ov::PartialShape{get_shape(input)}; @@ -790,13 +881,16 @@ void GgmlOvDecoder::add_extra_inputs() { // see llama_kv_cache_unified::get_n_kv and llama_kv_cache_unified::get_padding. // 2. `n_seq_active` and `seq_active_start`, used in FLASH_ATTN_EXT to indicate the active sequences in the batch - auto create_1d_input = [this](const std::string & name, int64_t value) { - m_model_extra_inputs[name] = {ov::element::i64, ov::Shape{1}, value, !m_is_static}; + auto create_1d_input = [this](const std::string & name, int64_t value, bool force_parameter = false) { + m_model_extra_inputs[name] = {ov::element::i64, ov::Shape{1}, value, force_parameter || !m_is_static}; }; if (m_compute_params.attention_size != -1) { create_1d_input("attention_size", m_compute_params.attention_size); } + if (m_compute_params.attention_size_static != -1) { + create_1d_input("attention_size_static", m_compute_params.attention_size_static); + } if (m_compute_params.attention_size_swa != -1) { create_1d_input("attention_size_swa", m_compute_params.attention_size_swa); } @@ -809,17 +903,32 @@ void GgmlOvDecoder::add_extra_inputs() { // create_1d_input("token_len", m_compute_params.token_len_per_seq * m_compute_params.n_seq_active); if (m_compute_params.cache_rs_reset_idx != -1) { - create_1d_input("cache_rs_reset_idx", m_compute_params.cache_rs_reset_idx); - create_1d_input("cache_rs_reset_len", m_compute_params.cache_rs_reset_len); + // Whether/which cache slot to reset varies per compute call (e.g. a new sequence starting + // vs. continued decoding). can_reuse_statically() does not invalidate the cached static + // model on ComputeParams changes, so these must stay runtime Parameters even when static + // (scale.cpp op_case 1 only uses them in value comparisons, never as Slice bounds, so this + // does not reintroduce dynamic shapes). + create_1d_input("cache_rs_reset_idx", m_compute_params.cache_rs_reset_idx, /*force_parameter=*/true); + create_1d_input("cache_rs_reset_len", m_compute_params.cache_rs_reset_len, /*force_parameter=*/true); } if (m_compute_params.s_copy_active_slot_len != -1) { create_1d_input("s_copy_active_slot_len", m_compute_params.s_copy_active_slot_len); + if (m_is_static) { + // Number of real tokens in the current prefill chunk. The last chunk is padded with + // fabricated token ids; attention masks them out, but the recurrent (GDN/conv) path + // would otherwise fold them into cache_r/cache_s permanently. Varies per chunk, so it + // must stay a runtime Parameter; it is only compared against a Range or used as Gather + // indices, so it does not make any shape dynamic. + create_1d_input("chunk_valid_len", get_static_n_tokens(), /*force_parameter=*/true); + } } for (const auto & [node_name, writeback] : m_compute_params.rs_writebacks) { create_1d_input("rs_slot_begin_" + node_name, writeback.slot_begin); - create_1d_input("rs_src_begin_" + node_name, writeback.src_begin); + if (!m_is_static) { + create_1d_input("rs_src_begin_" + node_name, writeback.src_begin); + } } } @@ -1785,13 +1894,23 @@ void GgmlOvDecoder::compute_node_dynamic_dims() { auto dynamic_dim_stride = src_logical_nb[dynamic_dim_idx] / ggml_type_size(node->src[0]->type) * ggml_type_size(node->type); int matched_dim_count = 0; + int first_matched_dim = -1; for (int i = 0; i < GGML_MAX_DIMS; i++) { if (node->nb[i] == dynamic_dim_stride && node->ne[i] == node->src[0]->ne[dynamic_dim_idx]) { + if (first_matched_dim == -1) { + first_matched_dim = i; + } m_node_dynamic_dims[node] = i; matched_dim_count++; } } - if (matched_dim_count != 1) { + if (matched_dim_count > 1 && node->src[0]->ne[dynamic_dim_idx] == 1) { + // Single-token capture: every trailing dim is size 1 with the same stride, so + // the match is ambiguous. The lowest index is the real axis; the rest are + // ggml's size-1 padding. Bailing out here would bake the captured token count + // into the static prefill model, which then runs with a different one. + m_node_dynamic_dims[node] = first_matched_dim; + } else if (matched_dim_count != 1) { m_node_dynamic_dims[node] = -1; GGML_LOG_WARN("ggml-openvino: cannot determine dynamic dim for CONT node '%s', src[0]: '%s'\n", node->name, node->src[0]->name); diff --git a/ggml/src/ggml-openvino/ggml-decoder.h b/ggml/src/ggml-openvino/ggml-decoder.h index 8e39a26c8b79..74cb7385029a 100644 --- a/ggml/src/ggml-openvino/ggml-decoder.h +++ b/ggml/src/ggml-openvino/ggml-decoder.h @@ -47,6 +47,7 @@ struct ComputeParams { int seq_active_start = 0; int attention_size = -1; int attention_size_swa = -1; + int attention_size_static = -1; // encoder/cross-attn KV fill level (whisper) int input_len = -1; int token_len_per_seq = -1; int past_kv_len = -1; @@ -84,14 +85,15 @@ struct ComputeParams { struct RsWriteback { int slot_begin = 0; // first cache slot written by the CPY - int src_begin = 0; // where the copied data starts in the source tensor (in rows of it) + int src_begin = 0; // first source row or column copied by the CPY }; std::map rs_writebacks; - // Offsets of the state cache writeback CPY nodes, keyed by node name. They change with the - // batch (kv head, active sequence count, token count) and, with rollback enabled - // (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot, each snapshot - // taking a different conv_input window. Passed to the cached model as runtime inputs. + // Destination slot offset of each state cache writeback CPY node, keyed by node name. It + // changes with the batch (kv head, active sequence count) and, with rollback enabled + // (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot. Passed to the + // cached model as a runtime input. Dynamic models also receive the source-side offset; static + // models use a fixed end-anchored offset in the translator. }; class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder { diff --git a/ggml/src/ggml-openvino/ggml-openvino-extra.cpp b/ggml/src/ggml-openvino/ggml-openvino-extra.cpp index 36c749244f83..36dfa4d9471b 100644 --- a/ggml/src/ggml-openvino/ggml-openvino-extra.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino-extra.cpp @@ -32,6 +32,8 @@ void ggml_openvino_device_config::init() { "GGML_OPENVINO_DEVICE", "GGML_OPENVINO_CACHE_DIR", "GGML_OPENVINO_DEBUG_NODE", + "GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR", + "GGML_OPENVINO_NPU_COMPILE_CONFIG", // Integer values (use ggml_openvino_getenv_int) "GGML_OPENVINO_PREFILL_CHUNK_SIZE", // Boolean toggles (treated as int flags via ggml_openvino_getenv_int) @@ -41,6 +43,9 @@ void ggml_openvino_device_config::init() { "GGML_OPENVINO_DUMP_IR", "GGML_OPENVINO_DEBUG_INPUT", "GGML_OPENVINO_DEBUG_OUTPUT", + // Force the static (NPU-shape) compute path on any device, e.g. GGML_OPENVINO_DEVICE=CPU, + // to test the static-shape translation without NPUW/real NPU hardware in the loop. + "GGML_OPENVINO_FORCE_STATIC", "GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS", "GGML_OPENVINO_ENABLE_CACHE", "GGML_OPENVINO_DISABLE_CACHE", @@ -50,7 +55,7 @@ void ggml_openvino_device_config::init() { "GGML_OPENVINO_MEMORY_OPTIMIZE", "GGML_OPENVINO_RELEASE_WEIGHTS", "GGML_OPENVINO_REDUCE_COMPILE_MEM", - "GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR", + "GGML_OPENVINO_LOG_UNSUPPORTED_OPS", }; for (const char * const & env_var : env_var_names) { @@ -85,6 +90,11 @@ void ggml_openvino_device_config::init() { compile_config["NPUW_CACHE_DIR"] = cache_dir; compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE)); } + const char * compilation_mode_params = + ggml_openvino_getenv_str("GGML_OPENVINO_NPU_COMPILE_CONFIG"); + if (compilation_mode_params && strlen(compilation_mode_params) > 0) { + compile_config["NPU_COMPILATION_MODE_PARAMS"] = compilation_mode_params; + } } else if (cache_dir && strlen(cache_dir) > 0) { compile_config.insert(ov::cache_dir(cache_dir)); compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE)); diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index e299e16c778a..4b1789713d1d 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -908,11 +908,27 @@ static bool has_non_contiguous_view_input(const ggml_tensor * op) { } static bool is_supported_flash_attn_pattern(const ggml_tensor * op) { - // pattern of q,k,v should be q->op==PERMUTE, q->src[0]->op==VIEW, q->src[0]->src[0]->view_src==nullptr + // Each Q/K/V input must follow one of: + // PERMUTE -> VIEW -> base (view_src==nullptr) (llama KV-cache path) + // PERMUTE -> RESHAPE -> base (view_src==nullptr) (whisper Q) + // VIEW -> base (view_src==nullptr) (whisper K/V from kv_pad) for (int i = 0; i < 3; i++) { const ggml_tensor * src = op->src[i]; - if (src->op != GGML_OP_PERMUTE || src->src[0] == nullptr || src->src[0]->op != GGML_OP_VIEW || - src->src[0]->src[0] == nullptr || src->src[0]->src[0]->view_src != nullptr) { + if (src->op == GGML_OP_PERMUTE) { + if (src->src[0] == nullptr) { + return false; + } + if (src->src[0]->op != GGML_OP_VIEW && src->src[0]->op != GGML_OP_RESHAPE) { + return false; + } + if (src->src[0]->src[0] == nullptr || src->src[0]->src[0]->view_src != nullptr) { + return false; + } + } else if (src->op == GGML_OP_VIEW) { + if (src->src[0] == nullptr || src->src[0]->view_src != nullptr) { + return false; + } + } else { return false; } } @@ -1030,18 +1046,29 @@ static bool is_msa_block_mask_expansion(const ggml_tensor * op) { return tensor_name_starts_with(src, "msa_block_mask"); } -static bool is_op_unsupported_case(const ggml_tensor * op) { +namespace { +struct ggml_openvino_op_support { + bool is_supported = true; + std::string reason; + + operator bool() const { + return is_supported; + } +}; +} // namespace + +static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) { if (is_msa_block_mask_expansion(op)) { - return true; + return {false, "MSA block mask expansion is not supported"}; } switch (op->op) { case GGML_OP_CONCAT: { if (op->type == GGML_TYPE_I64) { - return true; + return {false, "CONCAT with I64 type is not supported"}; } if (ggml_openvino_get_device_name() == "GPU" && op->type == GGML_TYPE_BF16 && has_view_op_input(op)) { - return true; + return {false, "CONCAT with BF16 type and VIEW input is not supported on GPU"}; } break; } @@ -1052,24 +1079,21 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // OpenVINO SET translation currently supports dst layouts that match src0 strides. if (op->src[0] == nullptr || nb1 != op->src[0]->nb[1] || nb2 != op->src[0]->nb[2] || nb3 != op->src[0]->nb[3]) { - // std::cout << "Unsupported SET op with dst nb1=" << nb1 << ", nb2=" << nb2 << ", nb3=" << nb3 - // << " that does not match src0 strides nb[1]=" - // << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[1]) : "null") - // << ", nb[2]=" << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[2]) : "null") - // << ", nb[3]=" << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[3]) : "null") - // << std::endl; - return true; + return {false, "SET op with dst nb1=" + std::to_string(nb1) + ", nb2=" + std::to_string(nb2) + ", nb3=" + std::to_string(nb3) + + " that does not match src0 strides nb[1]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[1]) : "null") + + ", nb[2]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[2]) : "null") + + ", nb[3]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[3]) : "null")}; } break; } case GGML_OP_GET_ROWS: case GGML_OP_SET_ROWS: { if (op->ne[3] != 1) { - return true; + return {false, "GET_ROWS/SET_ROWS with ne[3] != 1 (ne[3]=" + std::to_string(op->ne[3]) + ") is not supported"}; } if (op->op == GGML_OP_GET_ROWS && ggml_openvino_get_device_name() == "GPU" && op->src[0]->type == GGML_TYPE_BF16) { - return true; + return {false, "GET_ROWS with BF16 src0 is not supported on GPU"}; } if (op->ne[0] == 256 && (op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K || op->src[0]->type == GGML_TYPE_Q4_1 || op->src[0]->type == GGML_TYPE_Q5_1)) { @@ -1078,14 +1102,14 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // make_int8_weights/make_int4_weights: dequant is done in f16, not f32, to keep the // Convert/Subtract/Multiply chain fusable into GatherMatmulCompressed/FullyConnectedCompressed // for the shared non-test code paths). - return true; + return {false, "GET_ROWS/SET_ROWS with ne[0] == 256 and type " + std::string(ggml_type_name(op->src[0]->type)) + + " rejected due to f16-arithmetic dequant rounding errors that intermittently exceed 1e-7 NMSE threshold"}; } - break; } case GGML_OP_RESHAPE: { if (strncmp(op->name, "ffn_norm_exps", sizeof("ffn_norm_exps") - 1) == 0) { - return true; + return {false, "RESHAPE for ffn_norm_exps is not supported"}; } break; } @@ -1093,11 +1117,13 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { case GGML_OP_MUL: case GGML_OP_SUB: { if (op->src[1]->op == GGML_OP_PERMUTE) { - return true; + return {false, "ADD/MUL/SUB with PERMUTE src1 is not supported"}; } for (int i = 0; i < 4; i++) { if (op->src[0]->ne[i] != op->src[1]->ne[i] && (op->src[0]->ne[i] != 1 && op->src[1]->ne[i] != 1)) { - return true; + return {false, "ADD/MUL/SUB with incompatible broadcast shapes: src0->ne[" + std::to_string(i) + "]=" + + std::to_string(op->src[0]->ne[i]) + ", src1->ne[" + std::to_string(i) + "]=" + + std::to_string(op->src[1]->ne[i])}; } } break; @@ -1106,7 +1132,7 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // Keep support aligned with the CPU backend implementation, which only handles f32 inputs/output and i32 ids. if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32 || op->src[2]->type != GGML_TYPE_I32) { - return true; + return {false, "ADD_ID only supports F32 inputs/output and I32 ids"}; } break; } @@ -1116,14 +1142,27 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // until the fused GPU kernel is reliable. (falied case llama-arch-test mpt) if (ggml_openvino_get_device_name() == "GPU" && op->src[1]->ne[0] == op->ne[0] && op->src[1]->ne[1] == 1 && op->src[1]->ne[2] == 1 && op->src[1]->ne[3] == 1) { - return true; + return {false, "DIV per-channel scale broadcast is not supported on GPU"}; + } + break; + } + case GGML_OP_POOL_2D: { + const auto& name = ggml_openvino_get_device_name(); + if (name == "GPU") { + const int32_t * params = op->op_params; + const int k0 = params[1]; + const int k1 = params[2]; + const int p0 = params[5]; + const int p1 = params[6]; + if ((p0 > 0 || p1 > 0) && (k0 < 3 || k1 < 3)) { + return {false, "POOL_2D with padding and kernel size < 3 is not supported on " + name}; + } } break; } case GGML_OP_SUM_ROWS: { - // if the input is PERMUTE skip if (op->src[0]->op == GGML_OP_PERMUTE) { - return true; + return {false, "SUM_ROWS with PERMUTE input is not supported"}; } break; } @@ -1140,54 +1179,51 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // accuracy drift in the OpenVINO path. Restrict by scale=1.0 to avoid // affecting non-gemma3n models such as Llama-3.2. if (fabsf(scale - 1.0f) < 1e-6f && is_gemma3n_flash_attn_pattern(op)) { - return true; + return {false, "FLASH_ATTN_EXT gemma3n pattern on GPU is not supported"}; } if (op->src[4] != nullptr) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with sinks\n"); - return true; + return {false, "FLASH_ATTN_EXT with sinks is not supported"}; } if (!is_supported_flash_attn_pattern(op)) { - return true; + return {false, "FLASH_ATTN_EXT unsupported attention pattern"}; } if (max_bias > 0) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with max_bias > 0\n"); - return true; + return {false, "FLASH_ATTN_EXT with max_bias > 0 (max_bias=" + std::to_string(max_bias) + ") is not supported"}; } if (logit_softcap != 0) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with logit_softcap != 0\n"); - return true; + return {false, "FLASH_ATTN_EXT with logit_softcap != 0 (logit_softcap=" + std::to_string(logit_softcap) + ") is not supported"}; } break; } case GGML_OP_PERMUTE: { - if (op->type == GGML_TYPE_BF16) { - // err msg: [GPU] Could not find a suitable kernel for transpose - // GGML_LOG_WARN("OpenVINO backend does not support PERMUTE with BF16 type\n"); - return true; + if (op->type == GGML_TYPE_BF16 && ggml_openvino_get_device_name() == "GPU") { + return {false, "PERMUTE with BF16 type is not supported on GPU"}; } break; } case GGML_OP_CPY: { if (op->src[0]->type == GGML_TYPE_BF16 || op->src[1]->type == GGML_TYPE_BF16) { - // GGML_LOG_WARN("OpenVINO backend does not support CPY with non-contiguous data or bf16 types\n"); - return true; + return {false, "CPY with BF16 src type is not supported"}; } // CPY to a quantized destination (e.g. f32 -> q4_0) is numerically unstable with OpenVINO backend. if (ggml_is_quantized(op->type)) { - return true; + return {false, "CPY to quantized destination (e.g. f32 -> q4_0) is numerically unstable"}; } if (ggml_nelements(op->src[0]) != ggml_nelements(op->src[1])) { - return true; + return {false, "CPY with mismatched element counts is not supported: src0=" + std::to_string(ggml_nelements(op->src[0])) + + " != src1=" + std::to_string(ggml_nelements(op->src[1]))}; } // op test case with non-contiguous src or dst if ((op->ne[0] == 3 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2) || (op->ne[0] == 1 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2) || (op->ne[0] == 2 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2)) { - return true; + return {false, "CPY with non-contiguous shape [" + std::to_string(op->ne[0]) + ", " + + std::to_string(op->ne[1]) + ", " + std::to_string(op->ne[2]) + ", " + + std::to_string(op->ne[3]) + "] is not supported"}; } if (!cpy_output_view_is_supported(op)) { - return true; + return {false, "CPY with non-contiguous output view is not supported"}; } break; } @@ -1196,13 +1232,14 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { ggml_is_quantized(op->src[0]->type) && strcmp(op->src[0]->name, "a") == 0 && strcmp(op->src[1]->name, "b") == 0 && op->src[0]->ne[1] == 1 && op->src[1]->ne[1] == 64 && op->src[0]->ne[0] == 256 && op->src[1]->ne[0] == 256) { - return true; + return {false, "MUL_MAT quantized benchmark test case on GPU is not supported"}; } if (op->src[0]->ne[3] != op->src[1]->ne[3] && op->src[0]->ne[3] != 1 && op->src[1]->ne[3] != 1) { - return true; + return {false, "MUL_MAT with incompatible broadcast on ne[3]: src0->ne[3]=" + std::to_string(op->src[0]->ne[3]) + + ", src1->ne[3]=" + std::to_string(op->src[1]->ne[3])}; } if (op->src[0]->op == GGML_OP_VIEW && op->src[1]->op == GGML_OP_VIEW) { - return true; + return {false, "MUL_MAT with both inputs as VIEW is not supported"}; } break; } @@ -1210,16 +1247,17 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // Single-expert (or empty) MUL_MAT_ID is a degenerate shape that stresses GatherMatmul edge // cases and never occurs in real MoE; let it fall back to CPU. if (op->src[0] != nullptr && op->src[0]->ne[2] <= 1) { - return true; + return {false, "MUL_MAT_ID with single-expert or empty ne[2] <= 1 (ne[2]=" + + std::to_string(op->src[0]->ne[2]) + ") is not supported"}; } if (ggml_openvino_get_device_name() == "GPU" && op->src[0] != nullptr && op->src[0]->type == GGML_TYPE_BF16) { - return true; + return {false, "MUL_MAT_ID with BF16 weights on GPU is not supported"}; } // GPU MUL_MAT_ID uses a Gather+MatMul fallback because the GPU plugin rejects internal // GatherMatmul for these test shapes. Skip cases that would materialize a large selected // expert-weight temporary. if (ggml_openvino_get_device_name() == "GPU" && mul_mat_id_requires_large_tmp(op)) { - return true; + return {false, "MUL_MAT_ID requires large temporary on GPU"}; } break; } @@ -1229,51 +1267,46 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { const int mode = op_params[2]; if (op_params[15] != 0) { // FIXME: support ggml_rope_set_offset - return true; + return {false, "ggml_rope_set_offset is not supported"}; } if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_IMROPE) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with mode %d\n", mode); - return true; + return {false, "ROPE with mode " + std::to_string(mode) + " is not supported"}; } const int64_t head_dim = op->src[0]->ne[0]; const int64_t rope_dims = n_dims == 0 ? head_dim : n_dims; if (rope_dims <= 0 || rope_dims > head_dim || (rope_dims % 2) != 0) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with n_dims %d and src[0]->ne[0] %ld\n", n_dims, - // op->src[0]->ne[0]); - return true; + return {false, "ROPE with n_dims=" + std::to_string(n_dims) + ", head_dim=" + std::to_string(head_dim) + " is not supported"}; } if (op->type != GGML_TYPE_F32 && op->type != GGML_TYPE_F16) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with type %s\n", ggml_type_name(op->type)); - return true; + return {false, "ROPE with type " + std::string(ggml_type_name(op->type)) + " is not supported"}; } if (op->src[0]->op == GGML_OP_VIEW) { - if (op->src[0]->view_src->ne[1] != op->src[0]->ne[2]) { - // GGML_LOG_WARN( - // "OpenVINO backend does not support ROPE with src[0]->view_src->ne[1] %ld != src[0]->ne[2] " - // "%ld\n", - // op->src[0]->view_src->ne[1], op->src[0]->ne[2]); - return true; + const struct ggml_tensor * view = op->src[0]; + const struct ggml_tensor * view_src = view->view_src; + if (view_src->ne[1] != view->ne[1] || view_src->ne[2] != view->ne[2] || view_src->ne[3] != view->ne[3]) { + return {false, "ROPE with view_src->ne [" + std::to_string(view_src->ne[1]) + ", " + + std::to_string(view_src->ne[2]) + ", " + std::to_string(view_src->ne[3]) + + "] != view->ne [" + std::to_string(view->ne[1]) + ", " + + std::to_string(view->ne[2]) + ", " + std::to_string(view->ne[3]) + + "] is not supported"}; } } if (mode == GGML_ROPE_TYPE_IMROPE && (op->src[2] != 0 || ((const float *) op_params)[6] != 1 || ((const float *) op_params)[7] != 0 || ((const float *) op_params)[8] != 1)) { - // GGML_LOG_WARN("OpenVINO backend does not support IMROPE with freq_factors, freq_scale, ext_factor, and attn_factor\n"); - return true; + return {false, "IMROPE with freq_factors, freq_scale, ext_factor, and attn_factor is not supported"}; } break; } case GGML_OP_TRANSPOSE: { - // if the type is bf16, will return true if (op->type == GGML_TYPE_BF16) { - // GGML_LOG_WARN("OpenVINO backend does not support CONT with BF16 type\n"); - return true; + return {false, "TRANSPOSE with BF16 type is not supported"}; } break; } case GGML_OP_REPEAT: { if (ggml_openvino_get_device_name() == "GPU" && op->type == GGML_TYPE_BF16) { - return true; + return {false, "REPEAT with BF16 type is not supported on GPU"}; } break; } @@ -1285,15 +1318,15 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // return true; // } if (op->src[2]->op == GGML_OP_PERMUTE) { - return true; + return {false, "GATED_DELTA_NET with PERMUTE src2 is not supported"}; } // kda (per-key-dimension gating) not supported by fused GatedDeltaNet op if (op->src[3]->ne[0] != 1) { - return true; + return {false, "GATED_DELTA_NET with kda (per-key-dimension gating) is not supported"}; } // K > 1 (multiple state snapshots) not supported by fused op if (((const int32_t *) op->op_params)[0] > 1) { - return true; + return {false, "GATED_DELTA_NET with K > 1 (multiple state snapshots) is not supported"}; } break; } @@ -1307,17 +1340,17 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // Skip TOPK_MOE fused tests until it is fully supported. // The argsort_top_k VIEW wrapping ARGSORT is named "selected_experts" in test_topk_moe. if (strcmp(op->name, "selected_experts") == 0) { - return true; + return {false, "VIEW for selected_experts (argsort_top_k) is not supported"}; } break; } default: break; } - return false; + return {true, ""}; } -static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { +static ggml_openvino_op_support ggml_backend_openvino_device_supports_op_impl(ggml_backend_dev_t dev, const ggml_tensor * op) { GGML_ASSERT(dev->reg != nullptr); static std::unordered_set supported_types{ @@ -1367,48 +1400,41 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con case GGML_OP_UNARY: { auto supported = supported_unary_ops.find(ggml_get_unary_op(op)) != supported_unary_ops.end(); if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support unary op %s\n", ggml_unary_op_name(ggml_get_unary_op(op))); - return false; + return {false, "unary op " + std::string(ggml_unary_op_name(ggml_get_unary_op(op))) + " has no op translator"}; } if (ggml_get_unary_op(op) == GGML_UNARY_OP_EXP && op->type == GGML_TYPE_F32) { - return false; + return {false, "UNARY_EXP with F32 type is not supported"}; } break; } case GGML_OP_GLU: { auto supported = supported_glu_ops.find(ggml_get_glu_op(op)) != supported_glu_ops.end(); if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support GLU op %s\n", ggml_glu_op_name(ggml_get_glu_op(op))); - return false; + return {false, "GLU op " + std::string(ggml_glu_op_name(ggml_get_glu_op(op))) + " has no op translator"}; } // if (has_view_op_input(op)) { - // // GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n", - // // ggml_glu_op_name(ggml_get_glu_op(op))); - // return false; + // return {false, "GLU op " + std::string(ggml_glu_op_name(ggml_get_glu_op(op))) + " with view input is not supported"}; // } if (op->src[1] == nullptr && op->src[0]->ne[0] % 2 != 0) { // triggers bug in ov gpu - return false; + return {false, "GLU op with odd src0 ne[0] and null src1 is not supported"}; } break; } default: { auto supported = supported_ops.find(op->op) != supported_ops.end(); if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support op %s\n", ggml_op_name(op->op)); - return false; + return {false, "op " + std::string(ggml_op_name(op->op)) + " has no op translator"}; } static std::set ops_not_support_view_input{}; if (ops_not_support_view_input.find(op->op) != ops_not_support_view_input.end() && has_view_op_input(op)) { - // GGML_LOG_WARN("OpenVINO backend does not support op %s with view input\n", ggml_op_name(op->op)); - return false; + return {false, "op " + std::string(ggml_op_name(op->op)) + " with VIEW input is not supported"}; } } } if (supported_types.find(op->type) == supported_types.end()) { - // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(op->type)); - return false; + return {false, "tensor type " + std::string(ggml_type_name(op->type)) + " is not supported"}; } for (int i = 0; i < GGML_MAX_SRC; i++) { auto * src = op->src[i]; @@ -1416,21 +1442,32 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con break; } if (supported_types.find(src->type) == supported_types.end()) { - // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(src->type)); - return false; + return {false, "src[" + std::to_string(i) + "] type " + std::string(ggml_type_name(src->type)) + " is not supported"}; } const bool is_supported_3d_moe_expert = op->op == GGML_OP_MUL_MAT_ID && i == 0 && (src->type == GGML_TYPE_MXFP4 || src->ne[3] == 1); if (ggml_is_quantized(src->type) && src->ne[2] != 1 && !is_supported_3d_moe_expert) { - // GGML_LOG_WARN("OpenVINO backend does not support 3D quantized tensors\n"); - return false; + return {false, "3D quantized tensor for src[" + std::to_string(i) + "] is not supported"}; } } - if (is_op_unsupported_case(op)) { - return false; + auto op_support_case = is_op_supported_case(op); + if (!op_support_case.is_supported) { + return op_support_case; } - return true; + return {true, ""}; +} + +static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + auto res = ggml_backend_openvino_device_supports_op_impl(dev, op); + if (!res.is_supported) { + static const bool log_unsupported = ggml_openvino_getenv_int("GGML_OPENVINO_LOG_UNSUPPORTED_OPS") != 0; + if (log_unsupported) { + GGML_LOG_WARN("OpenVINO op unsupported: op '%s' (%s), type %s: %s\n", + op->name, ggml_op_name(op->op), ggml_type_name(op->type), res.reason.c_str()); + } + } + return res.is_supported; } static bool ggml_backend_openvino_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { diff --git a/ggml/src/ggml-openvino/openvino/op/cpy.cpp b/ggml/src/ggml-openvino/openvino/op/cpy.cpp index 5b387fc50d38..6f1e34779ac4 100644 --- a/ggml/src/ggml-openvino/openvino/op/cpy.cpp +++ b/ggml/src/ggml-openvino/openvino/op/cpy.cpp @@ -3,8 +3,11 @@ #include "../utils.h" #include +#include +#include #include -#include +#include +#include #include #include #include @@ -12,9 +15,14 @@ #include #include #include +#include #include +#include #include #include +#include +#include +#include namespace ov { namespace frontend { @@ -61,10 +69,27 @@ OutputVector translate_cpy(const NodeContext & context) { return rename_outputs_with_suffix({res}, context.get_name()); } - // Recurrent state cache writeback into a slot block of the cache. Where the block starts and - // where the copied data starts in the source are runtime inputs, so the cached model works for - // any kv head, active sequence count and token count. The result is the full updated cache. + // Recurrent state cache writeback into a slot block of the cache. Where the block starts is a + // runtime input, so the cached model works for any kv head and active sequence count. The + // result is the full updated cache. // op_case 1: gated-delta-net state, op_case 2: conv state, op_case 3: defrag remainder. + if (op_case == 3) { + // With -np 1 (and generally whenever there is no defrag remainder) this GET_ROWS gathers + // zero rows: nothing to write back, and the cache is unchanged. NPU rejects zero-size + // tensors, so short-circuit instead of building a degenerate Slice/Concat chain. + bool is_empty = false; + if (input_shape.rank().is_static()) { + for (const auto & d : input_shape) { + if (d.is_static() && d.get_length() == 0) { + is_empty = true; + break; + } + } + } + if (is_empty) { + return {context.get_input(1)}; + } + } const std::string slot_begin_name = "rs_slot_begin_" + context.get_name(); const bool slice_assign = context.has_input(slot_begin_name) && !context.is_stateful() && (op_case >= 1 && op_case <= 3); @@ -81,19 +106,49 @@ OutputVector translate_cpy(const NodeContext & context) { ov::Output begin = context.get_input(slot_begin_name); auto base = context.get_input(1); if (op_case == 1) { - // GDN packs [attn | state snapshots]; the state part runs from src_begin to the end. - auto src_begin = context.get_input("rs_src_begin_" + context.get_name()); - auto state_part = std::make_shared(context.get_input(0), src_begin, int_max, one, axis); + ov::Output state_begin; + const std::string src_begin_name = "rs_src_begin_" + context.get_name(); + if (context.has_input(src_begin_name)) { + state_begin = context.get_input(src_begin_name); + } else { + auto ssm_state_size = context.get_ssm_state_size(); + if (context.has_input("s_copy_active_slot_len")) { + auto len = context.get_input("s_copy_active_slot_len"); + auto state_rows = std::make_shared( + ov::op::v0::Constant::create(ov::element::i64, {1}, {ssm_state_size}), len); + state_begin = std::make_shared(state_rows); + } else { + state_begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {-ssm_state_size}); + } + } + auto state_part = + std::make_shared(context.get_input(0), state_begin, int_max, one, axis); src = std::make_shared(state_part, feature, false); } else if (op_case == 2) { - // conv_input is [previous conv state | new tokens]; copy the conv_kernel_size - 1 wide - // window starting at src_begin, which is the snapshot this writeback corresponds to. + // conv_input is [previous conv state | new tokens]; the snapshot is the conv_kernel_size - 1 + // columns ending at the last *valid* token. Gather (rather than Slice) keeps the output + // shape static even though the window start is a runtime value. auto window_size = (int64_t) input_shape[3].get_length(); - auto src_begin = context.get_input("rs_src_begin_" + context.get_name()); - auto src_end = std::make_shared( - src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size})); - auto window = std::make_shared(context.get_input(0), src_begin, src_end, one, - ov::op::v0::Constant::create(ov::element::i64, {1}, {3})); + ov::Output window; + auto col_axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3}); + const std::string src_begin_name = "rs_src_begin_" + context.get_name(); + if (context.has_input(src_begin_name)) { + auto src_begin = context.get_input(src_begin_name); + auto src_end = std::make_shared( + src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size})); + window = std::make_shared(context.get_input(0), src_begin, src_end, one, col_axis); + } else if (context.has_input("chunk_valid_len")) { + std::vector offsets(window_size); + std::iota(offsets.begin(), offsets.end(), 0); + auto indices = std::make_shared( + ov::op::v0::Constant::create(ov::element::i64, {(size_t) window_size}, offsets), + context.get_input("chunk_valid_len")); + window = std::make_shared(context.get_input(0), indices, col_axis); + } else { + auto window_begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {-window_size}); + window = + std::make_shared(context.get_input(0), window_begin, int_max, one, col_axis); + } const auto base_shape = base.get_partial_shape(); FRONT_END_OP_CONVERSION_CHECK(base_shape.rank().is_static() && base_shape.rank().get_length() == 4, "CPY conv state cache update requires rank-4 base cache"); @@ -157,6 +212,63 @@ OutputVector translate_cpy(const NodeContext & context) { auto input = process_view_input_new(context, 0); + if (op_case == 5 || op_case == 6) { + auto input_shape = context.get_input_shape(0); + auto output_shape = context.get_output_shape(); + auto dst_ggml_shape = context.get_view_input_ggml_shape(1, 0); + auto dst_stride = context.get_view_input_stride(1, 0); + size_t offset_bytes = context.get_view_input_offset(1, 0); + auto n_state = (int64_t) context.get_input_shape(0)[3].get_length(); + auto n_state_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_state}); + auto kv_buf = context.get_input(1); // shape {1,1,1,N} + + Output token_len_per_seq; + Output n_write_dyn; + if (context.has_input("token_len_per_seq")) { + token_len_per_seq = context.get_input("token_len_per_seq"); + n_write_dyn = std::make_shared(token_len_per_seq, n_state_c); + } else { + n_write_dyn = ov::op::v0::Constant::create(ov::element::i64, {1}, {(int64_t) dst_ggml_shape[3]}); + } + size_t elem_size = dst_stride[3]; + FRONT_END_OP_CONVERSION_CHECK(elem_size > 0, "CPY KV cache view update has invalid element size"); + int64_t start_elem = (int64_t) (offset_bytes / elem_size); + // op_case 5: decoder self-attention – write offset advances each step. + // op_case 6: encoder self-attn or cross-attn – offset fixed at compile time. + const bool is_decoder_self_attn = (op_case == 5); + auto ones_c = ov::op::v0::Constant::create(ov::element::i64, {3}, std::vector{1, 1, 1}); + auto new_shape = std::make_shared(ov::OutputVector{ones_c, n_write_dyn}, 0); + + auto reshaped = std::make_shared(input, new_shape, false); + auto data = std::make_shared(reshaped, context.get_output_type()); + // Indices [start_elem .. start_elem + n_write) on axis 3 of {1,1,1,N} + // For decoder self-attention the write offset advances each step, so compute it + // dynamically from the model inputs: start = (attention_size - token_len_per_seq) * n_state. + // For encoder self-attn and cross-attn the offset is fixed at graph-compile time. + ov::Output start; + if (is_decoder_self_attn && context.has_input("attention_size") && context.has_input("token_len_per_seq")) { + auto attention_size_in = context.get_input("attention_size"); + auto token_len_in = context.get_input("token_len_per_seq"); + auto past_tokens = std::make_shared(attention_size_in, token_len_in); + auto new_start = std::make_shared(past_tokens, n_state_c); + start = std::make_shared( + new_start, ov::op::v0::Constant::create(ov::element::i64, {1}, {start_elem})); + } else { + start = ov::op::v0::Constant::create(ov::element::i64, {1}, {start_elem}); + } + auto start_squeezed = std::make_shared(start); + auto end = std::make_shared(start_squeezed, n_write_dyn); + auto end_squeezed = std::make_shared(end); + auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto step_squeezed = std::make_shared(step); + auto indices = + std::make_shared(start_squeezed, end_squeezed, step_squeezed, ov::element::i64); + auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3}); + + auto kv_updated = std::make_shared(kv_buf, indices, data, axis); + return rename_outputs_with_suffix({kv_updated}, context.get_name()); + } + if (input_shape != output_shape) { auto new_shape = ov::op::v0::Constant::create( ov::element::i64, {static_cast(output_shape.rank().get_length())}, output_shape.to_shape()); diff --git a/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp b/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp index 582df0130b59..06547f3d2968 100644 --- a/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp +++ b/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp @@ -3,8 +3,8 @@ #include "../utils.h" #include "ggml-openvino/ggml-openvino-extra.h" +#include #include -#include #include #include #include @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -24,13 +25,62 @@ namespace ov { namespace frontend { namespace ggml { namespace op { +static ov::Output reshape_flat_kv(const ov::Output & kv_flat, + size_t view_offset_bytes, + size_t nb1_bytes, + int64_t n_head, + int64_t head_size, + const ov::Output & attention_size) { + int64_t n_state = n_head * head_size; + int64_t layer_start_elem = (int64_t) (view_offset_bytes / (nb1_bytes / n_state)); + // Dynamic slice: [layer_start_elem, layer_start_elem + n_kv * n_state) + auto start_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {layer_start_elem}); + auto n_state_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_state}); + // end = start + attention_size * n_state (both static + dynamic) + auto kv_len_elems = std::make_shared(attention_size, n_state_c); + auto end_c = std::make_shared(start_c, kv_len_elems); + auto step_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto axis_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {3}); + auto sliced = std::make_shared(kv_flat, start_c, end_c, step_c, axis_c); + + // KV cache is laid out as {n_kv, n_head, head_size} in memory + // Reshape to {1, n_kv, n_head, head_size}, then transpose to {1, n_head, n_kv, head_size} + // as required by SDPA. + auto one_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto n_head_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_head}); + auto head_size_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {head_size}); + // reshape: {n_kv*n_state} -> {1, n_kv, n_head, head_size} + auto new_shape = + std::make_shared(ov::OutputVector{one_c, attention_size, n_head_c, head_size_c}, 0); + auto reshaped = std::make_shared(sliced, new_shape, false); + // transpose: {1, n_kv, n_head, head_size} -> {1, n_head, n_kv, head_size} + auto perm = ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}); + auto ret = std::make_shared(reshaped, perm); + return ret; +} OutputVector translate_flash_attn_ext(const NodeContext & context) { - num_inputs_check(context, 4, 4); + num_inputs_check(context, 3, 4); + const bool has_mask = context.get_input_size() == 4; auto q_f32 = context.get_input(0); auto k = context.get_input(1); auto v = context.get_input(2); - auto mask = context.get_input(3); + const int op_case = context.get_op_case(); + + if (op_case == 1 || op_case == 2) { + int64_t n_state_head = (int64_t) context.get_view_input_ggml_shape(1, 0)[3]; + int64_t n_head = (int64_t) context.get_view_input_ggml_shape(1, 0)[1]; + size_t nb1 = context.get_view_input_stride(1, 0)[2]; + size_t offset = context.get_view_input_offset(1, 0); + ov::Output attention_size; + if (op_case == 1) { + attention_size = context.get_input("attention_size"); + } else { + attention_size = context.get_input("attention_size_static"); + } + k = reshape_flat_kv(k, offset, nb1, n_head, n_state_head, attention_size); + v = reshape_flat_kv(v, offset, nb1, n_head, n_state_head, attention_size); + } float * params = reinterpret_cast(context.get_output_op_params()); float scale = params[0]; @@ -43,16 +93,19 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) { ov::Output res; // For stateful - std::string mask_name = "KQ_mask_sliced"; - if (context.get_input_names()[3].find("swa") != std::string::npos) { - mask_name = "KQ_mask_swa_sliced"; - } - if (context.has_input(mask_name)) { - mask = context.get_input(mask_name); - } - - if (mask.get_element_type() != ov::element::f16) { - mask = std::make_shared(mask, ov::element::f16); + ov::Output mask; + if (has_mask) { + mask = context.get_input(3); + std::string mask_name = "KQ_mask_sliced"; + if (context.get_input_names()[3].find("swa") != std::string::npos) { + mask_name = "KQ_mask_swa_sliced"; + } + if (context.has_input(mask_name)) { + mask = context.get_input(mask_name); + } + if (mask.get_element_type() != ov::element::f16) { + mask = std::make_shared(mask, ov::element::f16); + } } //auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output kv) { @@ -108,10 +161,14 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) { // get [B, 1, 1, S_q, S_k], which NUMPY-broadcasts cleanly against the // [B, num_heads_kv, factor, S_q, S_k] scores: B==B, then 1→num_heads_kv and // 1→factor on the head dims. - auto mask_unsq1 = - std::make_shared(mask, ov::op::v0::Constant::create(ov::element::i64, {1}, {2})); - // mask_unsq1: [B, 1, 1, S_q, S_k] (rank 5) - ov::Output qk_masked = std::make_shared(qk_scaled, mask_unsq1); + ov::Output qk_masked; + if (has_mask) { + auto mask_unsq1 = + std::make_shared(mask, ov::op::v0::Constant::create(ov::element::i64, {1}, {2})); + qk_masked = std::make_shared(qk_scaled, mask_unsq1); + } else { + qk_masked = qk_scaled; + } auto softmax = std::make_shared(qk_masked, /*axis=*/-1); @@ -164,9 +221,16 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) { k = tile_kv(num_heads, num_heads_kv, head_size, k); v = tile_kv(num_heads, num_heads_kv, head_size, v); - auto sdpa = std::make_shared(q, k, v, mask, scale_node, false); - res = std::make_shared(sdpa, - ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3})); + constexpr auto causal = false; + if (has_mask) { + auto sdpa = std::make_shared(q, k, v, mask, scale_node, causal); + res = std::make_shared( + sdpa, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3})); + } else { + auto sdpa = std::make_shared(q, k, v, scale_node, causal); + res = std::make_shared( + sdpa, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3})); + } res = std::make_shared(res, ov::element::f32); return rename_outputs_with_suffix({res}, context.get_name()); } diff --git a/ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp b/ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp index 66c748283311..07eeb3c8fd6d 100644 --- a/ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp +++ b/ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp @@ -7,12 +7,15 @@ #include #include #include +#include #include #include #include #include +#include #include #include +#include #include #include #include @@ -80,6 +83,28 @@ OutputVector translate_gated_delta_net(const NodeContext & context) { g = std::make_shared(g, ov::op::v0::Constant::create(ov::element::i64, {1}, {3})); beta = std::make_shared(beta, ov::op::v0::Constant::create(ov::element::i64, {1}, {3})); + if (context.has_input("chunk_valid_len")) { + // The last prefill chunk is padded with fabricated tokens. The recurrence is + // S_t = S_{t-1} * exp(g_t) + k_t (x) ((v_t - S_{t-1}^T k_t) * beta_t) + // so forcing g = 0 and beta = 0 makes a padded step an exact identity and keeps the final + // state equal to the state after the last real token. Attention output at those positions + // is garbage but never read. + const auto & g_ps = g.get_partial_shape(); + FRONT_END_OP_CONVERSION_CHECK(g_ps.rank().is_static() && g_ps.rank().get_length() == 3 && g_ps[1].is_static(), + "GATED_DELTA_NET pad masking requires a static token dimension"); + const int64_t n_tokens = g_ps[1].get_length(); + std::vector positions(n_tokens); + std::iota(positions.begin(), positions.end(), 0); + auto valid = std::make_shared( + ov::op::v0::Constant::create(ov::element::i64, {(size_t) n_tokens}, positions), + context.get_input("chunk_valid_len")); + auto mask = std::make_shared( + std::make_shared(valid, g.get_element_type()), + ov::op::v0::Constant::create(ov::element::i64, {2}, std::vector{0, 2})); + g = std::make_shared(g, mask); + beta = std::make_shared(beta, mask); + } + // std::cout << "GatedDeltaNet input shapes: q=" << q.get_partial_shape() << ", k=" << k.get_partial_shape() // << ", v=" << v.get_partial_shape() << ", g=" << g.get_partial_shape() // << ", beta=" << beta.get_partial_shape() << ", state=" << state.get_partial_shape() << std::endl; diff --git a/ggml/src/ggml-openvino/openvino/op/glu_geglu_quick.cpp b/ggml/src/ggml-openvino/openvino/op/glu_geglu_quick.cpp new file mode 100644 index 000000000000..c6d64aed43aa --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/op/glu_geglu_quick.cpp @@ -0,0 +1,64 @@ +#include "../node_context.h" +#include "../op_table.h" +#include "../utils.h" + +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_glu_geglu_quick(const NodeContext & context) { + num_inputs_check(context, 1, 2); + + ov::Output src0; + ov::Output src1; + if (context.get_input_size() == 2) { + src0 = process_view_input_new(context, 0); + src1 = process_view_input_new(context, 1); + } else { + // split along last axis, nc = ne[0] / 2 + auto combined = process_view_input_new(context, 0); + auto combined_shape = combined.get_partial_shape(); + int64_t last_dim_val = combined_shape[combined_shape.rank().get_length() - 1].get_length(); + int64_t nc = last_dim_val / 2; + + auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); + auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto start0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto stop0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); + auto start1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); + auto stop1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {2 * nc}); + + src0 = std::make_shared(combined, start0, stop0, step, axis); + src1 = std::make_shared(combined, start1, stop1, step, axis); + } + + int32_t * params = context.get_output_op_params(); + const int32_t swapped = params[1]; + if (swapped) { + std::swap(src0, src1); + } + + // GELU_QUICK(x) = x * sigmoid(1.702 * x) + // Create the constant in the same type as src0 to avoid f16/f32 mismatch. + auto input_type = src0.get_element_type(); + auto coef = ov::op::v0::Constant::create(input_type, ov::Shape{}, {1.702f}); + auto scaled = std::make_shared(src0, coef); + auto sigmoid = std::make_shared(scaled); + auto gated = std::make_shared(src0, sigmoid); + auto res = std::make_shared(gated, src1); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/pool_2d.cpp b/ggml/src/ggml-openvino/openvino/op/pool_2d.cpp new file mode 100644 index 000000000000..fb6333175f02 --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/op/pool_2d.cpp @@ -0,0 +1,53 @@ +#include "../node_context.h" +#include "../op_table.h" +#include "../utils.h" + +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_pool_2d(const NodeContext & context) { + num_inputs_check(context, 1, 1); + const int32_t * params = context.get_output_op_params(); + + const int k0 = params[1]; + const int k1 = params[2]; + const int s0 = params[3]; + const int s1 = params[4]; + const int p0 = params[5]; + const int p1 = params[6]; + + const int op_case = context.get_op_case(); + ov::Output input = context.get_input(0); + ov::Strides strides{static_cast(s1), static_cast(s0)}; + ov::Shape pads_begin{static_cast(p1), static_cast(p0)}; + ov::Shape pads_end{static_cast(p1), static_cast(p0)}; + ov::Shape kernel{static_cast(k1), static_cast(k0)}; + ov::Output res; + + switch (op_case) { + case 1: // GGML_OP_POOL_MAX + { + res = std::make_shared(input, strides, pads_begin, pads_end, kernel); + break; + } + case 2: // GGML_OP_POOL_AVG + { + res = std::make_shared(input, strides, pads_begin, pads_end, kernel, false); + break; + } + default: + break; + } + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/roll.cpp b/ggml/src/ggml-openvino/openvino/op/roll.cpp new file mode 100644 index 000000000000..e8d1b8e50b34 --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/op/roll.cpp @@ -0,0 +1,36 @@ +#include "../node_context.h" +#include "../op_table.h" +#include "../utils.h" + +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_roll(const NodeContext & context) { + num_inputs_check(context, 1, 1); + const int32_t * params = context.get_output_op_params(); + + int64_t s0 = params[0]; + int64_t s1 = params[1]; + int64_t s2 = params[2]; + int64_t s3 = params[3]; + + auto input = context.get_input(0); + + auto shift = ov::op::v0::Constant::create( + ov::element::i64, ov::Shape{4}, std::vector{s3, s2, s1, s0}); + auto axes = ov::op::v0::Constant::create( + ov::element::i64, ov::Shape{4}, std::vector{0, 1, 2, 3}); + + auto roll = std::make_shared(input, shift, axes); + return rename_outputs_with_suffix({roll}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/view.cpp b/ggml/src/ggml-openvino/openvino/op/view.cpp index 138526cb49c6..56f5ceec9bb0 100644 --- a/ggml/src/ggml-openvino/openvino/op/view.cpp +++ b/ggml/src/ggml-openvino/openvino/op/view.cpp @@ -17,6 +17,13 @@ namespace op { OutputVector translate_view(const NodeContext & context) { num_inputs_check(context, 1, 1); + if (context.get_op_case() == 1) { + // Static-mode identity pass-through for VIEWs over a GATED_DELTA_NET combined output or + // the conv_input CONCAT; the consuming op (CPY/RMS_NORM) does its own runtime-correct + // slicing on the full tensor (see ggml-decoder.cpp compute_op_case, GGML_OP_VIEW). + return {context.get_input(0)}; + } + if (!context.is_static()) { // On the stateless/non-static path VIEW is normally a no-op (consumers re-slice). // EXCEPTION: the MoE expert aggregation slices each expert plane out of diff --git a/ggml/src/ggml-openvino/openvino/op_table.cpp b/ggml/src/ggml-openvino/openvino/op_table.cpp index 3c26fe83b1ad..9c9d8eeac781 100644 --- a/ggml/src/ggml-openvino/openvino/op_table.cpp +++ b/ggml/src/ggml-openvino/openvino/op_table.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -55,10 +56,12 @@ std::unordered_map get_supported_ops() { {"GGML_UNARY_OP_SIGMOID", op::translate_1to1_match_1_input }, {"GGML_UNARY_OP_EXP", op::translate_1to1_match_1_input }, {"GGML_UNARY_OP_NEG", op::translate_1to1_match_1_input }, + {"GGML_UNARY_OP_RELU", op::translate_1to1_match_1_input }, {"GGML_OP_VIEW", op::translate_view }, {"GGML_GLU_OP_SWIGLU", op::translate_glu_swiglu }, {"GGML_GLU_OP_SWIGLU_OAI", op::translate_glu_swiglu_oai }, {"GGML_GLU_OP_GEGLU", op::translate_glu_geglu }, + {"GGML_GLU_OP_GEGLU_QUICK", op::translate_glu_geglu_quick }, {"GGML_OP_SET_ROWS", op::translate_set_rows }, {"GGML_OP_CPY", op::translate_cpy }, {"GGML_OP_FLASH_ATTN_EXT", op::translate_flash_attn_ext }, @@ -72,6 +75,8 @@ std::unordered_map get_supported_ops() { {"GGML_OP_DIAG", op::translate_diag }, {"GGML_OP_TRI", op::translate_tri }, {"GGML_OP_SET", op::translate_set }, + {"GGML_OP_POOL_2D", op::translate_pool_2d }, + {"GGML_OP_ROLL", op::translate_roll }, // solve_tri has accuracy issues on GPU // {"GGML_OP_SOLVE_TRI", op::translate_solve_tri }, }; diff --git a/ggml/src/ggml-openvino/openvino/op_table.h b/ggml/src/ggml-openvino/openvino/op_table.h index d4b9292d6377..0a81a57a6677 100644 --- a/ggml/src/ggml-openvino/openvino/op_table.h +++ b/ggml/src/ggml-openvino/openvino/op_table.h @@ -38,6 +38,7 @@ GGML_OP_CONVERTER(translate_view); GGML_OP_CONVERTER(translate_glu_swiglu); GGML_OP_CONVERTER(translate_glu_swiglu_oai); GGML_OP_CONVERTER(translate_glu_geglu); +GGML_OP_CONVERTER(translate_glu_geglu_quick); GGML_OP_CONVERTER(translate_set_rows); GGML_OP_CONVERTER(translate_cpy); GGML_OP_CONVERTER(translate_argsort); @@ -53,6 +54,8 @@ GGML_OP_CONVERTER(translate_set); GGML_OP_CONVERTER(translate_diag); GGML_OP_CONVERTER(translate_tri); GGML_OP_CONVERTER(translate_solve_tri); +GGML_OP_CONVERTER(translate_pool_2d); +GGML_OP_CONVERTER(translate_roll); } // namespace op diff --git a/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.cpp b/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.cpp new file mode 100644 index 000000000000..21801c0f3992 --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.cpp @@ -0,0 +1,212 @@ +#include "fuse_to_conv.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace opp = ov::pass::pattern; + +namespace ov { +namespace frontend { +namespace ggml { +namespace pass { + +// This pass fuses an IM2COL + MatMul convolution into OpenVINO's Convolution op for performance gains. +// Reference the im2col.cpp translator for reference on the pattern being matched. + +FuseToConv::FuseToConv() { + const auto m_wei = opp::any_input(); + const auto m_act = opp::any_input(); + const auto m_matmul = opp::wrap_type({m_wei, m_act}); + + const auto callback = [=](ov::pass::pattern::Matcher & m) { + const auto & pm = m.get_pattern_value_map(); + + auto matmul_node = ov::as_type_ptr(pm.at(m_matmul).get_node_shared_ptr()); + if (!matmul_node || matmul_node->get_transpose_a() || !matmul_node->get_transpose_b()) { + return false; + } + + auto trace = matmul_node->input_value(1); + + // Optional Convert + if (auto n = ov::as_type_ptr(trace.get_node_shared_ptr())) { + trace = n->input_value(0); + } + + for (int i = 0; i < 2; ++i) { + auto n = ov::as_type_ptr(trace.get_node_shared_ptr()); + if (!n) { + return false; + } + trace = n->input_value(0); + } + + if (auto n = ov::as_type_ptr(trace.get_node_shared_ptr())) { + trace = n->input_value(0); + } else { + return false; + } + + if (auto n = ov::as_type_ptr(trace.get_node_shared_ptr())) { + trace = n->input_value(0); + } else { + return false; + } + + if (auto n = ov::as_type_ptr(trace.get_node_shared_ptr())) { + trace = n->input_value(0); + } else { + return false; + } + + auto eip = ov::as_type_ptr(trace.get_node_shared_ptr()); + if (!eip) { + return false; + } + const auto eip_strides = eip->get_strides(); // {stride_h, stride_w} + const auto eip_rates = eip->get_rates(); // {dil_h, dil_w} + + auto pad = ov::as_type_ptr(eip->input_value(0).get_node_shared_ptr()); + if (!pad) { + return false; + } + auto pads_begin_const = + ov::as_type_ptr(pad->input_value(1).get_node_shared_ptr()); + + const auto pads_begin_vals = pads_begin_const->cast_vector(); // {0, 0, pad_h, pad_w} + const std::ptrdiff_t pad_h = static_cast(pads_begin_vals[2]); + const std::ptrdiff_t pad_w = static_cast(pads_begin_vals[3]); + + auto image_input = pad->input_value(0); // [N, IC, 1, IW] NCHW + + auto w_trace = matmul_node->input_value(0); + if (auto n = ov::as_type_ptr(w_trace.get_node_shared_ptr())) { + w_trace = n->input_value(0); + } + for (int i = 0; i < 2; ++i) { + auto n = ov::as_type_ptr(w_trace.get_node_shared_ptr()); + if (!n) { + break; + } + w_trace = n->input_value(0); + } + + auto weight_const = ov::as_type_ptr(w_trace.get_node_shared_ptr()); + if (!weight_const) { + return false; + } + + // Reshape weight to [OC, IC, 1, KW] (OIHW). + const auto w_shape = weight_const->get_shape(); + ov::Shape conv_w_shape; + if (w_shape.size() == 3) { + conv_w_shape = {w_shape[0], w_shape[1], 1, w_shape[2]}; + } else if (w_shape.size() == 4) { + conv_w_shape = {w_shape[1], w_shape[2], 1, w_shape[3]}; + } else { + return false; + } + + auto weight_reshaped = register_new_node(weight_const->get_element_type(), conv_w_shape, + weight_const->get_data_ptr()); + + ov::Output weight_input = weight_reshaped; + if (weight_reshaped->get_element_type() != image_input.get_element_type()) { + weight_input = register_new_node(weight_reshaped, image_input.get_element_type()); + } + + auto conv = register_new_node( + image_input, weight_input, + ov::Strides{static_cast(eip_strides[0]), static_cast(eip_strides[1])}, + ov::CoordinateDiff{pad_h, pad_w}, ov::CoordinateDiff{pad_h, pad_w}, + ov::Strides{static_cast(eip_rates[0]), static_cast(eip_rates[1])}, + ov::op::PadType::EXPLICIT); + + constexpr auto target_type = ov::element::f32; + ov::Output conv_out = conv; + if (conv_out.get_element_type() != target_type) { + conv_out = register_new_node(conv_out, target_type); + } + + std::shared_ptr add_node; + ov::Output bias_input; + for (const auto & consumer_in : matmul_node->output(0).get_target_inputs()) { + auto cast = ov::as_type_ptr(consumer_in.get_node()->shared_from_this()); + if (!cast) { + continue; + } + for (const auto & add_in : cast->output(0).get_target_inputs()) { + auto add = ov::as_type_ptr(add_in.get_node()->shared_from_this()); + if (!add) { + continue; + } + for (size_t i = 0; i < 2; ++i) { + if (ov::as_type_ptr(add->input_value(i).get_node_shared_ptr())) { + bias_input = add->input_value(i); + add_node = add; + break; + } + } + if (add_node) { + break; + } + } + if (add_node) { + break; + } + } + + ov::Output final_out; + std::shared_ptr target_node; + + if (add_node) { + // Reshape bias [OC, 1] → [1, OC, 1, 1] for NCHW broadcasting. + ov::Output bias = bias_input; + if (bias.get_element_type() != target_type) { + bias = register_new_node(bias, target_type); + } + const auto oc = static_cast(conv_w_shape[0]); + auto bias_shape = register_new_node(ov::element::i64, ov::Shape{4}, + std::vector{1, oc, 1, 1}); + bias = register_new_node(bias, bias_shape, false); + final_out = register_new_node(conv_out, bias); + target_node = add_node; + } else { + final_out = conv_out; + target_node = matmul_node; + } + + // Reshape final output back to the target node's original shape if needed. + auto orig_shape = target_node->get_output_partial_shape(0); + if (orig_shape.is_static() && final_out.get_partial_shape() != orig_shape) { + auto shape_const = register_new_node(ov::element::i64, ov::Shape{orig_shape.size()}, + orig_shape.to_shape()); + final_out = register_new_node(final_out, shape_const, false); + } + + final_out.get_node_shared_ptr()->set_friendly_name(target_node->get_friendly_name()); + ov::copy_runtime_info(m.get_matched_nodes(), final_out.get_node_shared_ptr()); + ov::replace_node(target_node, final_out.get_node_shared_ptr()); + + return true; + }; + + register_matcher(std::make_shared(m_matmul, "ov::frontend::ggml::pass::FuseToConv"), callback); +} + +} // namespace pass +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.h b/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.h new file mode 100644 index 000000000000..feac14b13ff2 --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.h @@ -0,0 +1,17 @@ +#include "openvino/pass/matcher_pass.hpp" + +namespace ov { +namespace frontend { +namespace ggml { +namespace pass { + +class FuseToConv : public ov::pass::MatcherPass { +public: + OPENVINO_MATCHER_PASS_RTTI("ov::frontend::ggml::pass::FuseToConv") + FuseToConv(); +}; + +} // namespace pass +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/translate_session.cpp b/ggml/src/ggml-openvino/openvino/translate_session.cpp index 35598aba6be8..df3a72f3286c 100644 --- a/ggml/src/ggml-openvino/openvino/translate_session.cpp +++ b/ggml/src/ggml-openvino/openvino/translate_session.cpp @@ -5,6 +5,7 @@ #include "ggml-openvino/openvino/node_context.h" #include "ggml-openvino/openvino/utils.h" #include "input_model.h" +#include "pass/fuse_to_conv.h" #include "pass/mark_decompression_convert_constant_folding.h" #include "pass/mark_dequantization_subgraph.h" #include "pass/squeeze_matmul.h" @@ -109,7 +110,8 @@ ov::pass::MakeStateful::ParamResPairs get_kv_param_res_pairs( void add_sliced_mask_stateful(TensorMap & tensor_map) { auto create_sliced_mask = [&](const std::string & mask_name, const std::string & sliced_name) { if ((tensor_map.find(mask_name) != tensor_map.end()) && - (tensor_map.find("token_len_per_seq") != tensor_map.end())) { + (tensor_map.find("token_len_per_seq") != tensor_map.end()) && + (tensor_map.find("inp_pos") != tensor_map.end())) { auto token_len_per_seq = tensor_map.at("token_len_per_seq").get_node_shared_ptr(); auto mask = tensor_map.at(mask_name).get_node_shared_ptr(); std::shared_ptr mask_sliced = mask; @@ -137,6 +139,7 @@ void add_sliced_mask_stateful(TensorMap & tensor_map) { }; create_sliced_mask("self_kq_mask", "KQ_mask_sliced"); + create_sliced_mask("KQ_mask", "KQ_mask_sliced"); create_sliced_mask("self_kq_mask_swa", "KQ_mask_swa_sliced"); } @@ -395,6 +398,7 @@ std::shared_ptr TranslateSession::apply_transformations(std::shared_ptr( std::vector{ov::element::u8, ov::element::i8, ov::element::u4, ov::element::i4}); + manager.register_pass(); if (ggml_model_decoder->is_stateful()) { const auto kv_param_res_names = ggml_model_decoder->get_kv_param_res_names(); diff --git a/ggml/src/ggml-openvino/openvino/utils.cpp b/ggml/src/ggml-openvino/openvino/utils.cpp index 504d74b70679..8bb7678ee381 100644 --- a/ggml/src/ggml-openvino/openvino/utils.cpp +++ b/ggml/src/ggml-openvino/openvino/utils.cpp @@ -72,6 +72,7 @@ OutputVector rename_outputs_with_suffix(const OutputVector & outputs, const std: name += "_"; name += suffix; node->set_friendly_name(name); + // Uncomment to dump every node's inferred shape (used to hunt down dynamic dims on NPU). // std::cout << name << " " << output.get_partial_shape() << std::endl; } return outputs; diff --git a/ggml/src/ggml-openvino/utils.cpp b/ggml/src/ggml-openvino/utils.cpp index 4df8381dcbd9..93b1ccbe9075 100644 --- a/ggml/src/ggml-openvino/utils.cpp +++ b/ggml/src/ggml-openvino/utils.cpp @@ -16,6 +16,8 @@ #include #include #include +#include +#include #include #include #include @@ -48,7 +50,7 @@ enum ggml_status ov_graph_compute(ggml_cgraph * cgraph, ggml_backend_t backend) GgmlOvDecoder::dump_cgraph(cgraph, filename); } - const auto is_static = ggml_openvino_is_npu(); + const auto is_static = ggml_openvino_is_npu() || ggml_openvino_getenv_int("GGML_OPENVINO_FORCE_STATIC"); GGML_ASSERT(ctx->runtime_context != nullptr); std::shared_ptr r_ctx = std::static_pointer_cast(ctx->runtime_context); @@ -168,13 +170,24 @@ ov::Tensor create_ov_output_tensor(std::shared_ptr ggml_decoder, auto output_type = ggml_decoder->get_ov_type(ggml_tensor); ov::Shape output_shape; + void * output_data = ggml_tensor->data; if (ggml_decoder->is_static()) { output_shape = infer_request->get_output_tensor(output_index).get_shape(); } else { - output_shape = ggml_decoder->get_shape(ggml_tensor); + // For a CPY into a padded view_src (e.g. a padded KV cache buffer), the + // OV ScatterUpdate node outputs the full view_src shape, not the CPY node's + // own (smaller) shape. Using the CPY shape here causes set_output_tensor to + // fail with a shape-incompatibility error. Use view_src's shape and data + // pointer instead so the OV tensor matches the model output exactly. + if (ggml_tensor->op == GGML_OP_CPY && ggml_tensor->view_src != nullptr && + ggml_nbytes(ggml_tensor) != ggml_nbytes(ggml_tensor->view_src)) { + output_shape = ggml_decoder->get_shape(ggml_tensor->view_src); + output_data = ggml_tensor->view_src->data; + } else { + output_shape = ggml_decoder->get_shape(ggml_tensor); + } } - - ov::Tensor output_tensor(output_type, output_shape, ggml_tensor->data); + ov::Tensor output_tensor(output_type, output_shape, output_data); return output_tensor; } @@ -583,7 +596,9 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr(ggml_decoder_prefill); - auto input_model_decode = std::make_shared(ggml_decoder_decode); - - auto model_prefill = ov::frontend::ggml::FrontEnd::convert(input_model_prefill); - ggml_decoder_prefill->clear_model_weights(); - auto model_decode = ov::frontend::ggml::FrontEnd::convert(input_model_decode); - ggml_decoder_decode->clear_model_weights(); - conversion_end_time = ggml_time_us(); - - if (ggml_openvino_getenv_int("GGML_OPENVINO_DUMP_IR")) { - char timestamped_filename[64]; - auto timestamp = (long long) ggml_time_us(); - snprintf(timestamped_filename, sizeof(timestamped_filename), "model_prefill_%lld.xml", timestamp); - ov::serialize(model_prefill, timestamped_filename); - snprintf(timestamped_filename, sizeof(timestamped_filename), "model_decode_%lld.xml", timestamp); - ov::serialize(model_decode, timestamped_filename); - } + const bool dump_ir = ggml_openvino_getenv_int("GGML_OPENVINO_DUMP_IR"); + const auto dump_ir_timestamp = static_cast(ggml_time_us()); + + auto build_static_model = [&core, &config, dump_ir, dump_ir_timestamp]( + std::shared_ptr decoder, + const char * tag, + std::shared_ptr & model, + ov::CompiledModel & compiled_model, + std::shared_ptr & infer_request, + int64_t & local_conversion_end_time, + int64_t & local_compile_end_time) { + auto input_model = std::make_shared(decoder); + model = ov::frontend::ggml::FrontEnd::convert(input_model); + decoder->clear_model_weights(); + local_conversion_end_time = ggml_time_us(); + + if (dump_ir) { + char timestamped_filename[64]; + snprintf(timestamped_filename, sizeof(timestamped_filename), "model_%s_%lld.xml", tag, + dump_ir_timestamp); + ov::serialize(model, timestamped_filename); + } + compiled_model = core.compile_model(model, device, config); + infer_request = std::make_shared(compiled_model.create_infer_request()); + local_compile_end_time = ggml_time_us(); + }; + std::shared_ptr model_prefill; + std::shared_ptr model_decode; ov::CompiledModel compiled_model_prefill; ov::CompiledModel compiled_model_decode; - auto remote_context = ggml_openvino_get_remote_context(); - if (remote_context.has_value()) { - compiled_model_prefill = core.compile_model(model_prefill, remote_context.value(), config); - compiled_model_decode = core.compile_model(model_decode, remote_context.value(), config); - } else { - compiled_model_prefill = core.compile_model(model_prefill, device, config); - compiled_model_decode = core.compile_model(model_decode, device, config); - } - - auto infer_request_prefill = std::make_shared(compiled_model_prefill.create_infer_request()); - auto infer_request_decode = std::make_shared(compiled_model_decode.create_infer_request()); - compile_end_time = ggml_time_us(); + std::shared_ptr infer_request_prefill; + std::shared_ptr infer_request_decode; + int64_t prefill_conversion_end_time; + int64_t decode_conversion_end_time; + int64_t prefill_compile_end_time; + int64_t decode_compile_end_time; + auto prefill_future = std::async(std::launch::async, build_static_model, ggml_decoder_prefill, "prefill", + std::ref(model_prefill), std::ref(compiled_model_prefill), + std::ref(infer_request_prefill), std::ref(prefill_conversion_end_time), + std::ref(prefill_compile_end_time)); + auto decode_future = std::async(std::launch::async, build_static_model, ggml_decoder_decode, "decode", + std::ref(model_decode), std::ref(compiled_model_decode), + std::ref(infer_request_decode), std::ref(decode_conversion_end_time), + std::ref(decode_compile_end_time)); + prefill_future.get(); + decode_future.get(); + conversion_end_time = std::max(prefill_conversion_end_time, decode_conversion_end_time); + compile_end_time = std::max(prefill_compile_end_time, decode_compile_end_time); model = is_prefill ? model_prefill : model_decode; ggml_decoder = is_prefill ? ggml_decoder_prefill : ggml_decoder_decode; @@ -742,7 +774,7 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptrne[0]; + auto inp_len = get_inp_pos_n_tokens(cgraph, inp_pos); for (int chunk_index = 0; chunk_index * prefill_chunk_size < inp_len; chunk_index++) { for (size_t i = 0; i < ov_input_names_local.size(); i++) { auto param_name = ov_input_names_local[i]; @@ -762,6 +794,11 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptrsecond; + if (ggml_nbytes(ggml_tensor) == 0) { + // Zero-row in-place writeback (e.g. the empty s_copy defrag remainder). The OV + // Result is the full cache, so binding it over this 0-byte buffer overflows it. + continue; + } auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); infer_request->set_output_tensor(i, output_tensor); } @@ -798,6 +835,9 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptrsecond; + if (ggml_nbytes(ggml_tensor) == 0) { + continue; + } auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); infer_request->set_output_tensor(i, output_tensor); } @@ -1074,6 +1114,9 @@ ov::Tensor get_ov_input_tensor(std::shared_ptr ggml_decoder, cons ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr ggml_decoder, const std::string & param_name) { // NPU decoding stage + if (ggml_decoder->get_model_extra_inputs().count(param_name)) { + return get_ov_input_tensor(ggml_decoder, param_name); + } const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); @@ -1123,14 +1166,30 @@ ov::Tensor get_ov_input_tensor_static_prefill(std::shared_ptr ggm const std::string & param_name, int chunk_index) { // NPU prompt processing stage - const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); - const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); - const size_t input_len = ggml_decoder->get_input_len(); const size_t chunk_size = ggml_decoder->m_prefill_chunk_size; const size_t chunk_valid_size = std::min(chunk_size, input_len - chunk_index * chunk_size); const size_t chunk_pad_size = chunk_size - chunk_valid_size; + if (param_name == "chunk_valid_len") { + ov::Tensor input_tensor(ov::element::i64, ov::Shape{1}); + *input_tensor.data() = (int64_t) chunk_valid_size; + return input_tensor; + } + if (chunk_index > 0 && param_name == "cache_rs_reset_len") { + // The recurrent-state clear belongs to the start of the sequence. Re-applying it on every + // chunk would wipe the state accumulated by the preceding chunks, so disable it (a zero + // length makes scale.cpp's keep-mask select every slot) after the first chunk. + ov::Tensor input_tensor(ov::element::i64, ov::Shape{1}); + *input_tensor.data() = 0; + return input_tensor; + } + if (ggml_decoder->get_model_extra_inputs().count(param_name)) { + return get_ov_input_tensor(ggml_decoder, param_name); + } + const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); + const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); + if (GgmlOvDecoder::is_inp_pos(ggml_tensor, op) && GgmlOvDecoder::get_inp_pos_n_planes(op) > 1) { // IMROPE: inp_pos stacks n_planes (t/h/w/e) position planes, each of length // input_len; pad every plane independently so they stay aligned to chunk_size. @@ -1306,7 +1365,7 @@ void print_input_tensor_info(const std::string & name, const ov::Tensor & tensor << std::endl; switch (tensor.get_element_type()) { case ov::element::f32: { - if (name.find("self_kq_mask") == std::string::npos) { + if (name.find("self_kq_mask") == std::string::npos && name.find("KQ_mask") == std::string::npos) { std::cout << *(tensor.data()) << std::endl; } else { size_t rows = tensor.get_shape()[2]; @@ -1414,8 +1473,24 @@ const ggml_tensor * get_inp_pos_tensor(ggml_cgraph * cgraph) { throw std::runtime_error("get_inp_pos_tensor: inp_pos not found in cgraph"); } -bool get_is_prefill(const ggml_tensor * inp_pos) { - return inp_pos->ne[0] > 1; +int64_t get_inp_pos_n_tokens(ggml_cgraph * cgraph, const ggml_tensor * inp_pos) { + // IMROPE stacks n_planes (t/h/w/e) position planes into inp_pos, so ne[0] is + // n_planes * n_tokens. Callers that need a token count must divide the planes out. + int n_planes = 1; + for (int i = 0; i < cgraph->n_nodes; ++i) { + auto * op = cgraph->nodes[i]; + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (op->src[j] == inp_pos) { + n_planes = GgmlOvDecoder::get_inp_pos_n_planes(op); + break; + } + } + } + return inp_pos->ne[0] / n_planes; +} + +bool get_is_prefill(ggml_cgraph * cgraph, const ggml_tensor * inp_pos) { + return get_inp_pos_n_tokens(cgraph, inp_pos) > 1; } #pragma GCC diagnostic pop diff --git a/ggml/src/ggml-openvino/utils.h b/ggml/src/ggml-openvino/utils.h index 513fa83c9d6e..5aa74da38d3b 100644 --- a/ggml/src/ggml-openvino/utils.h +++ b/ggml/src/ggml-openvino/utils.h @@ -164,7 +164,9 @@ std::vector pad_input(const ggml_tensor * tensor, size_t padded_rows, size_t const ggml_tensor * get_inp_pos_tensor(struct ggml_cgraph * cgraph); -bool get_is_prefill(const ggml_tensor * inp_pos); +int64_t get_inp_pos_n_tokens(struct ggml_cgraph * cgraph, const ggml_tensor * inp_pos); + +bool get_is_prefill(struct ggml_cgraph * cgraph, const ggml_tensor * inp_pos); ov::Tensor get_ov_input_tensor(std::shared_ptr ggml_decoder, const std::string & param_name); ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr ggml_decoder, From f5e85d43a048f3d5adefb4c5e29867d8077fba62 Mon Sep 17 00:00:00 2001 From: Strongtut Date: Fri, 28 Aug 2026 05:37:37 -0700 Subject: [PATCH 013/109] metal : add fa-vec tunings for M4 (#27875) This adds fa_vec_tuned_table records for Apple M4 to ggml-metal-tuning.cpp. Includes F16, Q4_0, Q4_1, Q5_0, Q5_1, and Q8_0. (M4, 10 GPU Cores) Co-authored-by: Strongtut <8432058+Strongtut@users.noreply.github.com> --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 226 ++++++++++++++++++++++ 1 file changed, 226 insertions(+) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index c2139fe20b00..285590d1601b 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -542,6 +542,232 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } }, { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 2, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 2, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, From 86632248188c106d749fad34a1dcd237c95863d4 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 28 Aug 2026 16:34:26 +0300 Subject: [PATCH 014/109] context : disable non-fused GDN and LID ops (#27877) --- src/llama-context.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 179c526c2940..9aed80133275 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -231,10 +231,10 @@ llama_context::llama_context( cparams.fused_gdn_ar = true; cparams.fused_gdn_ch = true; - cparams.auto_fgdn = true; + cparams.auto_fgdn = false; - cparams.fused_lid = true; - cparams.auto_flid = true; + cparams.fused_lid = true; + cparams.auto_flid = false; cparams.fused_dsv4_hc_pre = true; cparams.fused_dsv4_hc_comb = true; From 90c26fcd4b2114b4aa39d09d69318cb8f438d27a Mon Sep 17 00:00:00 2001 From: ravel7524 <58877666+ravel7524@users.noreply.github.com> Date: Fri, 28 Aug 2026 10:52:49 -0400 Subject: [PATCH 015/109] Vulkan: add hoisting support for row IDs and expert count in shaders (#26686) * vulkan: add hoisting support for row IDs and expert count in shaders * use hoisted row ids in coopmat2 * vulkan: address review feedback on count_experts - use vk_op_count_experts_push_constants instead of a raw uint vector - apply the fastdiv trick to the ne00 div/mod in count_experts - compute the per-expert offsets with subgroupExclusiveAdd when the device supports it, keeping the serial path as fallback - document the data_d layout and the hoisted_row_id_words bound - drop a leftover debug print in ggml_vk_matmul_id * vulkan: use init_pushconst_fastdiv for count_experts push constants * vulkan: refine comments for row ID hoisting and data layout in count_experts shader * Whitespace --------- Co-authored-by: Jeff Bolz --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 48 +++++++--- .../vulkan-shaders/count_experts.comp | 95 ++++++++++++++++++- .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 35 ++++--- .../vulkan-shaders/mul_mm_cm2.comp | 25 ++++- .../vulkan-shaders/mul_mm_id_funcs.glsl | 15 +++ .../ggml-vulkan/vulkan-shaders/mul_mmq.comp | 35 ++++--- .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + 7 files changed, 212 insertions(+), 42 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 72e844aebfd5..28b2d875d200 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1348,6 +1348,8 @@ struct vk_mat_mat_id_push_constants { uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; uint32_t padded_N; + uint32_t n_experts; + uint32_t hoist_row_ids; }; struct vk_mat_vec_id_push_constants { uint32_t ncols; @@ -1428,6 +1430,10 @@ struct vk_op_count_experts_push_constants { uint32_t nb00; uint32_t nb01; uint32_t a_offset; + uint32_t n_experts; + uint32_t hoist_row_ids; + uint32_t ne00mp; + uint32_t ne00L; }; struct vk_op_glu_push_constants { @@ -1606,6 +1612,10 @@ template <> void init_pushconst_fastdiv(vk_op_glu_push_constants &p) { init_fastdiv_values(p.ne20, p.ne2_0mp, p.ne2_0L); } +template <> void init_pushconst_fastdiv(vk_op_count_experts_push_constants &p) { + init_fastdiv_values(p.ne00, p.ne00mp, p.ne00L); +} + struct vk_op_binary_push_constants { uint32_t ne; uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; @@ -5839,7 +5849,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_count_equal_i32, "count_equal_i32", count_equal_i32_len, count_equal_i32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, { device->subgroup_size }, 1); - ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true); + if (device->subgroup_arithmetic && device->subgroup_require_full_support) { + ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true); + } else { + ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true); + } for (auto &s : device->pipeline_solve_tri_f32) { const vk_solve_tri_pipeline_state &state = s.first; @@ -8970,13 +8984,13 @@ static void ggml_vk_matmul_id( uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11, - uint32_t padded_n) { + uint32_t padded_n, bool hoist_row_ids) { VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " << "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, - nei0, nei1, nbi1, ne11, padded_n }; + nei0, nei1, nbi1, ne11, padded_n, n_as, uint32_t(hoist_row_ids) }; ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as }); } @@ -10162,6 +10176,12 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& // const uint64_t ne23 = dst->ne[3]; const uint64_t n_as = ne02; + // n_as counts, n_as offsets, one total, then one packed row id per (expert, token). + // Hoisting requires 16-bit indices for the packing and a table that fits one binding. + const uint64_t hoisted_row_id_words = 2 * n_as + 1 + nei0 * nei1; + const bool hoist_row_ids = n_as <= 256 && nei0 <= 0xffff && nei1 <= 0xffff && + hoisted_row_id_words * sizeof(uint32_t) <= + ctx->device->properties.limits.maxStorageBufferRange; ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; @@ -10302,7 +10322,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } vk_pipeline count_experts = ctx->device->pipeline_count_experts; - uint32_t expert_count_size = sizeof(uint32_t) * n_as; + const size_t expert_data_size = sizeof(uint32_t) * + (hoist_row_ids ? hoisted_row_id_words : n_as); { if ( @@ -10318,8 +10339,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& ctx->prealloc_size_y = y_sz; ggml_vk_preallocate_buffers(ctx, subctx); } - if (ctx->prealloc_size_split_k < expert_count_size) { - ctx->prealloc_size_split_k = expert_count_size; + if (ctx->prealloc_size_split_k < expert_data_size) { + ctx->prealloc_size_split_k = expert_data_size; ggml_vk_preallocate_buffers(ctx, subctx); } @@ -10385,18 +10406,23 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } } // Count how many times each expert is used - vk_subbuffer expert_count_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + vk_subbuffer expert_count_buf = { ctx->prealloc_split_k, 0, expert_data_size }; if (ctx->prealloc_split_k_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } { - const std::vector pc = { (uint32_t)nei0, + vk_op_count_experts_push_constants pc = { (uint32_t)nei0, (uint32_t)nei1, (uint32_t)(nbi0 / ggml_type_size(ids->type)), (uint32_t)(nbi1 / ggml_type_size(ids->type)), - (uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)) }; + (uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)), + (uint32_t)n_as, + uint32_t(hoist_row_ids), + 0, 0 }; + init_pushconst_fastdiv(pc); ggml_vk_dispatch_pipeline(ctx, subctx, count_experts, - { vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc, { (uint32_t)n_as, 1, 1}); + { vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc, + { hoist_row_ids ? 1u : (uint32_t)n_as, 1, 1}); } if (x_non_contig) { @@ -10465,7 +10491,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& { d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf, ne01, ne21, ne10, ne10, stride_b_y, ne01, stride_batch_x, stride_batch_y, ne20*ne21, - n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n + n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n, hoist_row_ids ); // NOLINT if (x_non_contig || qx_needs_dequant) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp index ffc8608691f7..83c56fce5209 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp @@ -2,6 +2,11 @@ #extension GL_EXT_control_flow_attributes : enable +#ifdef USE_SUBGROUPS +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_KHR_shader_subgroup_arithmetic : enable +#endif + #include "types.glsl" layout (push_constant) uniform parameter @@ -11,6 +16,10 @@ layout (push_constant) uniform parameter uint32_t nb00; uint32_t nb01; uint32_t a_offset; + uint32_t n_experts; + uint32_t hoist_row_ids; + uint32_t ne00mp; + uint32_t ne00L; } p; #define BLOCK_SIZE 256 @@ -21,16 +30,98 @@ layout (binding = 0) readonly buffer A {uint data_a[];}; layout (binding = 1) writeonly buffer D {uint data_d[];}; shared uint vals[BLOCK_SIZE]; +shared uint offsets[BLOCK_SIZE]; +shared uint cursors[BLOCK_SIZE]; + +// see init_fastdiv_values in ggml-vulkan.cpp +uint fastdiv(uint n, uint mp, uint L) { + uint msbs, lsbs; + // msbs = mulhi(n, mp) + umulExtended(n, mp, msbs, lsbs); + return (msbs + n) >> L; +} +// data_d layout when p.hoist_row_ids is set: +// [0, n_experts) per-expert row count +// [n_experts, 2*n_experts) per-expert start offset into the row id region +// [2*n_experts] total row count +// [2*n_experts + 1, ) row ids grouped by expert, packed as (i01 << 16) | (i00 & 0xffff) +// Otherwise only data_d[expert_id] is written, holding that expert's row count. void main() { const uint expert_id = gl_WorkGroupID.x; const uint num_elements = p.ne00 * p.ne01; const uint tid = gl_LocalInvocationID.x; + if (p.hoist_row_ids != 0) { + if (tid < p.n_experts) { + vals[tid] = 0; + } + barrier(); + + for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) { + const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L); + const uint i00 = idx - i01 * p.ne00; + const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]; + if (expert < p.n_experts) { + atomicAdd(vals[expert], 1); + } + } + barrier(); + +#ifdef USE_SUBGROUPS + if (gl_SubgroupID == 0) { + // pad the trip count so the subgroup ops stay in uniform control flow + const uint n_experts_padded = (p.n_experts + gl_SubgroupSize - 1) & ~(gl_SubgroupSize - 1); + uint base = 0; + for (uint expert = gl_SubgroupInvocationID; expert < n_experts_padded; expert += gl_SubgroupSize) { + const bool in_range = expert < p.n_experts; + const uint count = in_range ? vals[expert] : 0; + const uint offset = base + subgroupExclusiveAdd(count); + if (in_range) { + data_d[expert] = count; + data_d[p.n_experts + expert] = offset; + offsets[expert] = offset; + cursors[expert] = 0; + } + base += subgroupAdd(count); + } + if (subgroupElect()) { + data_d[2 * p.n_experts] = base; + } + } +#else + if (tid == 0) { + uint offset = 0; + for (uint expert = 0; expert < p.n_experts; ++expert) { + const uint count = vals[expert]; + data_d[expert] = count; + data_d[p.n_experts + expert] = offset; + offsets[expert] = offset; + cursors[expert] = 0; + offset += count; + } + data_d[2 * p.n_experts] = offset; + } +#endif + barrier(); + + for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) { + const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L); + const uint i00 = idx - i01 * p.ne00; + const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]; + if (expert < p.n_experts) { + const uint row = atomicAdd(cursors[expert], 1); + const uint packed_row_id = (i01 << 16) | (i00 & 0xffffu); + data_d[2 * p.n_experts + 1 + offsets[expert] + row] = packed_row_id; + } + } + return; + } + uint count = 0; for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) { - const uint i01 = idx / p.ne00; - const uint i00 = idx % p.ne00; + const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L); + const uint i00 = idx - i01 * p.ne00; const uint a = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]; count += uint(a == expert_id); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 3df88044a5ee..c1ccac7aa202 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -88,6 +88,9 @@ layout (push_constant) uniform parameter uint nei1; uint nbi1; uint ne11; + uint padded_N; + uint n_experts; + uint hoist_row_ids; #else uint base_work_group_z; uint num_batches; @@ -214,27 +217,31 @@ void main() { const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B_EFF * LOAD_VEC_BATCH_B / BK; #ifdef MUL_MAT_ID -#ifdef MUL_MAT_ID_USE_SUBGROUPS - if (bitCount(p.nei0) == 1) { - load_row_ids(expert_idx, true, ic); + if (p.hoist_row_ids != 0) { + load_row_ids_hoisted(expert_idx, ic); } else { - load_row_ids(expert_idx, false, ic); - } +#ifdef MUL_MAT_ID_USE_SUBGROUPS + if (bitCount(p.nei0) == 1) { + load_row_ids(expert_idx, true, ic); + } else { + load_row_ids(expert_idx, false, ic); + } #else - _ne1 = 0; - for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) { - for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) { - if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { - if (_ne1 >= ic * BN) { - row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1); + _ne1 = 0; + for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) { + for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) { + if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { + if (_ne1 >= ic * BN) { + row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1); + } + _ne1++; } - _ne1++; } } - } - barrier(); + barrier(); #endif + } // Workgroup has no work if (ic * BN >= _ne1) return; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp index a2e15f6f5ced..cf78474a9258 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -67,6 +67,10 @@ layout (push_constant) uniform parameter #endif // N dimension for the B matrix can be >= p.N uint padded_N; +#ifdef MUL_MAT_ID + uint n_experts; + uint hoist_row_ids; +#endif } p; @@ -225,6 +229,23 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) { } barrier(); } + +void load_row_ids_hoisted(uint expert_idx, uint ic) { + _ne1 = uint(data_expert_count[expert_idx]); + + const uint tile_begin = ic * BN; + const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0; + const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]); + const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin; + + for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) { + const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]); + const uint ii0 = packed_row_id & 0xffffu; + const uint ii1 = packed_row_id >> 16; + row_ids[i] = u16vec4(fastmod(ii0, p.ne11), ii1, ii0, 0); + } + barrier(); +} #endif void main() { @@ -266,7 +287,9 @@ void main() { const uint ik = gl_WorkGroupID.x / blocks_m; #ifdef MUL_MAT_ID - if (bitCount(p.nei0) == 1) { + if (p.hoist_row_ids != 0) { + load_row_ids_hoisted(expert_idx, ic); + } else if (bitCount(p.nei0) == 1) { load_row_ids(expert_idx, true, ic); } else { load_row_ids(expert_idx, false, ic); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl index 26c5c12a49a2..54ad60b2efba 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl @@ -71,4 +71,19 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) { barrier(); } #endif // MUL_MAT_ID_USE_SUBGROUPS + +void load_row_ids_hoisted(uint expert_idx, uint ic) { + _ne1 = uint(data_expert_count[expert_idx]); + + const uint tile_begin = ic * BN; + const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0; + const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]); + const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin; + + for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) { + const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]); + row_ids[i] = u16vec2(packed_row_id & 0xffffu, packed_row_id >> 16); + } + barrier(); +} #endif // MUL_MAT_ID diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp index aae1c2e8ae9f..c2d84c05b40d 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp @@ -56,6 +56,9 @@ layout (push_constant) uniform parameter uint nei1; uint nbi1; uint ne11; + uint padded_N; + uint n_experts; + uint hoist_row_ids; #else uint base_work_group_z; uint num_batches; @@ -157,27 +160,31 @@ void main() { const uint loadstride_b = BLOCK_SIZE * LOAD_VEC_B / BK; #ifdef MUL_MAT_ID -#ifdef MUL_MAT_ID_USE_SUBGROUPS - if (bitCount(p.nei0) == 1) { - load_row_ids(expert_idx, true, ic); + if (p.hoist_row_ids != 0) { + load_row_ids_hoisted(expert_idx, ic); } else { - load_row_ids(expert_idx, false, ic); - } +#ifdef MUL_MAT_ID_USE_SUBGROUPS + if (bitCount(p.nei0) == 1) { + load_row_ids(expert_idx, true, ic); + } else { + load_row_ids(expert_idx, false, ic); + } #else - _ne1 = 0; - for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) { - for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) { - if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { - if (_ne1 >= ic * BN) { - row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1); + _ne1 = 0; + for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) { + for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) { + if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { + if (_ne1 >= ic * BN) { + row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1); + } + _ne1++; } - _ne1++; } } - } - barrier(); + barrier(); #endif + } // Workgroup has no work if (ic * BN >= _ne1) return; 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 0da943da9563..d375c2d12771 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1039,6 +1039,7 @@ void process_shaders() { string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("count_experts", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}})); + string_to_spv("count_experts_subgroup", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}, {"USE_SUBGROUPS", "1"}})); for (std::string dim_str : {"", "_3d"}) { for (bool bda : {false, true}) { From a43c3986b44225c8633fe938464278d20dff0ec0 Mon Sep 17 00:00:00 2001 From: Tekin Ertekin Date: Fri, 28 Aug 2026 20:09:08 +0300 Subject: [PATCH 016/109] ggml : fix conv_transpose_2d for multiple batches (#26132) * ggml : fix conv_transpose_2d for multiple batches ggml_compute_forward_conv_transpose_2d_impl only computed the first batch (ne[3] of the destination); every batch after the first was left as zero. Both the src1 permutation and the main compute loop now iterate over the batch dimension, and the work buffer size in ggml_graph_plan is scaled by the src1 batch count so the extra permuted batches fit. A multi-batch test case is added to test-backend-ops. Fixes ggml-org/ggml#1448 * metal : fix conv_transpose_2d for multiple batches The kernel only computed batch 0 of the input (src1->ne[3]); every output batch after the first was left as zero, so multi-batch conv_transpose_2d results diverged from the CPU reference. The grid now covers all batches (OW x OH x OC x N), the kernel decodes the batch from the grid z coordinate and offsets both the input and destination indices accordingly. nb3 is passed in the kernel args. Assisted-by: pi:llama.cpp/Qwen3.8-27B --------- Co-authored-by: Georgi Gerganov --- ggml/src/ggml-cpu/ggml-cpu.c | 3 +- ggml/src/ggml-cpu/ops.cpp | 58 ++++++++++++++------------ ggml/src/ggml-metal/ggml-metal-impl.h | 1 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 4 +- ggml/src/ggml-metal/kernels/conv.metal | 7 ++-- tests/test-backend-ops.cpp | 1 + 6 files changed, 43 insertions(+), 31 deletions(-) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 87ac0a702efc..b9c0fa3ddc01 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2936,12 +2936,13 @@ struct ggml_cplan ggml_graph_plan( const int64_t ne10 = node->src[1]->ne[0]; // W const int64_t ne11 = node->src[1]->ne[1]; // H const int64_t ne12 = node->src[1]->ne[2]; // Channels In + const int64_t ne13 = node->src[1]->ne[3]; // Batch GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32); GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32); cur += ggml_type_size(node->src[0]->type) * ne00 * ne01 * ne02 * ne03; - cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12; + cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12 * ne13; } break; case GGML_OP_TOP_K: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index b869f4bddde0..b47ce5463c61 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -7267,18 +7267,21 @@ static void ggml_compute_forward_conv_transpose_2d_impl( } } - // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh) + // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches { kernel_t * const wdata = (kernel_t *) params->wdata + nk; - for (int i12 = 0; i12 < ne12; i12++) { - for (int i11 = 0; i11 < ne11; i11++) { - const float * const src = (float *)((char *) src1->data + i12*nb12 + i11*nb11); - kernel_t * dst_data = wdata + i11*ne10*ne12; - for (int i10 = 0; i10 < ne10; i10++) { - if constexpr (std::is_same_v) { - dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]); - } else { - dst_data[i10*ne12 + i12] = src[i10]; + for (int i13 = 0; i13 < ne13; i13++) { + kernel_t * const wdata_b = wdata + i13*ne10*ne11*ne12; + for (int i12 = 0; i12 < ne12; i12++) { + for (int i11 = 0; i11 < ne11; i11++) { + const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11); + kernel_t * dst_data = wdata_b + i11*ne10*ne12; + for (int i10 = 0; i10 < ne10; i10++) { + if constexpr (std::is_same_v) { + dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]); + } else { + dst_data[i10*ne12 + i12] = src[i10]; + } } } } @@ -7305,24 +7308,27 @@ static void ggml_compute_forward_conv_transpose_2d_impl( kernel_t * const wdata_src = wdata + nk; for (int i2 = ip0; i2 < ip1; i2++) { // Cout - float * dst_data = (float *)((char *) dst->data + i2*nb2); kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03; - for (int i11 = 0; i11 < ne11; i11++) { - for (int i10 = 0; i10 < ne10; i10++) { - const int i1n = i11*ne10*ne12 + i10*ne12; - for (int i01 = 0; i01 < ne01; i01++) { - for (int i00 = 0; i00 < ne00; i00++) { - float v = 0; - if constexpr (std::is_same_v) { - ggml_vec_dot_f16(ne03, &v, 0, - wdata_src + i1n, 0, - wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); - } else { - ggml_vec_dot_f32(ne03, &v, 0, - wdata_src + i1n, 0, - wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); + for (int i3 = 0; i3 < ne3; i3++) { // batch + float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2); + kernel_t * wdata_src_b = wdata_src + i3*ne10*ne11*ne12; + for (int i11 = 0; i11 < ne11; i11++) { + for (int i10 = 0; i10 < ne10; i10++) { + const int i1n = i11*ne10*ne12 + i10*ne12; + for (int i01 = 0; i01 < ne01; i01++) { + for (int i00 = 0; i00 < ne00; i00++) { + float v = 0; + if constexpr (std::is_same_v) { + ggml_vec_dot_f16(ne03, &v, 0, + wdata_src_b + i1n, 0, + wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); + } else { + ggml_vec_dot_f32(ne03, &v, 0, + wdata_src_b + i1n, 0, + wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); + } + dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v; } - dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v; } } } diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 9becf04797ba..49102afe9c02 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -660,6 +660,7 @@ typedef struct { uint64_t nb0; uint64_t nb1; uint64_t nb2; + uint64_t nb3; } ggml_metal_kargs_conv_transpose_2d; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 75de0f6dd08a..f6f2fdc86c6a 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -4645,6 +4645,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { const int32_t OW = op->ne[0]; const int32_t OH = op->ne[1]; const int32_t OC = op->ne[2]; + const int32_t N = op->src[1]->ne[3]; ggml_metal_kargs_conv_transpose_2d args = { /*.IC =*/ IC, @@ -4657,6 +4658,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { /*.nb0 =*/ nb0, /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, + /*.nb3 =*/ nb3, }; auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_2d(lib, op); @@ -4671,7 +4673,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { const size_t smem = GGML_PAD(KW * KH * sizeof(float), 16); ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); - ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC, KW, KH, 1); + ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC * N, KW, KH, 1); return 1; } diff --git a/ggml/src/ggml-metal/kernels/conv.metal b/ggml/src/ggml-metal/kernels/conv.metal index 5685b5cd4915..a5d5aa9d9293 100644 --- a/ggml/src/ggml-metal/kernels/conv.metal +++ b/ggml/src/ggml-metal/kernels/conv.metal @@ -366,7 +366,8 @@ kernel void kernel_conv_transpose_2d( const int64_t out_x = tgpig[0]; const int64_t out_y = tgpig[1]; - const int64_t out_c = tgpig[2]; + const int64_t batch = tgpig[2] / args.OC; + const int64_t out_c = tgpig[2] % args.OC; const int64_t kw = tpitg[0]; const int64_t kh = tpitg[1]; @@ -390,7 +391,7 @@ kernel void kernel_conv_transpose_2d( if (in_x >= args.IW) continue; - const int64_t input_idx = (args.IW * args.IH) * in_c + (args.IW) * in_y + in_x; + const int64_t input_idx = (args.IW * args.IH) * (args.IC * batch + in_c) + (args.IW) * in_y + in_x; const int64_t kernel_idx = (args.KH * args.KW * args.OC) * in_c + (args.KH * args.KW) * out_c + (args.KW) * kh + kw; v += (float)src0[kernel_idx] * src1[input_idx]; @@ -408,7 +409,7 @@ kernel void kernel_conv_transpose_2d( total += shared_sum[i]; } - device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2); + device float * dst_ptr = (device float *) (dst + batch*args.nb3 + out_c*args.nb2 + out_y * args.nb1 + out_x*args.nb0); dst_ptr[0] = total; } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index df0b9aa66689..6be83ac161bd 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8785,6 +8785,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_conv_transpose_2d({3, 2, 3, 1}, {2, 2, 1, 3}, 1, kernel_type)); test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 1}, {3, 3, 1, 9}, 2, kernel_type)); test_cases.emplace_back(new test_conv_transpose_2d({129, 63, 35, 1}, {3, 3, 48, 35}, 1, kernel_type)); + test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 2}, {3, 3, 1, 9}, 2, kernel_type)); // for multiple batches } test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4, 500, 1, 1})); From b387ddfd84b4b1f79a6e09910748195e3320e89e Mon Sep 17 00:00:00 2001 From: Eric A Stalee <87948564+Eric-A-Stalee@users.noreply.github.com> Date: Fri, 28 Aug 2026 12:12:33 -0500 Subject: [PATCH 017/109] vulkan: fix missing view-alias dependencies in ggml_vk_graph_optimize (#27812) * vulkan: fix missing view-alias dependencies in ggml_vk_graph_optimize is_src_of doesn't treat two views of one tensor as dependent, so the optimizer reorders nodes across aliased reads and writes. Result: silently wrong tokens under greedy decoding, different output on every server start, and invalid speculative-decoding acceptance, with nothing logged. Hits Qwen3.8's recurrent state (and any model with view-aliased state) on AMD and NVIDIA Vulkan. CUDA is clean. Compare view_src bases on both sides. Fixes #27805 * vulkan: don't treat view/no-op nodes as aliasing dependencies Nodes whose op is NONE, RESHAPE, TRANSPOSE, VIEW or PERMUTE execute nothing, so aliasing through them is not a real dependency. The previous base comparison matched them anyway, which only costs the optimizer reordering freedom. Co-authored-by: Jeff Bolz * vulkan: make the lambda parameter const and capture is_empty in is_src_of Code will not compile without these changes. is_src_of has an empty capture list, so is_empty was not visible inside it, and is_empty took a non-const pointer, while is_src_of receives const ones. Other call sites pass non-const pointers, which still convert as usual. --------- Co-authored-by: Jeff Bolz --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 28b2d875d200..320127cdc5ac 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17800,20 +17800,32 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * return; } - auto const &is_empty = [](ggml_tensor * node) -> bool { + auto const &is_empty = [](const ggml_tensor * node) -> bool { return 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; }; - auto const &is_src_of = [](const ggml_tensor *dst, const ggml_tensor *src) -> bool { + auto const &is_src_of = [&is_empty](const ggml_tensor *dst, const ggml_tensor *src) -> bool { + auto const &base = [](const ggml_tensor * tensor) { + return tensor->view_src ? tensor->view_src : tensor; + }; for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) { if (dst->src[s] == src) { return true; } + if (is_empty(dst) || is_empty(src)) { + continue; + } + // A source view of dst may read storage written through a different view by src. + if (dst->src[s] && base(dst->src[s]) == base(src)) { + return true; + } + // Moving dst forward may overwrite storage still read through a view by src. + if (src->src[s] && base(dst) == base(src->src[s])) { + return true; + } } // implicit dependency if they view the same tensor - const ggml_tensor *dst2 = dst->view_src ? dst->view_src : dst; - const ggml_tensor *src2 = src->view_src ? src->view_src : src; - if (dst2 == src2) { + if (base(dst) == base(src)) { return true; } return false; From 6fe74980162af0ed5e559870d5deccafaa034e7c Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Fri, 28 Aug 2026 19:24:46 +0200 Subject: [PATCH 018/109] model: qwen4exp: reduce number of graph splits (#27880) --- src/models/models.h | 5 ++++- src/models/qwen4exp.cpp | 32 +++++++++++++++++++++++--------- 2 files changed, 27 insertions(+), 10 deletions(-) diff --git a/src/models/models.h b/src/models/models.h index af60764c2f7f..9b87a40d5af9 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2360,9 +2360,12 @@ struct llama_model_qwen4exp : public llama_model_base { int64_t channels, int il); + ggml_tensor * build_inp_ple( + const llama_memory_hybrid_idx_context * mctx_hyb); + ggml_tensor * build_ple( llm_graph_input_rs * inp, - const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * emb, ggml_tensor * hidden, int il); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index acfdd5b50038..abf6a0502fbf 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -296,6 +296,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inpL = build_inp_embd(model.tok_embd); cb(inpL, "model.input_embed", -1); + ggml_build_forward_expand(gf, inpL); auto * inp = build_inp_mem_hybrid(); @@ -312,6 +313,13 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = build_inp_out_ids(); + ggml_tensor * ple_emb = nullptr; + if (hparams.ple_n_heads > 0) { + ple_emb = build_inp_ple(mctx_hyb); + // make sure ple_emb and build_inp_embd are in the same graph split + ggml_build_forward_expand(gf, ple_emb); + } + // the wide residual starts as hc identical copies of the embedding ggml_tensor * res_hc = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), @@ -322,7 +330,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa res->t_layer_inp[il] = res_hc; if (hparams.is_ple(il)) { - res_hc = build_ple(inp->get_recr(), mctx_hyb, res_hc, il); + res_hc = build_ple(inp->get_recr(), ple_emb, res_hc, il); } ggml_tensor * inject = nullptr; @@ -1090,13 +1098,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at( return conv_input; } -ggml_tensor * llama_model_qwen4exp::graph::build_ple( - llm_graph_input_rs * inp, - const llama_memory_hybrid_idx_context * mctx_hyb, - ggml_tensor * hidden, - int il) { - const int64_t hc = hparams.dsv4_hc_mult; - const int64_t hc_dim = hc * n_embd; +ggml_tensor * llama_model_qwen4exp::graph::build_inp_ple( + const llama_memory_hybrid_idx_context * mctx_hyb) { const int64_t n_heads = hparams.ple_n_heads; // the attention cells see every ubatch regardless of the layer types @@ -1111,7 +1114,18 @@ ggml_tensor * llama_model_qwen4exp::graph::build_ple( // gather then flatten the heads: get_rows lays the head dimension out slowest, as the reference does ggml_tensor * emb = ggml_get_rows(ctx0, model.per_layer_tok_embd, rows); emb = ggml_reshape_2d(ctx0, emb, hparams.ple_head_dim * n_heads, n_tokens); - cb(emb, "ple_embd", il); + cb(emb, "ple_embd", -1); + + return emb; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_ple( + llm_graph_input_rs * inp, + ggml_tensor * emb, + ggml_tensor * hidden, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; ggml_tensor * key = build_lora_mm(model.layers[il].ple_key, emb); ggml_tensor * value = build_lora_mm(model.layers[il].ple_value, emb); From 50f068ffffc3e0e4c9c2e4139281c6075224f429 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Fri, 28 Aug 2026 20:51:05 +0200 Subject: [PATCH 019/109] bench: add --tensor-read-lazy (#27881) * bench: add --tensor-read-lazy * rm the alias * rename to LLAMA_LAZY_MODE_* --- common/arg.cpp | 6 +-- common/common.cpp | 2 +- common/common.h | 2 +- include/llama.h | 10 ++--- src/llama-model-loader.cpp | 4 +- src/llama-model-loader.h | 2 +- src/llama-model.cpp | 2 +- src/llama.cpp | 2 +- tools/llama-bench/README.md | 1 + tools/llama-bench/llama-bench.cpp | 65 +++++++++++++++++++++++++++++-- 10 files changed, 78 insertions(+), 18 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index e346863e51fd..4469612cd5b2 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2735,9 +2735,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "- auto: on, but only for tensors larger than 4 GiB\n" "- off: always keep them resident", [](common_params & params, const std::string & value) { - /**/ if (value == "on") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_ON; } - else if (value == "auto") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; } - else if (value == "off") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF; } + /**/ if (value == "on") { params.lazy_mode = LLAMA_LAZY_MODE_ON; } + else if (value == "auto") { params.lazy_mode = LLAMA_LAZY_MODE_AUTO; } + else if (value == "off") { params.lazy_mode = LLAMA_LAZY_MODE_OFF; } else { throw std::invalid_argument("invalid value"); } } ).set_env("LLAMA_ARG_TENSOR_READ_LAZY")); diff --git a/common/common.cpp b/common/common.cpp index 347e8e9fc416..d162a38800e0 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1688,7 +1688,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.main_gpu = params.main_gpu; mparams.split_mode = params.split_mode; mparams.load_mode = params.load_mode; - mparams.tensor_read_lazy = params.tensor_read_lazy; + mparams.lazy_mode = params.lazy_mode; mparams.tensor_split = params.tensor_split; mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; diff --git a/common/common.h b/common/common.h index a333f702ac1d..4e9448bb106a 100644 --- a/common/common.h +++ b/common/common.h @@ -483,7 +483,7 @@ struct common_params { enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs enum llama_load_mode load_mode = LLAMA_LOAD_MODE_AUTO; // how to load the model - enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; // on-demand reading of tensors marked by the arch + enum llama_lazy_mode lazy_mode = LLAMA_LAZY_MODE_AUTO; // on-demand reading of tensors marked by the arch common_cpu_params cpuparams; common_cpu_params cpuparams_batch; diff --git a/include/llama.h b/include/llama.h index 49a758db2680..ef7a012c43a1 100644 --- a/include/llama.h +++ b/include/llama.h @@ -214,10 +214,10 @@ extern "C" { LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str); - enum llama_tensor_read_lazy { - LLAMA_TENSOR_READ_LAZY_OFF = 0, // always read the whole tensor up front - LLAMA_TENSOR_READ_LAZY_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap) - LLAMA_TENSOR_READ_LAZY_ON = 2, // read the rows of tensors marked by the arch on demand (requires mmap) + enum llama_lazy_mode { + LLAMA_LAZY_MODE_OFF = 0, // always read the whole tensor up front + LLAMA_LAZY_MODE_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap) + LLAMA_LAZY_MODE_ON = 2, // read the rows of tensors marked by the arch on demand (requires mmap) }; enum llama_context_type { @@ -321,7 +321,7 @@ extern "C" { enum llama_split_mode split_mode; // how to split the model across multiple GPUs enum llama_load_mode load_mode; // how to load the model - enum llama_tensor_read_lazy tensor_read_lazy; // on-demand reading of tensors marked by the arch + enum llama_lazy_mode lazy_mode; // on-demand reading of tensors marked by the arch // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE int32_t main_gpu; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index d9241022cf41..1b1f852a010d 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1287,10 +1287,10 @@ struct ggml_tensor * llama_model_loader::create_tensor( return NULL; } - if ((flags & TENSOR_READ_LAZY) && use_mmap && tensor_read_lazy != LLAMA_TENSOR_READ_LAZY_OFF) { + if ((flags & TENSOR_READ_LAZY) && use_mmap && lazy_mode != LLAMA_LAZY_MODE_OFF) { // in auto mode, small tensors are cheap enough to keep resident constexpr size_t auto_lazy_min_size = 4ull * 1024 * 1024 * 1024; - if (tensor_read_lazy == LLAMA_TENSOR_READ_LAZY_ON || ggml_nbytes(cur) > auto_lazy_min_size) { + if (lazy_mode == LLAMA_LAZY_MODE_ON || ggml_nbytes(cur) > auto_lazy_min_size) { const auto & w = require_weight(tn.str().c_str()); lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur)); diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 407260e9907e..20f744253892 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -84,7 +84,7 @@ struct llama_model_loader { bool load_mtp; // set by the caller before the create_tensor() calls - enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF; + enum llama_lazy_mode lazy_mode = LLAMA_LAZY_MODE_OFF; llama_files files; llama_ftype ftype; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index fc83658dd7ff..65a6702cefe7 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2681,7 +2681,7 @@ llama_model_params llama_model_default_params() { /*.n_gpu_layers =*/ -1, /*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER, /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, - /*.tensor_read_lazy =*/ LLAMA_TENSOR_READ_LAZY_AUTO, + /*.lazy_mode =*/ LLAMA_LAZY_MODE_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, diff --git a/src/llama.cpp b/src/llama.cpp index 9c841ee35213..6ec5d315dcec 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -318,7 +318,7 @@ static std::pair llama_model_load(struct gguf_context * meta llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode, params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides); - ml.tensor_read_lazy = params.tensor_read_lazy; + ml.lazy_mode = params.lazy_mode; ml.print_info(); std::unique_ptr model_ptr(llama_model_create(ml, params)); diff --git a/tools/llama-bench/README.md b/tools/llama-bench/README.md index 42cb14859f07..a1404d2e33e3 100644 --- a/tools/llama-bench/README.md +++ b/tools/llama-bench/README.md @@ -67,6 +67,7 @@ test parameters: -nkvo, --no-kv-offload <0|1> (default: 0) -fa, --flash-attn (default: auto) -dev, --device (default: auto) + --tensor-read-lazy (default: auto) -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode) -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode) -embd, --embeddings <0|1> (default: 0) diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index bc14d15c7968..1b4bbde4d4ff 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -271,6 +271,19 @@ static const char * split_mode_str(llama_split_mode mode) { } } +static const char * lazy_mode_str(llama_lazy_mode mode) { + switch (mode) { + case LLAMA_LAZY_MODE_OFF: + return "off"; + case LLAMA_LAZY_MODE_AUTO: + return "auto"; + case LLAMA_LAZY_MODE_ON: + return "on"; + default: + GGML_ABORT("invalid tensor read lazy mode"); + } +} + static std::string pair_str(const std::pair & p) { static char buf[32]; snprintf(buf, sizeof(buf), "%d,%d", p.first, p.second); @@ -341,6 +354,7 @@ struct cmd_params { std::vector n_cpu_moe; std::vector split_mode; std::vector load_mode; + std::vector lazy_mode; std::vector main_gpu; std::vector no_kv_offload; std::vector flash_attn; @@ -385,6 +399,7 @@ static const cmd_params cmd_params_defaults = { /* n_cpu_moe */ { 0 }, /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, + /* lazy_mode */ { LLAMA_LAZY_MODE_AUTO }, /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, @@ -460,6 +475,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); + printf(" --tensor-read-lazy (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str()); printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); @@ -786,6 +802,32 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { break; } params.load_mode.insert(params.load_mode.end(), modes.begin(), modes.end()); + } else if (arg == "--tensor-read-lazy") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + + std::vector modes; + for (const auto & m : p) { + llama_lazy_mode mode; + if (m == "on") { + mode = LLAMA_LAZY_MODE_ON; + } else if (m == "auto") { + mode = LLAMA_LAZY_MODE_AUTO; + } else if (m == "off") { + mode = LLAMA_LAZY_MODE_OFF; + } else { + invalid_param = true; + break; + } + modes.push_back(mode); + } + if (invalid_param) { + break; + } + params.lazy_mode.insert(params.lazy_mode.end(), modes.begin(), modes.end()); } else if (arg == "-mg" || arg == "--main-gpu") { if (++i >= argc) { invalid_param = true; @@ -1137,6 +1179,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.load_mode.empty()) { params.load_mode = cmd_params_defaults.load_mode; } + if (params.lazy_mode.empty()) { + params.lazy_mode = cmd_params_defaults.lazy_mode; + } if (params.main_gpu.empty()) { params.main_gpu = cmd_params_defaults.main_gpu; } @@ -1203,6 +1248,7 @@ struct cmd_params_instance { int n_cpu_moe; llama_split_mode split_mode; llama_load_mode load_mode; + llama_lazy_mode lazy_mode; int main_gpu; bool no_kv_offload; llama_flash_attn_type flash_attn; @@ -1224,6 +1270,7 @@ struct cmd_params_instance { } mparams.split_mode = split_mode; mparams.load_mode = load_mode; + mparams.lazy_mode = lazy_mode; mparams.main_gpu = main_gpu; mparams.tensor_split = tensor_split.data(); mparams.no_host = no_host; @@ -1271,7 +1318,8 @@ struct cmd_params_instance { return model == other.model && n_gpu_layers == other.n_gpu_layers && n_cpu_moe == other.n_cpu_moe && split_mode == other.split_mode && main_gpu == other.main_gpu && tensor_split == other.tensor_split && - load_mode == other.load_mode && devices == other.devices && no_host == other.no_host && + load_mode == other.load_mode && lazy_mode == other.lazy_mode && + devices == other.devices && no_host == other.no_host && vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides); } @@ -1305,6 +1353,7 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & ncmoe : params.n_cpu_moe) for (const auto & sm : params.split_mode) for (const auto & lm : params.load_mode) + for (const auto & lzm : params.lazy_mode) for (const auto & mg : params.main_gpu) for (const auto & devs : params.devices) for (const auto & ts : params.tensor_split) @@ -1344,6 +1393,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1380,6 +1430,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1416,6 +1467,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1457,6 +1509,7 @@ struct test { int n_cpu_moe; llama_split_mode split_mode; llama_load_mode load_mode; + llama_lazy_mode lazy_mode; int main_gpu; bool no_kv_offload; llama_flash_attn_type flash_attn; @@ -1496,6 +1549,7 @@ struct test { n_cpu_moe = inst.n_cpu_moe; split_mode = inst.split_mode; load_mode = inst.load_mode; + lazy_mode = inst.lazy_mode; main_gpu = inst.main_gpu; no_kv_offload = inst.no_kv_offload; flash_attn = inst.flash_attn; @@ -1563,7 +1617,8 @@ struct test { "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", "main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split", - "tensor_buft_overrides", "load_mode", "embeddings", + "tensor_buft_overrides", "load_mode", "lazy_mode", + "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", "test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts" @@ -1588,7 +1643,7 @@ struct test { if (field == "avg_ts" || field == "stddev_ts") { return FLOAT; } - if (field == "load_mode") { + if (field == "load_mode" || field == "lazy_mode") { return STRING; } return STRING; @@ -1658,6 +1713,7 @@ struct test { tensor_split_str, tensor_buft_overrides_str, llama_load_mode_name(load_mode), + lazy_mode_str(lazy_mode), std::to_string(embeddings), std::to_string(no_op_offload), std::to_string(no_host), @@ -1972,6 +2028,9 @@ struct markdown_printer : public printer { if (params.load_mode.size() > 1 || params.load_mode != cmd_params_defaults.load_mode) { fields.emplace_back("load_mode"); } + if (params.lazy_mode.size() > 1 || params.lazy_mode != cmd_params_defaults.lazy_mode) { + fields.emplace_back("lazy_mode"); + } if (params.embeddings.size() > 1 || params.embeddings != cmd_params_defaults.embeddings) { fields.emplace_back("embeddings"); } From d7bd3bfcad3e29c7e49fd26f38c79ee3e9a3fd6b Mon Sep 17 00:00:00 2001 From: kurquhar Date: Fri, 28 Aug 2026 14:01:59 -0700 Subject: [PATCH 020/109] snapdragon: python SDK setup (Windows) (#27903) * port setup-build.ps1 to setup_sdk.py, to facilitate installation of Hexagon and OpenCL SDKs on Windows * rename setup_sdk.py -> setup-sdk.py * flake8 fix: print() -> logger.info() --------- Co-authored-by: Kristopher Urquhart --- docs/backend/snapdragon/windows.md | 13 +- scripts/snapdragon/build.py | 18 ++- scripts/snapdragon/sdk.py | 62 ++++++++ scripts/snapdragon/setup-sdk.py | 233 +++++++++++++++++++++++++++++ 4 files changed, 319 insertions(+), 7 deletions(-) create mode 100644 scripts/snapdragon/sdk.py create mode 100644 scripts/snapdragon/setup-sdk.py diff --git a/docs/backend/snapdragon/windows.md b/docs/backend/snapdragon/windows.md index 3f7d60dd95f8..886cfda3f653 100644 --- a/docs/backend/snapdragon/windows.md +++ b/docs/backend/snapdragon/windows.md @@ -24,7 +24,18 @@ must be included in the .cat file digitally signed with a trusted certificate. This document covers details on how to generate personal certificate files (.pfx) and how to configure the system to allow for test signatures (aka test-signing). -## Install the latest Adreno OpenCL SDK +## Install Windows SDKs + +The recommended method is `setup-sdk.py`: + +``` +> python scripts\snapdragon\setup-sdk.py --list-sdk-releases +> python scripts\snapdragon\setup-sdk.py --hexagon --opencl +``` + +It installs the selected SDKs under `C:\Qualcomm` and sets their corresponding environment variables for the current user. Start a new terminal after it completes; native Windows builds check all SDK paths before CMake runs. + +Select the SDKs to install with `--hexagon` and `--opencl`; use both to prepare a dual-backend build. To select a different available version, pass it to the SDK option, for example `--hexagon 6.4.0.2`. SDK versions install side by side, so you can switch versions without deleting an existing installation. Use `--force` to reinstall the selected SDKs. Use a new CMake build directory after each switch because CMake caches the SDK paths. Either use the trimmed down version (optimized for CI) from diff --git a/scripts/snapdragon/build.py b/scripts/snapdragon/build.py index 02dcb930cb67..5e9fab3d08b6 100755 --- a/scripts/snapdragon/build.py +++ b/scripts/snapdragon/build.py @@ -11,6 +11,8 @@ import shutil import logging +from sdk import validate_windows_sdks + logger = logging.getLogger("build") @@ -65,6 +67,13 @@ def main(): logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.") sys.exit(1) + if target_type == "windows": + logger.info("Windows target selected. Forcing native compilation...") + args.no_docker = True + if platform.system() != "Windows": + logger.warning("Warning: Windows compilation is intended to run on Windows arm64 hosts.") + validate_windows_sdks() + # Determine preset and check if it's debug preset = args.preset if preset: @@ -120,12 +129,6 @@ def main(): jobs = args.jobs if args.jobs else os.cpu_count() or 4 - if target_type == "windows": - logger.info("Windows target selected. Forcing native compilation...") - args.no_docker = True - if platform.system() != "Windows": - logger.warning("Warning: Windows compilation is intended to run on Windows arm64 hosts.") - if args.no_docker: # Native/local host build logger.info("Running native/local CMake build...") @@ -258,3 +261,6 @@ def main(): except KeyboardInterrupt: logger.info("\nInterrupted by user.") sys.exit(130) + except RuntimeError as err: + logger.error("Error: %s", err) + sys.exit(1) diff --git a/scripts/snapdragon/sdk.py b/scripts/snapdragon/sdk.py new file mode 100644 index 000000000000..bb3cb77b2ad4 --- /dev/null +++ b/scripts/snapdragon/sdk.py @@ -0,0 +1,62 @@ +import os +from pathlib import Path + + +SDK_CONFIGS = ( + { + "name": "Hexagon SDK", + "repo": "snapdragon-toolchain/hexagon-sdk", + "default_version": "6.6.0.0", + "parent_dir": "Hexagon_SDK", + "archive_prefix": "hexagon-sdk-v", + "markers": ("hexagon_sdk.json",), + }, + { + "name": "OpenCL SDK", + "repo": "snapdragon-toolchain/opencl-sdk", + "default_version": "2.3.2", + "parent_dir": "OpenCL_SDK", + "archive_prefix": "adreno-opencl-sdk-v", + "markers": ("include/CL", "lib/OpenCL.lib"), + }, +) + + +def is_valid_sdk(config, target_dir): + return target_dir.is_dir() and all((target_dir / marker).exists() for marker in config["markers"]) + + +def get_hexagon_tools_dir(hexagon_dir): + tools_parent = hexagon_dir / "tools" / "HEXAGON_Tools" + if not tools_parent.is_dir(): + raise RuntimeError(f"Expected Hexagon tools directory in {tools_parent}") + tools_dirs = [path for path in tools_parent.iterdir() if path.is_dir()] + if len(tools_dirs) != 1: + raise RuntimeError(f"Expected one Hexagon tools directory in {tools_parent}") + return tools_dirs[0] + + +def validate_windows_sdks(): + hexagon_config, opencl_config = SDK_CONFIGS + hexagon_dir = os.environ.get("HEXAGON_SDK_ROOT") + tools_dir = os.environ.get("HEXAGON_TOOLS_ROOT") + opencl_dir = os.environ.get("OPENCL_SDK_ROOT") + missing = [] + + expected_tools_dir = None + if not hexagon_dir or not is_valid_sdk(hexagon_config, Path(hexagon_dir)): + missing.append("HEXAGON_SDK_ROOT") + else: + try: + expected_tools_dir = get_hexagon_tools_dir(Path(hexagon_dir)) + except RuntimeError: + pass + if not tools_dir or not expected_tools_dir or Path(tools_dir) != expected_tools_dir: + missing.append("HEXAGON_TOOLS_ROOT") + if not opencl_dir or not is_valid_sdk(opencl_config, Path(opencl_dir)): + missing.append("OPENCL_SDK_ROOT") + if missing: + raise RuntimeError( + f"Missing or invalid Windows SDK paths: {', '.join(missing)}. " + "Run scripts/snapdragon/setup-sdk.py first." + ) diff --git a/scripts/snapdragon/setup-sdk.py b/scripts/snapdragon/setup-sdk.py new file mode 100644 index 000000000000..ad828c079a6a --- /dev/null +++ b/scripts/snapdragon/setup-sdk.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +# +# Install Windows on Snapdragon SDKs for llama.cpp. +# + +import sys +import os +import argparse +import shutil +import logging +import json +import hashlib +import tarfile +import tempfile +from pathlib import Path +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from sdk import SDK_CONFIGS, get_hexagon_tools_dir, is_valid_sdk + + +logger = logging.getLogger("setup_sdk") + +DEFAULT_SDK_BASE_DIR = r"C:\Qualcomm" + + +def get_sdk_releases(config): + request = Request( + f"https://api.github.com/repos/{config['repo']}/releases?per_page=100", + headers={"Accept": "application/vnd.github+json", "User-Agent": "llama.cpp"}, + ) + try: + with urlopen(request, timeout=30) as response: + releases = json.load(response) + except (HTTPError, URLError, TimeoutError) as err: + raise RuntimeError(f"Cannot query {config['name']} releases: {err}") from err + + result = [] + for release in releases: + if release["draft"] or release["prerelease"]: + continue + version = release["tag_name"].removeprefix("v") + archive_name = f"{config['archive_prefix']}{version}-arm64-wos.tar.xz" + for asset in release["assets"]: + if asset["name"] != archive_name: + continue + result.append({ + "version": version, + "name": asset["name"], + "url": asset["browser_download_url"], + "sha256": (asset.get("digest") or "").removeprefix("sha256:"), + }) + return result + + +def list_sdk_releases(): + for config in SDK_CONFIGS: + logger.info("%s:", config["name"]) + releases = get_sdk_releases(config) + if not releases: + logger.info(" no Windows on Snapdragon releases found") + continue + for release in releases: + logger.info(" %s: %s", release["version"], release["name"]) + + +def get_sdk_release(config, version): + version = version or config["default_version"] + version = version.removeprefix("v") + for release in get_sdk_releases(config): + if release["version"] == version: + if not release["sha256"]: + raise RuntimeError(f"{config['name']} {version} does not provide a SHA-256 digest") + return release + raise RuntimeError( + f"No Windows on Snapdragon release for {config['name']} {version}. " + "Run scripts/snapdragon/setup-sdk.py --list-sdk-releases to see available versions." + ) + + +def sha256sum(path): + digest = hashlib.sha256() + with open(path, "rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def download_sdk(release, archive): + while True: + if archive.exists() and sha256sum(archive) == release["sha256"]: + logger.info("Using existing archive %s", archive) + return + + offset = archive.stat().st_size if archive.exists() else 0 + headers = {"User-Agent": "llama.cpp"} + if offset: + headers["Range"] = f"bytes={offset}-" + logger.info("Resuming download of %s at %d MiB", release["name"], offset // (1024 * 1024)) + else: + logger.info("Downloading %s", release["name"]) + + try: + with urlopen(Request(release["url"], headers=headers), timeout=30) as response: + mode = "ab" if offset and response.status == 206 else "wb" + with open(archive, mode) as file: + shutil.copyfileobj(response, file) + except HTTPError as err: + if err.code != 416: + raise RuntimeError(f"Cannot download {release['name']}: {err}") from err + archive.unlink(missing_ok=True) + continue + except (URLError, TimeoutError) as err: + raise RuntimeError(f"Cannot download {release['name']}: {err}") from err + + if sha256sum(archive) == release["sha256"]: + return + raise RuntimeError(f"SHA-256 mismatch for {archive}. Re-run the command to resume the download.") + + +def extract_sdk(config, archive, target_dir): + if not hasattr(tarfile, "data_filter"): + raise RuntimeError("SDK extraction requires Python 3.10.12 or later") + + with tempfile.TemporaryDirectory(prefix=f".{target_dir.name}.tmp-", dir=target_dir.parent) as staging_path: + staging_dir = Path(staging_path) + with tarfile.open(archive, "r:xz") as tar: + tar.extractall(staging_dir, filter=tarfile.data_filter) + + candidates = [staging_dir] + [path for path in staging_dir.iterdir() if path.is_dir()] + extracted_dirs = [path for path in candidates if is_valid_sdk(config, path)] + if len(extracted_dirs) != 1: + raise RuntimeError(f"{config['name']} archive does not contain the expected files") + extracted_dir = extracted_dirs[0] + + backup_dir = None + if target_dir.exists(): + backup_dir = target_dir.parent / f".{target_dir.name}.backup" + if backup_dir.exists(): + raise RuntimeError(f"Cannot replace {target_dir}: backup directory {backup_dir} already exists") + target_dir.replace(backup_dir) + try: + extracted_dir.replace(target_dir) + except Exception: + if backup_dir: + backup_dir.replace(target_dir) + raise + if backup_dir: + shutil.rmtree(backup_dir) + + +def install_sdk(config, version, base_dir, force): + version = (version or config["default_version"]).removeprefix("v") + target_dir = base_dir / config["parent_dir"] / version + if is_valid_sdk(config, target_dir) and not force: + logger.info("Using existing %s at %s", config["name"], target_dir) + return target_dir + + release = get_sdk_release(config, version) + target_dir.parent.mkdir(parents=True, exist_ok=True) + archive = target_dir.parent / release["name"] + download_sdk(release, archive) + logger.info("Extracting %s to %s", config["name"], target_dir) + extract_sdk(config, archive, target_dir) + archive.unlink(missing_ok=True) + return target_dir + + +def set_user_environment(values): + if os.name != "nt": + raise RuntimeError("SDK setup must run on Windows") + + import winreg + + with winreg.CreateKey(winreg.HKEY_CURRENT_USER, "Environment") as key: + for name, value in values.items(): + winreg.SetValueEx(key, name, 0, winreg.REG_SZ, str(value)) + os.environ[name] = str(value) + + import ctypes + + result = ctypes.c_ulong() + ctypes.windll.user32.SendMessageTimeoutW(0xffff, 0x001a, 0, "Environment", 0x0002, 5000, ctypes.byref(result)) + + +def setup_sdks(args): + base_dir = Path(args.sdk_base_dir).expanduser().resolve() + hexagon_config, opencl_config = SDK_CONFIGS + environment = {} + + if args.hexagon is not None: + hexagon_dir = install_sdk(hexagon_config, args.hexagon, base_dir, args.force) + environment["HEXAGON_SDK_ROOT"] = hexagon_dir + environment["HEXAGON_TOOLS_ROOT"] = get_hexagon_tools_dir(hexagon_dir) + if args.opencl is not None: + opencl_dir = install_sdk(opencl_config, args.opencl, base_dir, args.force) + environment["OPENCL_SDK_ROOT"] = opencl_dir + + set_user_environment(environment) + logger.info("SDK environment variables were updated. Start a new terminal before building.") + + +def main(): + logging.basicConfig(level=logging.INFO, format="%(message)s") + parser = argparse.ArgumentParser(description="Install Windows on Snapdragon SDKs for llama.cpp.") + parser.add_argument("--list-sdk-releases", action="store_true", help="List available Windows on Snapdragon SDK releases") + parser.add_argument("--sdk-base-dir", default=DEFAULT_SDK_BASE_DIR, help=r"SDK installation directory (default: C:\Qualcomm)") + parser.add_argument("--hexagon", nargs="?", const=SDK_CONFIGS[0]["default_version"], metavar="VERSION", help="Install the Hexagon SDK, optionally selecting a version") + parser.add_argument("--opencl", nargs="?", const=SDK_CONFIGS[1]["default_version"], metavar="VERSION", help="Install the OpenCL SDK, optionally selecting a version") + parser.add_argument("--force", action="store_true", help="Reinstall selected SDKs even when they already exist") + args = parser.parse_args() + + if args.list_sdk_releases: + if args.sdk_base_dir != DEFAULT_SDK_BASE_DIR or args.hexagon is not None or args.opencl is not None or args.force: + parser.error("Installation options cannot be combined with --list-sdk-releases") + list_sdk_releases() + return + if args.hexagon is None and args.opencl is None: + parser.error("Select at least one SDK with --hexagon or --opencl") + if os.name != "nt": + parser.error("SDK setup must run on Windows") + setup_sdks(args) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + logger.info("\nInterrupted by user.") + sys.exit(130) + except RuntimeError as err: + logger.error("Error: %s", err) + sys.exit(1) From 77f132cb1df1de6357617aeaf0ca04c02cf15fb1 Mon Sep 17 00:00:00 2001 From: Jeff Bolz Date: Sat, 29 Aug 2026 02:09:24 -0500 Subject: [PATCH 021/109] vulkan: Change mul_mat_id to pad K rather than N (#27925) The N padding is needed for mul_mat, but not mul_mat_id. For mul_mat_id, we indirect the row index through a shared memory lookup table which avoids any OOB row coordinate. But that callback doesn't bounds check K, so we actually need K padding instead. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 92 ++++++++++--------- .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 1 - .../vulkan-shaders/mul_mm_cm2.comp | 19 +++- .../ggml-vulkan/vulkan-shaders/mul_mmq.comp | 1 - tests/test-backend-ops.cpp | 7 +- 5 files changed, 68 insertions(+), 52 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 320127cdc5ac..39b4cd359803 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1347,7 +1347,6 @@ struct vk_mat_mat_id_push_constants { uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; - uint32_t padded_N; uint32_t n_experts; uint32_t hoist_row_ids; }; @@ -2403,9 +2402,8 @@ struct ggml_backend_vk_context { // Cache most recent tensor that was converted into prealloc_y, and what pipeline it used to convert. vk_pipeline_struct * prealloc_y_last_pipeline_used {}; const ggml_tensor * prealloc_y_last_tensor_used {}; - // True when prealloc_y holds the padded fp16 layout used by the coopmat2 B decode-vector callback. - // If false, then it's contiguous. - bool prealloc_y_last_decode_vector_staging {}; + // True when the K dimension in prealloc_y is padded. + bool prealloc_y_last_k_padded {}; // Track which nodes have been used since the last sync, and whether they were written to std::vector unsynced_nodes_written; @@ -8984,13 +8982,13 @@ static void ggml_vk_matmul_id( uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11, - uint32_t padded_n, bool hoist_row_ids) { + bool hoist_row_ids) { VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " << "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, - nei0, nei1, nbi1, ne11, padded_n, n_as, uint32_t(hoist_row_ids) }; + nei0, nei1, nbi1, ne11, n_as, uint32_t(hoist_row_ids) }; ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as }); } @@ -9455,27 +9453,27 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub if (y_non_contig) { if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0)); ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } if (quantize_y) { if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { 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); 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; + ctx->prealloc_y_last_k_padded = false; } } @@ -9734,27 +9732,27 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, d_Qy, d_Y); ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } if (quantize_y) { if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_quantize_q8_1(ctx, subctx, d_Qy, d_Y, y_ne); 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; + ctx->prealloc_y_last_k_padded = false; } } @@ -10234,8 +10232,6 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) || !ggml_vk_dim01_contiguous(src1); - const uint32_t y_staged_row_stride = y_decode_vector_staging ? (uint32_t)ggml_vk_align_size(ne10, 4) : (uint32_t)ne10; - const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; bool quantize_y = ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && !y_non_contig && (ne11 * ne10) % 4 == 0; @@ -10250,19 +10246,25 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } const bool qx_needs_dequant = mmp == nullptr || x_non_contig; - const bool qy_needs_dequant = !quantize_y && ((src1->type != f16_type && !y_f32_kernel) || y_non_contig); + bool qy_needs_dequant = !quantize_y && ((src1->type != f16_type && !y_f32_kernel) || y_non_contig); if (qx_needs_dequant) { // Fall back to dequant + f16 mulmat mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, (ggml_prec)dst->op_params[0]); } - // Not implemented - GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT - const ggml_type effective_src1_type = quantize_y ? GGML_TYPE_Q8_1 : (y_f32_kernel ? GGML_TYPE_F32 : src1->type); const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, nei1, qx_needs_dequant ? f16_type : src0->type, effective_src1_type)); + // Coopmat2 MUL_MAT_ID BK specialization constants in ggml_vk_load_shaders are at most 64. + const uint32_t y_staged_row_stride = ctx->device->coopmat2 && !quantize_y ? ggml_vk_align_size(ne10, 64) : ne10; + const bool y_needs_k_padding = ne10 != y_staged_row_stride; + const bool y_needs_reformat = y_non_contig || y_needs_k_padding; + qy_needs_dequant = qy_needs_dequant || y_needs_k_padding; + + // Not implemented + GGML_ASSERT(y_needs_reformat || !qy_needs_dequant); // NOLINT + const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && nei1 > 8; vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type); @@ -10270,10 +10272,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) { pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline); } - // Reserve extra storage in the N dimension for the Y matrix, so we can avoid bounds-checking - uint32_t padded_n = qy_needs_dequant ? ROUNDUP_POW2(ne11, pipeline->wg_denoms[1]) :ne11; const uint64_t x_ne = ggml_nelements(src0); - const uint64_t y_ne = (uint64_t)y_staged_row_stride * padded_n * ne12 * ne13; + const uint64_t y_ne = (uint64_t)y_staged_row_stride * ne11 * ne12 * ne13; const uint64_t d_ne = ggml_nelements(dst); const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); @@ -10292,7 +10292,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& y_staged_dst.type = f16_type; y_staged_dst.nb[0] = ggml_type_size(f16_type); y_staged_dst.nb[1] = y_staged_dst.nb[0] * y_staged_row_stride; - y_staged_dst.nb[2] = y_staged_dst.nb[1] * padded_n; + y_staged_dst.nb[2] = y_staged_dst.nb[1] * ne11; y_staged_dst.nb[3] = y_staged_dst.nb[2] * y_staged_dst.ne[2]; return y_staged_dst; }; @@ -10302,10 +10302,10 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } else { to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); } - if (y_non_contig) { + if (y_needs_reformat) { ggml_tensor y_staged_dst; const ggml_tensor * y_staged_dst_ptr = nullptr; - if (y_decode_vector_staging) { + if (y_needs_k_padding) { y_staged_dst = make_y_staged_dst(); y_staged_dst_ptr = &y_staged_dst; } @@ -10432,14 +10432,18 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_X, 0, x_sz } }, pc, { (uint32_t)x_ne, 1, 1}); } - if (y_non_contig) { + if (y_needs_reformat) { if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging != y_decode_vector_staging) { + ctx->prealloc_y_last_k_padded != y_needs_k_padding) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } - if (y_decode_vector_staging) { + if (y_needs_k_padding) { + GGML_ASSERT(y_sz % 4 == 0); + // Zero B padding because clamping only A can produce 0 * Inf or NaN. + subctx->s->buffer->buf.fillBuffer(d_Y->buffer, 0, y_sz, 0); + ggml_vk_sync_buffers(ctx, subctx); const ggml_tensor y_staged_dst = make_y_staged_dst(); const uint32_t y_staged_dst_type_size = ggml_type_size(y_staged_dst.type); ggml_vk_cpy_to_strided( @@ -10454,27 +10458,27 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = y_decode_vector_staging; + ctx->prealloc_y_last_k_padded = y_needs_k_padding; } } if (quantize_y) { if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { 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); 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; + ctx->prealloc_y_last_k_padded = false; } } ggml_vk_sync_buffers(ctx, subctx); uint32_t stride_batch_x = ne00*ne01; - uint32_t stride_b_y = y_decode_vector_staging ? y_staged_row_stride : ne10; - uint32_t stride_batch_y = y_decode_vector_staging ? y_staged_row_stride * padded_n : ne10*ne11; + uint32_t stride_b_y = y_needs_k_padding ? y_staged_row_stride : ne10; + uint32_t stride_batch_y = y_needs_k_padding ? y_staged_row_stride * ne11 : ne10*ne11; if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); @@ -10491,13 +10495,13 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& { d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf, ne01, ne21, ne10, ne10, stride_b_y, ne01, stride_batch_x, stride_batch_y, ne20*ne21, - n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n, hoist_row_ids + n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, hoist_row_ids ); // NOLINT if (x_non_contig || qx_needs_dequant) { ctx->prealloc_x_need_sync = true; } - if (y_non_contig || quantize_y) { + if (y_needs_reformat || quantize_y) { ctx->prealloc_y_need_sync = true; } ctx->prealloc_split_k_need_sync = true; @@ -10648,27 +10652,27 @@ static void ggml_vk_mul_mat_vec_id_q_f16(ggml_backend_vk_context * ctx, vk_conte GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, d_Qy, d_Y); ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } if (quantize_y) { if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_quantize_q8_1(ctx, subctx, d_Qy, d_Y, y_ne); 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; + ctx->prealloc_y_last_k_padded = false; } } @@ -15520,7 +15524,7 @@ static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx, vk_contex ctx->prealloc_y = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_y); ctx->prealloc_y_last_pipeline_used = nullptr; ctx->prealloc_y_last_tensor_used = nullptr; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } if (ctx->prealloc_split_k == nullptr || (ctx->prealloc_size_split_k > 0 && ctx->prealloc_split_k->size < ctx->prealloc_size_split_k)) { VK_LOG_MEMORY("ggml_vk_preallocate_buffers(split_k_size: " << ctx->prealloc_size_split_k << ")"); @@ -16145,7 +16149,7 @@ static void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) { VK_LOG_DEBUG("ggml_vk_graph_cleanup()"); ctx->prealloc_y_last_pipeline_used = {}; ctx->prealloc_y_last_tensor_used = nullptr; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; ctx->unsynced_nodes_written.clear(); ctx->unsynced_nodes_read.clear(); @@ -16197,7 +16201,7 @@ static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) { ctx->prealloc_y_last_pipeline_used = nullptr; ctx->prealloc_y_last_tensor_used = nullptr; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; ctx->prealloc_size_x = 0; ctx->prealloc_size_y = 0; @@ -17395,7 +17399,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->prealloc_y_last_pipeline_used = nullptr; ctx->prealloc_y_last_tensor_used = nullptr; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; if (ctx->prealloc_size_add_rms_partials) { ggml_vk_preallocate_buffers(ctx, nullptr); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index c1ccac7aa202..63c4aaebcb1a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -88,7 +88,6 @@ layout (push_constant) uniform parameter uint nei1; uint nbi1; uint ne11; - uint padded_N; uint n_experts; uint hoist_row_ids; #else diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp index cf78474a9258..27f3178e7f26 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -56,6 +56,8 @@ layout (push_constant) uniform parameter uint nei1; uint nbi1; uint ne11; + uint n_experts; + uint hoist_row_ids; #else uint base_work_group_z; uint num_batches; @@ -64,12 +66,8 @@ layout (push_constant) uniform parameter uint ne12; uint broadcast2; uint broadcast3; -#endif // N dimension for the B matrix can be >= p.N uint padded_N; -#ifdef MUL_MAT_ID - uint n_experts; - uint hoist_row_ids; #endif } p; @@ -332,7 +330,9 @@ void main() { tensorLayoutNV<2> tensorLayoutA = createTensorLayoutNV(2); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutAClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); tensorLayoutNV<2> tensorLayoutB = createTensorLayoutNV(2); +#ifndef MUL_MAT_ID tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutBClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); +#endif tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); #if QUANT_K > 1 @@ -345,12 +345,19 @@ void main() { // Use end_k rather than p.K as the dimension because that's what // we need to bound check against when using split_k. - // Bounds check B against padded_N, but bounds check D against N. tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k); +#ifdef MUL_MAT_ID + // MUL_MAT_ID pads each B row to stride_b so partial K tiles read zeros without clamping. + tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, BN, p.stride_b); +#else + // Bounds check B against padded_N, but bounds check D against N. tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, p.padded_N, end_k); +#endif tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.N, p.M); tensorLayoutAClamp = setTensorLayoutDimensionNV(tensorLayoutAClamp, p.M, end_k); +#ifndef MUL_MAT_ID tensorLayoutBClamp = setTensorLayoutDimensionNV(tensorLayoutBClamp, p.padded_N, end_k); +#endif tensorLayoutD = setTensorLayoutStrideNV(tensorLayoutD, p.stride_d, 1); @@ -527,7 +534,9 @@ void main() { tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); +#ifndef MUL_MAT_ID tensorLayoutBClamp = setTensorLayoutStrideNV(tensorLayoutBClamp, stride_b, 1); +#endif uint k_iters = (end_k - start_k + BK - 1) / BK; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp index c2d84c05b40d..1fbcbf6c9332 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp @@ -56,7 +56,6 @@ layout (push_constant) uniform parameter uint nei1; uint nbi1; uint ne11; - uint padded_N; uint n_experts; uint hoist_row_ids; #else diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 6be83ac161bd..4a7a0623174c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9367,7 +9367,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, b, 50, 200, 64)); } - test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1)); + // For issue 27873 + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ2_XXS, GGML_TYPE_F32, 1, 1, false, 1, 8192, 4096)); + + for (int k : {1, 63, 65}) { + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, k)); + } test_cases.emplace_back(new test_mul_mat_id_fusion(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 32, 32, 32, 3)); // gpt-oss issue with Vulkan mmq_id From 5ea1b124e7dfcdb80d7291be188efc7d0b485d66 Mon Sep 17 00:00:00 2001 From: Jhen-Jie Hong Date: Sat, 29 Aug 2026 15:12:23 +0800 Subject: [PATCH 022/109] metal : add fa-vec tunings for M1 Max (#27932) --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 188 ++++++++++++++++++++++ 1 file changed, 188 insertions(+) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 285590d1601b..6cfa73e65153 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -279,6 +279,194 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 128, 128, 2, 1 }, { 1, 4 } }, From c9ca51c1f6b18427cde490c7c7eba11d87a96b2d Mon Sep 17 00:00:00 2001 From: Jeff Bolz Date: Sat, 29 Aug 2026 02:59:48 -0500 Subject: [PATCH 023/109] vulkan: combine duplicated fastdiv functions, rename the one optimizing small divs (#27526) * vulkan: combine duplicated fastdiv functions, rename the one optimizing small divs * remove one more fastdiv --- .../ggml-vulkan/vulkan-shaders/conv2d_mm.comp | 9 +-------- .../ggml-vulkan/vulkan-shaders/conv3d_mm.comp | 9 +-------- .../vulkan-shaders/count_experts.comp | 9 +-------- .../vulkan-shaders/generic_unary_head.glsl | 14 ++------------ .../ggml-vulkan/vulkan-shaders/glu_head.glsl | 8 ++------ .../ggml-vulkan/vulkan-shaders/sum_rows.glsl | 10 ++-------- ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl | 18 +++++++++++++++--- 7 files changed, 24 insertions(+), 53 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp index 99400098bf2b..c64004cdc48e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp @@ -19,6 +19,7 @@ #endif #include "types.glsl" +#include "utils.glsl" // shape notation: [dim(N), ..., dim(0)] -- stride(dim(j)) >= stride(dim(i)) if i > j layout(binding = 0) readonly buffer A { @@ -193,14 +194,6 @@ uint32_t Br = tid / BS_NPQ; uint32_t Bc = tid % BS_NPQ; const uint32_t BrpWg = WG_SIZE / BS_NPQ; -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - #ifdef COOPMAT2 #define ACC_TYPE float16_t diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp index f66f299f6dae..d5ce4290b930 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp @@ -15,6 +15,7 @@ #endif #include "types.glsl" +#include "utils.glsl" // shape notation: [dim(N), ..., dim(0)] -- stride(dim(j)) >= stride(dim(i)) if i > j layout(binding = 0) readonly buffer A { @@ -178,14 +179,6 @@ uint32_t Br = tid / BS_NPQ; uint32_t Bc = tid % BS_NPQ; const uint32_t BrpWg = WG_SIZE / BS_NPQ; -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - void split_crs(uint32_t crs_idx, out uint32_t ic, out uint32_t kd, out uint32_t kh, out uint32_t kw) { const uint32_t KHKW = KH * KW; const uint32_t KDKHKW = KD * KHKW; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp index 83c56fce5209..ef659959d950 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp @@ -8,6 +8,7 @@ #endif #include "types.glsl" +#include "utils.glsl" layout (push_constant) uniform parameter { @@ -33,14 +34,6 @@ shared uint vals[BLOCK_SIZE]; shared uint offsets[BLOCK_SIZE]; shared uint cursors[BLOCK_SIZE]; -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - // data_d layout when p.hoist_row_ids is set: // [0, n_experts) per-expert row count // [n_experts, 2*n_experts) per-expert start offset into the row id region diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.glsl index 9d4176f3f967..e13de9a00f2f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.glsl @@ -1,6 +1,8 @@ #extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_control_flow_attributes : require +#include "utils.glsl" + layout (push_constant) uniform parameter { uint ne; @@ -32,18 +34,6 @@ uint get_idx() { uint get_aoffset() { return p.misalign_offsets >> 16; } uint get_doffset() { return p.misalign_offsets & 0xFFFF; } -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - -uint fastdiv_L(uint packed, uint slot) { - return (packed >> (slot * 8)) & 0x3Fu; -} - uint src0_idx(uint idx) { const uint i03 = fastdiv(idx, p.ne0_012mp, fastdiv_L(p.ne0_Ls, 0)); const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/glu_head.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/glu_head.glsl index c3cae736f977..fc2951ec2e56 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/glu_head.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/glu_head.glsl @@ -1,5 +1,7 @@ #extension GL_EXT_shader_16bit_storage : require +#include "utils.glsl" + layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; @@ -39,9 +41,3 @@ uint get_aoffset() { return p.misalign_offsets >> 16; } uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; } uint get_doffset() { return p.misalign_offsets & 0xFF; } -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.glsl index 2b841baa6bf2..1cb0f7827a38 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.glsl @@ -1,4 +1,6 @@ +#include "utils.glsl" + // vk_op_sum_rows_push_constants layout (push_constant) uniform parameter { @@ -15,11 +17,3 @@ layout (push_constant) uniform parameter uint get_aoffset() { return p.misalign_offsets >> 16; } uint get_doffset() { return p.misalign_offsets & 0xFFFF; } -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl index dc4a1e6d96ba..8aac64d75932 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl @@ -9,14 +9,26 @@ uint fastmod(uint a, uint b) { return a % b; } -uint fastdiv(uint a, uint b) { +// see init_fastdiv_values in ggml-vulkan.cpp +uint fastdiv(uint n, uint mp, uint L) { + uint msbs, lsbs; + // msbs = mulhi(n, mp) + umulExtended(n, mp, msbs, lsbs); + return (msbs + n) >> L; +} + +uint fastdiv_L(uint packed, uint slot) { + return (packed >> (slot * 8)) & 0x3Fu; +} + +uint fastdiv_small(uint a, uint b) { return (a < b) ? 0 : (a / b); } void get_indices(uint idx, out uint i00, out uint i01, out uint i02, out uint i03, uint ne00, uint ne01, uint ne02, uint ne03) { - i03 = fastdiv(idx, (ne02*ne01*ne00)); + i03 = fastdiv_small(idx, (ne02*ne01*ne00)); const uint i03_offset = i03 * ne02*ne01*ne00; - i02 = fastdiv((idx - i03_offset), (ne01*ne00)); + i02 = fastdiv_small((idx - i03_offset), (ne01*ne00)); const uint i02_offset = i02*ne01*ne00; i01 = (idx - i03_offset - i02_offset) / ne00; i00 = idx - i03_offset - i02_offset - i01*ne00; From cc83d7b4824f73cfdda4dfbb47ee39804f71b328 Mon Sep 17 00:00:00 2001 From: Nick Farrell Date: Sat, 29 Aug 2026 19:00:09 +1000 Subject: [PATCH 024/109] sycl: make --fit respect --fit-target better (#27629) improve the --fit algorithm to take into account the actual peak required VRAM for a given context size on a SYCL backend. This includes both properly accounting for how much VRAM is required when the allocated context is fully used (which makes the reported context drop below what it did before, but stop it OOMing) as well as preventing some overly-conservative calculations which meant too much VRAM was being reserved. Tested on a Arc b70 with unsloth's qwen3.8 (Q4_K_XL), able to get 262144 context, fully usable, with q8_0 KV and MTP and 4k ubatch size using --fit-target 1 --- ggml/src/ggml-sycl/fattn-common.hpp | 22 ++++---- ggml/src/ggml-sycl/fattn-onednn.cpp | 84 ++++++++++++++++++----------- ggml/src/ggml-sycl/fattn-onednn.hpp | 6 ++- ggml/src/ggml-sycl/fattn.cpp | 73 +++++++++++++++++++++++++ ggml/src/ggml-sycl/fattn.hpp | 18 +++++++ ggml/src/ggml-sycl/ggml-sycl.cpp | 5 +- 6 files changed, 167 insertions(+), 41 deletions(-) diff --git a/ggml/src/ggml-sycl/fattn-common.hpp b/ggml/src/ggml-sycl/fattn-common.hpp index c6cc13cfb005..82813f7a99a7 100644 --- a/ggml/src/ggml-sycl/fattn-common.hpp +++ b/ggml/src/ggml-sycl/fattn-common.hpp @@ -6,6 +6,7 @@ #include "convert.hpp" #include "vecdotq.hpp" #include "fattn-buffers.hpp" +#include "fattn.hpp" #include "ggml.h" @@ -926,6 +927,7 @@ void launch_fattn( ggml_sycl_fattn_alloc K_f16(fbuf.K); ggml_sycl_fattn_alloc V_f16(fbuf.V); + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); ggml_sycl_pool_alloc KV_max(pool); ggml_sycl_pool_alloc dst_tmp(pool); ggml_sycl_pool_alloc dst_tmp_meta(pool); @@ -944,10 +946,11 @@ void launch_fattn( const size_t bs = ggml_blck_size(K->type); const size_t ts = ggml_type_size(K->type); - K_f16.alloc(ggml_nelements(K)); + sycl::half * K_f16_ptr = extra.K_buffer_ptr ? (sycl::half *) extra.K_buffer_ptr + : K_f16.alloc(ggml_nelements(K)); if (ggml_is_contiguously_allocated(K)) { to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst); - to_fp16(K_data, K_f16.ptr, ggml_nelements(K), main_stream); + to_fp16(K_data, K_f16_ptr, ggml_nelements(K), main_stream); nb11 = nb11 * bs * sizeof(sycl::half) / ts; nb12 = nb12 * bs * sizeof(sycl::half) / ts; @@ -958,13 +961,13 @@ void launch_fattn( const int64_t s01 = nb11 / ts; const int64_t s02 = nb12 / ts; const int64_t s03 = nb13 / ts; - to_fp16(K_data, K_f16.ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream); + to_fp16(K_data, K_f16_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream); nb11 = K->ne[0] * sizeof(sycl::half); nb12 = K->ne[1] * nb11; nb13 = K->ne[2] * nb12; } - K_data = (char *) K_f16.ptr; + K_data = (char *) K_f16_ptr; } if (need_f16_V && V->type != GGML_TYPE_F16) { @@ -977,11 +980,12 @@ void launch_fattn( const size_t bs = ggml_blck_size(V->type); const size_t ts = ggml_type_size(V->type); - V_f16.alloc(ggml_nelements(V)); + sycl::half * V_f16_ptr = extra.V_buffer_ptr ? (sycl::half *) extra.V_buffer_ptr + : V_f16.alloc(ggml_nelements(V)); if (ggml_is_contiguously_allocated(V)) { to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst); - to_fp16(V_data, V_f16.ptr, ggml_nelements(V), main_stream); - V_data = (char *) V_f16.ptr; + to_fp16(V_data, V_f16_ptr, ggml_nelements(V), main_stream); + V_data = (char *) V_f16_ptr; nb21 = nb21 * bs * sizeof(sycl::half) / ts; nb22 = nb22 * bs * sizeof(sycl::half) / ts; @@ -992,13 +996,13 @@ void launch_fattn( const int64_t s01 = nb21 / ts; const int64_t s02 = nb22 / ts; const int64_t s03 = nb23 / ts; - to_fp16(V_data, V_f16.ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream); + to_fp16(V_data, V_f16_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream); nb21 = V->ne[0] * sizeof(sycl::half); nb22 = V->ne[1] * nb21; nb23 = V->ne[2] * nb22; } - V_data = (char *) V_f16.ptr; + V_data = (char *) V_f16_ptr; } } diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index d41c2ddce345..4349363a3d3e 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -14,9 +14,21 @@ // set minimum query length to treat as prefill (32) #define GGML_SYCL_FA_ONEDNN_MIN_Q 32 -bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { +bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V) { + if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) { + return false; + } + auto bindable = [](const ggml_tensor * t) { + return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 && + t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0; + }; + return bindable(K) && bindable(V); +} + +bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit) { #if !GGML_SYCL_DNNL GGML_UNUSED(dst); + GGML_UNUSED(use_shape_limit); return false; #else if (!g_ggml_sycl_fa_onednn) { @@ -44,7 +56,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { if (!k_ok || !v_ok) { return false; } - if (Q->ne[1] < 32 || K->ne[1] < 1024) { + if (use_shape_limit && (Q->ne[1] < 32 || K->ne[1] < 1024)) { return false; } for (const ggml_tensor * t : {K, V}) { @@ -94,7 +106,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { return false; } // Prefill only. - if (Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) { + if (use_shape_limit && Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) { return false; } return true; @@ -240,9 +252,16 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso dnnl::engine eng = ctx.engine_dnnl(stream); dnnl::stream strm = ctx.stream_dnnl(stream); + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); + // Q: always f32 -- copy to dense f16. - ggml_sycl_pool_alloc Qf(ctx.pool(), (size_t) H * q * d); - cont_to_f16_sycl((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream); + std::optional> Qf_pool; + sycl::half * Qf_ptr = (sycl::half *) extra.Q_buffer_ptr; + if (!Qf_ptr) { + Qf_pool.emplace(ctx.pool(), (size_t) H * q * d); + Qf_ptr = Qf_pool->get(); + } + cont_to_f16_sycl((const char *) Q->data, Qf_ptr, d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream); // K/V: bind the f16 cache in place. llama.cpp permutes it to [token][head][dim], so its head // plane is strided rather than dense, which is what an explicit stride vector expresses. @@ -253,11 +272,12 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso std::array v_str = k_str; std::optional> Kf_pool; std::optional> Vf_pool; + // Helper: hand out reserved space, or fall back to the pool. + auto stage_k = [&](size_t n) { if (extra.K_buffer_ptr) { return (sycl::half *) extra.K_buffer_ptr; } + Kf_pool.emplace(ctx.pool(), n); return Kf_pool->get(); }; + auto stage_v = [&](size_t n) { if (extra.V_buffer_ptr) { return (sycl::half *) extra.V_buffer_ptr; } + Vf_pool.emplace(ctx.pool(), n); return Vf_pool->get(); }; - auto bindable = [](const ggml_tensor * t) { - return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 && - t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0; - }; auto elem_strides = [](const ggml_tensor * t) { const int64_t s1 = (int64_t) (t->nb[1] / t->nb[0]); const int64_t s2 = (int64_t) (t->nb[2] / t->nb[0]); @@ -266,22 +286,19 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso return std::array{ s3, s2, s2, s1, 1 }; }; - if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && bindable(K) && bindable(V)) { + if (ggml_sycl_fattn_onednn_binds_kv(K, V)) { K_ptr = (sycl::half *) K->data; V_ptr = (sycl::half *) V->data; k_str = elem_strides(K); v_str = elem_strides(V); } else if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) { - Kf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d); - Vf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d); - cont_to_f16_sycl((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream); - cont_to_f16_sycl((const char *) V->data, Vf_pool->get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream); - K_ptr = Kf_pool->get(); - V_ptr = Vf_pool->get(); + K_ptr = stage_k((size_t) Hkv * seq * d); + V_ptr = stage_v((size_t) Hkv * seq * d); + cont_to_f16_sycl((const char *) K->data, K_ptr, d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream); + cont_to_f16_sycl((const char *) V->data, V_ptr, d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream); } else if (ggml_is_quantized(K->type)) { // Quantized K/V: dequant to dense F16 using pool, same lifetime as F16 path. - Kf_pool.emplace(ctx.pool(), ggml_nelements(K)); - K_ptr = Kf_pool->get(); + K_ptr = stage_k((size_t) ggml_nelements(K)); { const char * K_data = (const char *)K->data; const bool k_non_dense = ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]) && K->ne[2] > 1; @@ -315,8 +332,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso // data pointer), their logical values differ because the quantized // elements at different positions/offsets represent different K/V // data. Master's F16 path also never aliases K and V. - Vf_pool.emplace(ctx.pool(), ggml_nelements(V)); - V_ptr = Vf_pool->get(); + V_ptr = stage_v((size_t) ggml_nelements(V)); { const char * V_data = (const char *)V->data; const bool v_non_dense = ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]) && V->ne[2] > 1; @@ -347,12 +363,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso } } else { // F32: strided copy to dense F16 via cont_to_f16_sycl. - Kf_pool.emplace(ctx.pool(), ggml_nelements(K)); - K_ptr = Kf_pool->get(); + K_ptr = stage_k((size_t) ggml_nelements(K)); cont_to_f16_sycl((const char *) K->data, K_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], K->nb[1], K->nb[2], K->nb[3], stream); - Vf_pool.emplace(ctx.pool(), ggml_nelements(V)); - V_ptr = Vf_pool->get(); + V_ptr = stage_v((size_t) ggml_nelements(V)); cont_to_f16_sycl((const char *) V->data, V_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], V->nb[1], V->nb[2], V->nb[3], stream); } @@ -366,11 +380,21 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso // instead -- the value is captured into the command, so no host memory has to outlive the // call, and the enqueue stays async. const sycl::half scale_h = (sycl::half) (1.0f / kq_scale); - ggml_sycl_pool_alloc scbuf(ctx.pool(), 1); - sycl::half * const scale_dev = scbuf.get(); + std::optional> scbuf; + sycl::half * scale_dev = (sycl::half *) extra.scale_buffer_ptr; + if (!scale_dev) { + scbuf.emplace(ctx.pool(), 1); + scale_dev = scbuf->get(); + } stream->single_task([=]() { *scale_dev = scale_h; }); - ggml_sycl_pool_alloc outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d] + // f16 contiguous SDPA out [mb,H,q,d] + std::optional> outf_pool; + sycl::half * outf_ptr = (sycl::half *) extra.out_buffer_ptr; + if (!outf_ptr) { + outf_pool.emplace(ctx.pool(), (size_t) H * q * d); + outf_ptr = outf_pool->get(); + } // compile once per (device, shape, KV strides), reuse across layers/calls. Stride 2 always // repeats stride 1 and stride 4 is always 1, so the key covers every entry that can differ. @@ -392,7 +416,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso } auto id2ptr = [&](size_t r) -> void * { - if (r == E.id_q) return Qf.get(); + if (r == E.id_q) return Qf_ptr; if (r == E.id_k) return K_ptr; if (r == E.id_v) return V_ptr; if (r == E.id_scale) return scale_dev; @@ -404,10 +428,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso for (auto & lt : E.ins) { ti.emplace_back(lt, eng, id2ptr(lt.get_id())); } - tensor to(E.out, eng, outf.get()); + tensor to(E.out, eng, outf_ptr); E.cp.execute(strm, ti, {to}); - permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream); + permute_sdpa_out_sycl(outf_ptr, (float *) dst->data, mb, H, q, d, stream); // Single device needs no sync: the dnnl stream wraps this same in-order queue, so the SDPA // serializes with the staging kernels before it and the permute/pool reuse after it. The // garbage output formerly blamed on the missing sync here was the scale use-after-return diff --git a/ggml/src/ggml-sycl/fattn-onednn.hpp b/ggml/src/ggml-sycl/fattn-onednn.hpp index d3019e876889..9669d1bd27a6 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.hpp +++ b/ggml/src/ggml-sycl/fattn-onednn.hpp @@ -5,7 +5,11 @@ // Static-only check: fused-XMX oneDNN Graph SDPA path==flash-attn op // (f16 KV, no softcap/ALiBi, single stream, tuned head_dim, prefill-sized q.) -bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst); +bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit = true); + +// True when the oneDNN path binds an F16 KV cache in place instead of staging a dense copy of +// it. Depends only on the types and strides of K and V, so the answer holds for every call. +bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V); // Run flash attention through oneDNN's fused xmx SDPA // execute the cached SDPA partition, write the f32 dst. Falls back to the TILE kernel on any failure. diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.cpp index a85bca7cb3cd..b73e6d46ffa0 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.cpp @@ -378,3 +378,76 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst) { return ggml_sycl_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE; } + +static uintptr_t ggml_sycl_fattn_reserve_halves(ggml_sycl_fattn_extra & extra, size_t n_halves) { + if (n_halves == 0) { + return 0; + } + extra.end = GGML_PAD(extra.end, SYCL_BUFFER_ALIGNMENT); + const uintptr_t block = extra.end; + extra.end += n_halves * sizeof(sycl::half); + return block; +} + +ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst) { + ggml_sycl_fattn_extra extra; + + extra.end = (uintptr_t) dst->data + ggml_nbytes(dst); + + if (dst->op != GGML_OP_FLASH_ATTN_EXT) { + return extra; + } + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + if (!Q || !K || !V) { + return extra; + } + + const int64_t d = K->ne[0]; + const int64_t H = Q->ne[2]; + const int64_t q = Q->ne[1]; + + // calculate the worst-case memory consumption across all kernels + const bool onednn_supported = ggml_sycl_flash_attn_ext_onednn_supported(dst, /* use_shape_limit */ false); + + const bool tile_needs_K = K->type != GGML_TYPE_F16; + const bool tile_needs_V = V->type != GGML_TYPE_F16; + + const bool V_is_K_view = V->view_src && + (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)); + + size_t need_K = 0, need_V = 0, need_Q = 0, need_out = 0, need_scale = 0; + if (onednn_supported) { + need_Q = (size_t) H * q * d; + need_out = (size_t) H * q * d; + need_scale = 1; + // an f16 cache is bound in place, so it needs no staging copy + if (!ggml_sycl_fattn_onednn_binds_kv(K, V)) { + need_K = (size_t) ggml_nelements(K); + need_V = (size_t) ggml_nelements(V); + } + } + if (tile_needs_K) { + need_K = std::max(need_K, (size_t) ggml_nelements(K)); + } + if (tile_needs_V) { + need_V = std::max(need_V, (size_t) ggml_nelements(V)); + } + + extra.Q_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_Q); + extra.K_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_K); + extra.V_buffer_ptr = (V_is_K_view && !onednn_supported && need_V) + ? extra.K_buffer_ptr + : ggml_sycl_fattn_reserve_halves(extra, need_V); + extra.scale_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_scale); + extra.out_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_out); + + return extra; +} + +size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst) { + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); + return (size_t) (extra.end - (uintptr_t) dst->data); +} diff --git a/ggml/src/ggml-sycl/fattn.hpp b/ggml/src/ggml-sycl/fattn.hpp index c093970a3fed..f803aa2a804a 100644 --- a/ggml/src/ggml-sycl/fattn.hpp +++ b/ggml/src/ggml-sycl/fattn.hpp @@ -19,6 +19,24 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst); +// Scratch that flash attention needs beyond the output tensor +struct ggml_sycl_fattn_extra { + uintptr_t K_buffer_ptr = 0; // F16 copy of the K cache + uintptr_t V_buffer_ptr = 0; // F16 copy of the V cache + uintptr_t Q_buffer_ptr = 0; // dense F16 copy of Q, oneDNN only + uintptr_t scale_buffer_ptr = 0; // the softmax scale as an F16 scalar, oneDNN only + uintptr_t out_buffer_ptr = 0; // F16 SDPA output before conversion to F32, oneDNN only + uintptr_t end = 0; // one past the last reserved byte; sizes the allocation +}; + +// ggml_sycl_fattn_get_extra() is the single source of truth for the layout: it both sizes +// the reservation and hands out the pointers, so the two cannot disagree. +// Each field is the address of one reserved block, or 0 if that block was not reserved, +// in which case the caller allocates from the scratch pool instead. +ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst); + +size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst); + void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst); #endif // GGML_SYCL_FATTN_HPP diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d834e..dc8a1744323e 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -955,7 +955,10 @@ static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_typ } static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - size_t size = ggml_nbytes(tensor); + // Reserve the additional scratch so it's visible to the graph allocator + size_t size = tensor->op == GGML_OP_FLASH_ATTN_EXT + ? ggml_sycl_flash_attn_ext_get_alloc_size(tensor) + : ggml_nbytes(tensor); int64_t ne0 = tensor->ne[0]; if (ggml_is_quantized(tensor->type)) { From 17252c769a63c1cb650ce98ae309cf4de0da7778 Mon Sep 17 00:00:00 2001 From: Niklas Wenzel Date: Sat, 29 Aug 2026 14:50:13 +0200 Subject: [PATCH 025/109] metal : add remaining fa-vec tunings for M4 Pro (#27915) --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 150 ++++++++++++++++++++++ 1 file changed, 150 insertions(+) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 6cfa73e65153..26b634d21432 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -981,6 +981,156 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 512, 512, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 2, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 3, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 4 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 3, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 4 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } }, From 3173a56471c1753650cd806694145ffd6dcace67 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sat, 29 Aug 2026 17:55:15 +0300 Subject: [PATCH 026/109] metal : assert shared memory padding (#27951) * metal : assert shared memory padding * cont : add ref --- ggml/src/ggml-metal/ggml-metal-device.cpp | 3 ++- ggml/src/ggml-metal/ggml-metal-device.m | 3 +++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index a82caa5e4303..4e855be44670 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -593,7 +593,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me // - sgptg floats for shared_x_dt (nsg) // - sgptg floats for shared_dA (nsg) // Total: nsg * (32 + 2) floats - res.smem = (32 + 2)*sizeof(float)*nsg; + res.smem = GGML_PAD((32 + 2)*sizeof(float)*nsg, 16); return res; } @@ -1029,6 +1029,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(g } res.smem = (size_t) ne02*ne20*sizeof(uint16_t); + res.smem = GGML_PAD(res.smem, 16); return res; } diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 41ce90dc8a93..85c0f576b9fb 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -800,6 +800,9 @@ void ggml_metal_encoder_set_buffer(ggml_metal_encoder_t encoder, struct ggml_met } void ggml_metal_encoder_set_threadgroup_memory_size(ggml_metal_encoder_t encoder, size_t size, int idx) { + // ref: https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/setthreadgroupmemorylength(_:index:) + GGML_ASSERT(size % 16 == 0); + [encoder->obj setThreadgroupMemoryLength:size atIndex:idx]; } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index f6f2fdc86c6a..89c8483b3714 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -948,7 +948,7 @@ int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); - ggml_metal_encoder_set_threadgroup_memory_size(enc, nsg * sizeof(float), 0); + ggml_metal_encoder_set_threadgroup_memory_size(enc, GGML_PAD(nsg * sizeof(float), 16), 0); ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, nth, 1, 1); From c841aeeb8bb2fe417038dadfa9b007cf1a9ef950 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Sat, 29 Aug 2026 10:46:27 -0700 Subject: [PATCH 027/109] opencl: use a better matmul path on two Adreno GPU generations (#27640) * opencl: default the Adreno xmem F16xF32 GEMM on for X2E kernel_mul_mm_f16_f32_l4_lm is the slowest matmul this backend has on Adreno: on the X2-90 it runs the gpt-oss-20b attention projections at roughly a quarter of what the tuned dense q4_0 GEMM reaches on the same device. That matters for any model whose non-expert weights stay f16 -- the stock gpt-oss-20b release is exactly that, and its prefill spends 40.8% of GPU time in that one kernel. The xmem route already existed but was left opt-in, so nobody hit it. Worth about 25% prefill on gpt-oss-20b on an Adreno X2-90. Gated to X2E: the Adreno 840 measures neutral. Decode is untouched -- the dispatch gate needs N >= 16. It is worth nothing on the q8attn variant, whose attention weights already take the dp4a dense GEMM. The env var was presence-tested before, so =0 previously enabled it; it is now atoi()'d. MUL_MAT 963 OK / 0 FAIL on both arms. * opencl: bypass the tiled f32 GEMM on the Adreno A7X The A7X (E031.41) compiler executes kernel_mul_mm_f32_f32_l4_lm at roughly a tenth of what the same silicon reaches in its own f16 and q4_K kernels. It allocates 488 B/WI of private memory against 304 for the same source on the following generation, i.e. the older register allocator spills in the K-loop. Models with per-layer F32 projection pairs kept F32 by quantization policy land on this kernel twice per layer, and it dominates their prefill on that part. Route batched f32xf32 (ne11 > 8) around the tiled path on the A7X and let it fall through to the per-row f32 kernel, which that compiler handles fine; small batches keep the tiled path. Weights stay GPU-resident, so decode placement is untouched -- declining the op in supports_op instead was measured first and rejected, because the per-layer CPU round-trips cost more decode than the prefill it gained. Worth about 9% prefill on gemma-3n-E4B on an Adreno 740, with MUL_MAT counts identical on and off. No other generation is affected. Override with GGML_OPENCL_A7X_F32_LM_BYPASS=0. * opencl: enable xmem GEMM for adreno by default --------- Co-authored-by: Li He --- ggml/src/ggml-opencl/ggml-opencl.cpp | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 6ae83449b082..426aac523164 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -6059,9 +6059,13 @@ static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev) { } #ifdef GGML_OPENCL_USE_ADRENO_KERNELS - // determine whether to use Adreno xmem GEMM - backend_ctx->adreno_xmem_gemm_enabled = getenv("GGML_OPENCL_ADRENO_XMEM_GEMM") != nullptr && - backend_ctx->gpu_family == GPU_FAMILY::ADRENO; + // Adreno xmem F16xF32 GEMM, default on adreno, opt out with GGML_OPENCL_ADRENO_XMEM_GEMM=0. + // This helps models with f16 attention weights, e.g., gpt-oss-20b-f16 + { + const char * xmem_env = getenv("GGML_OPENCL_ADRENO_XMEM_GEMM"); + backend_ctx->adreno_xmem_gemm_enabled = backend_ctx->gpu_family == GPU_FAMILY::ADRENO && + (xmem_env ? atoi(xmem_env) != 0 : true); + } #endif // determine whether to use large buffer for Adreno @@ -19534,9 +19538,18 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co // GEMM using local memory // Current BK = 16, so ne00 % 16 == 0 + // + // Certain A7X compiler (E031.41) executes kernel_mul_mm_f32_f32_l4_lm poorly; + // matrices with ne11 <= 8 appears OK. + // Fallback to the MV style kernels for A7x and ne11 > 8. + // Override with GGML_OPENCL_A7X_F32_LM_BYPASS=0. + static const char * a7x_f32lm_env = getenv("GGML_OPENCL_A7X_F32_LM_BYPASS"); + static const bool a7x_f32lm_bypass = (a7x_f32lm_env == nullptr || a7x_f32lm_env[0] != '0'); if (src1t == GGML_TYPE_F32 && ne00 % 16 == 0 && - ne11 > 1) { + ne11 > 1 && + !(a7x_f32lm_bypass && src0t == GGML_TYPE_F32 && ne11 > 8 && + backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X)) { switch(src0t) { case GGML_TYPE_F32: { kernel = backend_ctx->kernel_mul_mm_f32_f32_l4_lm; From c589f0ed10c643678c4707dd160c21ac7633ebc0 Mon Sep 17 00:00:00 2001 From: codemonkey <441345965@qq.com> Date: Sun, 30 Aug 2026 07:44:53 +0800 Subject: [PATCH 028/109] metal : add fa-vec tunings for M2 (#27940) --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 82 +++++++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 26b634d21432..4abdafb48f30 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -467,6 +467,88 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, 1, 1 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, 1, 3 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, 2, 3 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 512, 512, 3, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 96, 96, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 128, 128, 2, 1 }, { 1, 4 } }, From 57291f2644af8c9df0dd8d44395881c5bdcf0ecd Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Sun, 30 Aug 2026 09:04:20 +0530 Subject: [PATCH 029/109] ggml: allow passing alloc dependencies in graph_optimize (#27301) * ggml: allow passing alloc dependencies in graph_optimize * add alloc dep tests * add TODO about using flat array --- ggml/src/ggml-backend-impl.h | 12 ++++- ggml/src/ggml-backend.cpp | 66 +++++++++++++++++++++--- ggml/src/ggml-cuda/ggml-cuda.cu | 4 +- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 4 +- ggml/src/ggml-metal/ggml-metal.cpp | 4 +- ggml/src/ggml-virtgpu/ggml-backend.cpp | 3 +- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 3 +- tests/test-alloc.cpp | 70 +++++++++++++++++++++++++- 8 files changed, 151 insertions(+), 15 deletions(-) diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 40cea024c3d2..56f0090cce66 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -103,6 +103,16 @@ extern "C" { // Backend (stream) // + // passed to graph_optimize so the backend can add allocation dependencies: + // if the backend executes parts of the graph out of order (e.g. on concurrent streams), + // it must keep the affected tensors allocated until a node where execution is known to have joined + struct ggml_backend_graph_optimize_params { + // keep `tensor` allocated at least until `until` (a node of the same graph) has been computed + // can be called multiple times for the same tensor: the longest lifetime applies + void (*add_alloc_dep)(void * user_data, struct ggml_tensor * tensor, struct ggml_tensor * until); + void * user_data; + }; + struct ggml_backend_i { const char * (*get_name)(ggml_backend_t backend); @@ -137,7 +147,7 @@ extern "C" { void (*event_wait) (ggml_backend_t backend, ggml_backend_event_t event); // (optional) sort/optimize the nodes in the graph - void (*graph_optimize) (ggml_backend_t backend, struct ggml_cgraph * cgraph); + void (*graph_optimize) (ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params); }; struct ggml_backend { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index e519bdf50a1b..78eb10dfe992 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #ifdef __APPLE__ @@ -558,10 +559,10 @@ void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event) backend->iface.event_wait(backend, event); } -static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph) { +static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params) { GGML_ASSERT(backend); if (backend->iface.graph_optimize != NULL) { - backend->iface.graph_optimize(backend, cgraph); + backend->iface.graph_optimize(backend, cgraph, params); } } @@ -1441,11 +1442,40 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra sched->prev_leaf_backend_ids = tmp; } + // optimize the split graphs and collect the allocation dependencies added by the backends + // this needs to happen before we make graph_copy, so they are in sync + // TODO: this may create many small allocations in the scheduler, restructure to use a flat array + std::unordered_map> alloc_deps; + + struct ggml_backend_graph_optimize_params opt_params = { + /* .add_alloc_dep = */ [](void * user_data, ggml_tensor * tensor, ggml_tensor * until) { + auto & deps = *(std::unordered_map> *) user_data; + std::vector & keep = deps[until]; + if (std::find(keep.begin(), keep.end(), tensor) == keep.end()) { + keep.push_back(tensor); + } + }, + /* .user_data = */ &alloc_deps, + }; + + for (int i = 0; i < sched->n_splits; i++) { + struct ggml_backend_sched_split * split = &sched->splits[i]; + split->graph = ggml_graph_view(graph, split->i_start, split->i_end); + + ggml_backend_graph_optimize(sched->backends[split->backend_id], &split->graph, &opt_params); + } + + // each dep is added to graph_copy as a GGML_OP_NONE node with the kept tensors as srcs + int n_dep_nodes = 0; + for (const auto & it : alloc_deps) { + n_dep_nodes += (it.second.size() + GGML_MAX_SRC - 1) / GGML_MAX_SRC; + } + int total_inputs = sched->n_graph_inputs; for (int i = 0; i < sched->n_splits; i++) { total_inputs += sched->splits[i].n_inputs; } - int graph_size = std::max(graph->n_nodes, graph->n_leafs) + total_inputs * 2 * sched->n_copies; + int graph_size = std::max(graph->n_nodes, graph->n_leafs) + total_inputs * 2 * sched->n_copies + n_dep_nodes; // remember the actual graph_size for performing reallocation checks later [GGML_SCHED_DEBUG_REALLOC] sched->debug_prev_graph_size = sched->debug_graph_size; @@ -1463,13 +1493,10 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra struct ggml_cgraph * graph_copy = &sched->graph; + int n_dep_nodes_added = 0; + for (int i = 0; i < sched->n_splits; i++) { struct ggml_backend_sched_split * split = &sched->splits[i]; - split->graph = ggml_graph_view(graph, split->i_start, split->i_end); - - // Optimize this split of the graph. This needs to happen before we make graph_copy, - // so they are in sync. - ggml_backend_graph_optimize(sched->backends[split->backend_id], &split->graph); // add inputs to the graph copy so that they are allocated by ggml-alloc at the start of the split for (int j = 0; j < split->n_inputs; j++) { @@ -1494,9 +1521,32 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra assert(graph_copy->size > graph_copy->n_nodes); sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]); graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j]; + + if (alloc_deps.empty()) { + continue; + } + + // add a dependency node so that the kept tensors are not freed before this node is computed + auto it = alloc_deps.find(graph->nodes[j]); + if (it != alloc_deps.end()) { + const std::vector & keep = it->second; + for (size_t k = 0; k < keep.size(); k += GGML_MAX_SRC) { + struct ggml_tensor * dep = ggml_view_tensor(sched->ctx, keep[k]); + for (size_t s = 0; s < GGML_MAX_SRC && k + s < keep.size(); s++) { + dep->src[s] = keep[k + s]; + } + assert(graph_copy->size > graph_copy->n_nodes); + sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id; + graph_copy->nodes[graph_copy->n_nodes++] = dep; + n_dep_nodes_added++; + } + } } } + // a mismatch means a backend added a dep with an `until` tensor that is not a node of the optimized graph + GGML_ASSERT(n_dep_nodes_added == n_dep_nodes); + if (sched->n_copies > 1) { // add input copies as leafs so that they are allocated first for (int i = 0; i < sched->n_graph_inputs; i++) { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc621..bd9754c2ffdf 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4328,7 +4328,9 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev } } -static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { +static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { + GGML_UNUSED(params); + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; #ifdef USE_CUDA_GRAPH diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 53e860755910..e7dcdc3d5513 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4984,7 +4984,9 @@ static std::vector ggml_hexagon_graph_optimize_reorder(const std::vectorn_nodes; constexpr int MAX_FUSE = 16; diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 9756d47050c3..4d58dc821cf4 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -558,7 +558,9 @@ static void ggml_backend_metal_event_wait(ggml_backend_t backend, ggml_backend_e ggml_metal_event_wait(ctx, ev); } -static void ggml_backend_metal_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { +static void ggml_backend_metal_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { + GGML_UNUSED(params); + ggml_metal_t ctx = (ggml_metal_t)backend->context; ggml_metal_graph_optimize(ctx, cgraph); diff --git a/ggml/src/ggml-virtgpu/ggml-backend.cpp b/ggml/src/ggml-virtgpu/ggml-backend.cpp index 12756c9282f7..996c57e358b6 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend.cpp @@ -17,7 +17,8 @@ static ggml_status ggml_backend_remoting_graph_compute(ggml_backend_t backend, g return apir_backend_graph_compute(gpu, cgraph); } -static void ggml_backend_remoting_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { +static void ggml_backend_remoting_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { + UNUSED(params); virtgpu * gpu = DEV_TO_GPU(backend->device); #if true UNUSED(gpu); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 39b4cd359803..8fbb1359f406 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17795,8 +17795,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } // Sort the graph for improved parallelism. -static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph) +static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, struct ggml_backend_graph_optimize_params * params) { + GGML_UNUSED(params); VK_LOG_DEBUG("ggml_vk_graph_optimize(" << graph->n_nodes << " nodes)"); ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 6d5428493e70..8f1a98aa03c3 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -19,6 +19,8 @@ struct dummy_backend_context { size_t alignment = 8; ggml_backend_buffer_i buffer_interface; + ggml_backend_device device; + ggml_backend backend; std::vector buffers; size_t allocated_total() const { @@ -83,7 +85,27 @@ static void dummy_backend_buffer_get_tensor(ggml_backend_buffer_t, const ggml_te static void dummy_backend_buffer_clear(ggml_backend_buffer_t, uint8_t) {} -// dummy_backend (not really a full backend, just provides what gallocr needs) +// ggml_backend_device interface + +static enum ggml_backend_dev_type dummy_backend_device_get_type(ggml_backend_dev_t) { + return GGML_BACKEND_DEVICE_TYPE_CPU; +} + +static bool dummy_backend_device_supports_op(ggml_backend_dev_t, const ggml_tensor *) { + return true; +} + +static bool dummy_backend_device_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { + return device->context == buft->context; +} + +// ggml_backend interface + +static const char * dummy_backend_get_name(ggml_backend_t) { + return "dummy_backend"; +} + +// dummy_backend struct dummy_backend { std::unique_ptr context; @@ -104,6 +126,16 @@ static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment b.context->buffer_interface.get_tensor = dummy_backend_buffer_get_tensor; b.context->buffer_interface.clear = dummy_backend_buffer_clear; + b.context->device.context = b.context.get(); + b.context->device.iface.get_type = dummy_backend_device_get_type; + b.context->device.iface.supports_op = dummy_backend_device_supports_op; + b.context->device.iface.supports_buft = dummy_backend_device_supports_buft; + + b.context->backend.context = b.context.get(); + b.context->backend.device = &b.context->device; + b.context->backend.iface.get_name = dummy_backend_get_name; + + b.buffer_type.device = &b.context->device; b.buffer_type.context = b.context.get(); b.buffer_type.iface.get_name = dummy_backend_buffer_type_get_name; b.buffer_type.iface.alloc_buffer = dummy_backend_buffer_type_alloc_buffer; @@ -583,6 +615,41 @@ static void test_reallocation() { } } +static void test_backend_graph_optimize(ggml_backend_t, ggml_cgraph * graph, ggml_backend_graph_optimize_params * params) { + GGML_ASSERT(graph->n_nodes == 3); + params->add_alloc_dep(params->user_data, graph->nodes[0], graph->nodes[2]); +} + +static bool graph_reuses_allocation(bool add_alloc_dep) { + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[4]; + x[0] = make_input_with_size(ctx, 16); + x[1] = ggml_scale(ctx, x[0], 2.0f); + x[2] = ggml_scale(ctx, x[1], 2.0f); + x[3] = ggml_scale(ctx, x[2], 2.0f); + + ggml_set_output(x[3]); + ggml_build_forward_expand(graph, x[3]); + + dummy_backend backend = dummy_backend_init(SIZE_MAX); + if (add_alloc_dep) { + backend.context->backend.iface.graph_optimize = test_backend_graph_optimize; + } + + ggml_backend_t backend_ptr = &backend.context->backend; + ggml_backend_buffer_type_t buft = &backend.buffer_type; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(&backend_ptr, &buft, 1, 8, false, true)); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph)); + + return x[1]->data == x[2]->data; +} + +static void test_graph_optimize_alloc_dep() { + GGML_ASSERT(graph_reuses_allocation(false)); + GGML_ASSERT(!graph_reuses_allocation(true)); +} + static void run(const char * name, void (*f)()) { printf("%s ", name); fflush(stdout); @@ -604,5 +671,6 @@ int main() { run("test_multiple_buffer_types", test_multiple_buffer_types); run("test_buffer_size_zero", test_buffer_size_zero); run("test_reallocation", test_reallocation); + run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep); return 0; } From bdf3955159d7184f44b76091973eeff532890a35 Mon Sep 17 00:00:00 2001 From: fairydreaming <166155368+fairydreaming@users.noreply.github.com> Date: Sun, 30 Aug 2026 07:47:15 +0200 Subject: [PATCH 030/109] memory : copy Hadamard matrix to k_rot tensor only if it has buffer assigned to prevent crashes during context shift of unquantized K cache (#27967) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Stanisław Szymczyk Co-authored-by: AesSedai <7980540+AesSedai@users.noreply.github.com> --- src/llama-kv-cache.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 8fafcd15304e..65afbd8c3778 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2034,7 +2034,7 @@ void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) { kv_self->set_input_k_shift(k_shift); } - if (k_rot) { + if (k_rot && k_rot->buffer) { kv_self->set_input_k_rot(k_rot); } } From d882575cc8a6fe6808fb6fedb75b372c4da3b812 Mon Sep 17 00:00:00 2001 From: QuintinShaw Date: Sun, 30 Aug 2026 13:56:35 +0800 Subject: [PATCH 031/109] metal : fix null-pipeline crash for F16 src1 mul_mat/mul_mat_id (#25648) * metal : fail closed on mul_mat shapes with missing F16 kernels * metal : abort on nil pipeline in encoder_set_pipeline * metal : address review comments * metal : share mul_mat mm dispatch with supports_op --- ggml/src/ggml-metal/ggml-metal-common.cpp | 17 +++++++++++ ggml/src/ggml-metal/ggml-metal-common.h | 4 +++ ggml/src/ggml-metal/ggml-metal-device.m | 37 ++++++++++++++++++++++- ggml/src/ggml-metal/ggml-metal-ops.cpp | 19 ++---------- 4 files changed, 59 insertions(+), 18 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp index 2eb9820bff91..6f1638a1147e 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.cpp +++ b/ggml/src/ggml-metal/ggml-metal-common.cpp @@ -1,10 +1,27 @@ #include "ggml-metal-common.h" +#include "ggml.h" #include "ggml-impl.h" #include "ggml-backend-impl.h" #include +bool ggml_metal_op_mul_mat_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm) { + const int64_t ne00 = op->src[0]->ne[0]; + const int64_t ne11 = op->src[1]->ne[1]; + + return !ggml_is_transposed(op->src[0]) && + !ggml_is_transposed(op->src[1]) && + has_simdgroup_mm && ne00 >= 64 && ne11 > 8; +} + +bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm) { + const int64_t ne00 = op->src[0]->ne[0]; + const int64_t ne21 = op->src[2]->ne[1]; + + return has_simdgroup_mm && ne00 >= 64 && ne21 >= 32; +} + // represents a memory range (i.e. an interval from a starting address p0 to an ending address p1 in a given buffer pb) // the type indicates whether it is a source range (i.e. ops read data from it) or a destination range (i.e. ops write data to it) struct ggml_mem_range { diff --git a/ggml/src/ggml-metal/ggml-metal-common.h b/ggml/src/ggml-metal/ggml-metal-common.h index 3acbc6ae174a..66abdb52efe3 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.h +++ b/ggml/src/ggml-metal/ggml-metal-common.h @@ -47,6 +47,10 @@ bool ggml_mem_ranges_check(ggml_mem_ranges_t mrs, const struct ggml_tensor * ten // if it proves to work well, we can start using it for other backends in the future void ggml_graph_optimize(struct ggml_cgraph * gf); +// mat-mat vs mat-vec dispatch; used by both supports_op and ggml_metal_op_mul_mat* +bool ggml_metal_op_mul_mat_use_mm (const struct ggml_tensor * op, bool has_simdgroup_mm); +bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm); + #ifdef __cplusplus } #endif diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 85c0f576b9fb..a053887a3370 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -3,6 +3,7 @@ #import "ggml-impl.h" #import "ggml-backend-impl.h" #import "ggml-metal-impl.h" +#import "ggml-metal-common.h" #include @@ -788,6 +789,10 @@ void ggml_metal_encoder_debug_group_pop (ggml_metal_encoder_t encoder) { } void ggml_metal_encoder_set_pipeline(ggml_metal_encoder_t encoder, struct ggml_metal_pipeline_with_params pipeline) { + if (!pipeline.pipeline) { + GGML_ABORT("%s: nil Metal pipeline (missing kernel; see compile_pipeline log above)\n", __func__); + } + [encoder->obj setComputePipelineState:pipeline.pipeline->obj]; } @@ -1410,6 +1415,30 @@ void ggml_metal_device_get_memory(ggml_metal_device_t dev, size_t * free, size_t } } +static bool ggml_metal_supports_mul_mat_op( + bool has_simdgroup_reduction, + const struct ggml_tensor * op, + bool src0_f16_has_mv, + bool mm_path) { + if (!has_simdgroup_reduction || op->src[0]->type == GGML_TYPE_NVFP4) { + return false; + } + + if (op->src[1]->type != GGML_TYPE_F16) { + return true; + } + + if (op->src[0]->type == GGML_TYPE_BF16) { + return false; + } + + if (src0_f16_has_mv && op->src[0]->type == GGML_TYPE_F16) { + return true; + } + + return mm_path; +} + bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_tensor * op) { const bool has_simdgroup_mm = dev->props.has_simdgroup_mm; const bool has_simdgroup_reduction = dev->props.has_simdgroup_reduction; @@ -1713,9 +1742,15 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_GATED_DELTA_NET: return has_simdgroup_reduction && op->src[2]->ne[0] % 32 == 0; case GGML_OP_SOLVE_TRI: + return has_simdgroup_reduction && op->src[0]->type == GGML_TYPE_F32; case GGML_OP_MUL_MAT: + return ggml_metal_supports_mul_mat_op( + has_simdgroup_reduction, op, true, + ggml_metal_op_mul_mat_use_mm(op, has_simdgroup_mm)); case GGML_OP_MUL_MAT_ID: - return has_simdgroup_reduction && op->src[0]->type != GGML_TYPE_NVFP4; + return ggml_metal_supports_mul_mat_op( + has_simdgroup_reduction, op, false, + ggml_metal_op_mul_mat_id_use_mm(op, has_simdgroup_mm)); case GGML_OP_SET: case GGML_OP_CPY: case GGML_OP_DUP: diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 89c8483b3714..7671d1d01564 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2362,10 +2362,6 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { const int16_t r2 = ne12/ne02; const int16_t r3 = ne13/ne03; - // find the break-even point where the matrix-matrix kernel becomes more efficient compared - // to the matrix-vector kernel - const int ne11_mm_min = 8; - // first try to use small-batch mat-mv kernels // these should be efficient for BS [2, ~8] if (op->src[1]->type == GGML_TYPE_F32 && (ne00%128 == 0) && @@ -2468,12 +2464,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + r0ptg - 1)/r0ptg), ((ne11 + r1ptg - 1)/r1ptg), ne12*ne13, 32, nsg, 1); - } else if ( - !ggml_is_transposed(op->src[0]) && - !ggml_is_transposed(op->src[1]) && - // for now the matrix-matrix multiplication kernel only works on A14+/M1+ SoCs - // AMD GPU and older A-chips will reuse matrix-vector multiplication kernel - props_dev->has_simdgroup_mm && ne00 >= 64 && ne11 > ne11_mm_min) { + } else if (ggml_metal_op_mul_mat_use_mm(op, props_dev->has_simdgroup_mm)) { //GGML_LOG_INFO("matrix: ne00 = %6d, ne01 = %6d, ne02 = %6d, ne11 = %6d, ne12 = %6d\n", ne00, ne01, ne02, ne11, ne12); // some Metal matrix data types require aligned pointers @@ -2622,13 +2613,7 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) { const uint32_t r2 = 1; const uint32_t r3 = 1; - // find the break-even point where the matrix-matrix kernel becomes more efficient compared - // to the matrix-vector kernel - // ne20 = n_used_experts - // ne21 = n_rows (batch size) - const int ne21_mm_id_min = 32; - - if (props_dev->has_simdgroup_mm && ne00 >= 64 && (ne21 >= ne21_mm_id_min)) { + if (ggml_metal_op_mul_mat_id_use_mm(op, props_dev->has_simdgroup_mm)) { // some Metal matrix data types require aligned pointers // ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf (Table 2.5) //switch (op->src[0]->type) { From 370cb12e8bb7d7e6a93372706b24a7356f5e6991 Mon Sep 17 00:00:00 2001 From: Titaniumtown Date: Sat, 29 Aug 2026 22:57:08 -0700 Subject: [PATCH 032/109] sycl: split long rows in TOP_K instead of one work-group per row (#27847) --- ggml/src/ggml-sycl/ggml-sycl.cpp | 294 +++++++++++++++++++++++-------- 1 file changed, 217 insertions(+), 77 deletions(-) diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index dc8a1744323e..d58ffd00dafb 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -2402,7 +2402,138 @@ static void argsort_f32_i32_sycl(const float *x, int *dst, const int ncols, } } +// Scan and block merge, shared by every launch shape below so a partitioned row uses the +// same insertion order as an unpartitioned one. +// +// src_map != nullptr: report src_map[col] instead of col, so a merge pass can carry the +// original column index through. +// out_vals != nullptr: also emit the k winning values, for a later merge pass. +// swap01: emit in the output order the single-pass path uses. +static void top_k_scan_merge_f32( + const float * src_vals, + const int32_t * src_map, + const int begin, + const int end, + const int k, + const int block_size, + float * shared_vals, + int * shared_idx, + float * out_vals, + int32_t * out_idx, + const bool swap01, + const sycl::nd_item<1> & item_ct1 +) { + const int tid = item_ct1.get_local_id(0); + + // The running top-k lives in SLM (shared local memory) rather than a private array: + // an array indexed by a runtime position cannot be register-allocated, so a private + // one lands in scratch, i.e. device memory, and insertion is this kernel's dominant + // cost. + // + // Lane-strided (lv[i * block_size]) rather than lane-blocked (lv[i]) so a given i is + // contiguous across lanes; a k-strided layout would put every lane of a shift step in + // the same SLM bank. + float * lv = shared_vals + tid; + int * li = shared_idx + tid; + + for (int i = 0; i < k; i++) { + lv[i * block_size] = -FLT_MAX; + li[i * block_size] = -1; + } + + // The k-th best, cached in a register. The reject test is taken for the large + // majority of elements scanned, and in that case touches no memory. + float kth = -FLT_MAX; + + for (int col = begin + tid; col < end; col += block_size) { + float val = src_vals[col]; + + if (val > kth) { + int pos = k - 1; + while (pos > 0 && val > lv[(pos - 1) * block_size]) { + pos--; + } + + for (int i = k - 1; i > pos; i--) { + lv[i * block_size] = lv[(i - 1) * block_size]; + li[i * block_size] = li[(i - 1) * block_size]; + } + lv[pos * block_size] = val; + li[pos * block_size] = src_map ? src_map[col] : col; + + kth = lv[(k - 1) * block_size]; + } + } + + item_ct1.barrier(sycl::access::fence_space::local_space); + + if (tid != 0) { + return; + } + + // Same treatment for the merge accumulator, past the per-lane region. + float * fv = shared_vals + (size_t) k * block_size; + int * fi = shared_idx + (size_t) k * block_size; + + for (int i = 0; i < k; i++) { + fv[i] = -FLT_MAX; + fi[i] = -1; + } + + float fkth = -FLT_MAX; + + // Candidates are visited in the same (t, i) order as before, so tie-breaking is + // unchanged. + for (int t = 0; t < block_size; t++) { + for (int i = 0; i < k; i++) { + float val = shared_vals[i * block_size + t]; + + if (val <= fkth) { + // Lane t's list is sorted descending, so once one of its entries loses + // to the k-th best, every later entry loses too. fkth only rises, so + // that stays true for the rest of the merge. This turns the merge from + // block_size*k steps into roughly block_size plus the candidates + // accepted. + break; + } + + int idx = shared_idx[i * block_size + t]; + + int pos = k - 1; + while (pos > 0 && val > fv[pos - 1]) { + pos--; + } + + for (int j = k - 1; j > pos; j--) { + fv[j] = fv[j - 1]; + fi[j] = fi[j - 1]; + } + fv[pos] = val; + fi[pos] = idx; + + fkth = fv[k - 1]; + } + } + + if (out_vals) { + for (int i = 0; i < k; i++) { + out_vals[i] = fv[i]; + } + } + + for (int i = 0; i < k; i++) { + out_idx[i] = fi[i]; + } + + if (swap01 && k > 1) { + int32_t temp = out_idx[0]; + out_idx[0] = out_idx[1]; + out_idx[1] = temp; + } +} + static void top_k_f32_sycl( + ggml_backend_sycl_context & ctx, const float * src, int32_t * dst_indices, const int64_t ncols, @@ -2410,98 +2541,107 @@ static void top_k_f32_sycl( const int k, dpct::queue_ptr main_stream ) { - const int block_size = 128; + // A row is scanned by exactly one work-group, so a vocabulary-sized row leaves the + // rest of the device idle. What the scan is short of is memory requests in flight, + // not bandwidth or per-request latency, so lanes in flight is the lever: split the + // row across independent work-groups, have each emit its partition's top-k, and + // merge those nsplit*k candidates in a second launch. + // + // split_block trades parallelism against SLM residency. Its cost is + // (split_block + 1) * k * 8 bytes of SLM per group, so at the k <= 32 ceiling 128 + // lanes need about 33 KB, which leaves a single resident group per Xe-core. Revisit + // if the supported k ever grows. + constexpr int split_block = 128; + constexpr int max_splits = 128; + constexpr int min_cols = 8192; - const sycl::range<1> block_dims(block_size); - const sycl::range<1> grid_dims(nrows); + int nsplit = 1; + if (ncols >= min_cols) { + // A partition is then always >= split_block = 128 columns, hence always more than + // the k <= 32 ceiling, so no pass is ever padded with -FLT_MAX sentinels. + const int64_t want = ncols / split_block; + nsplit = (int) (want > max_splits ? max_splits : want); + } - main_stream->submit([&](sycl::handler &cgh) { - sycl::local_accessor shared_vals(sycl::range<1>(block_size * k), cgh); - sycl::local_accessor shared_idx(sycl::range<1>(block_size * k), cgh); + if (nsplit > 1) { + const int nchunk = (int) ((ncols + nsplit - 1) / nsplit); + const size_t ncand = (size_t) nrows * nsplit * k; - cgh.parallel_for( - sycl::nd_range<1>(grid_dims * block_dims, block_dims), - [=](sycl::nd_item<1> item_ct1) { - const int row = item_ct1.get_group(0); - const int tid = item_ct1.get_local_id(0); + ggml_sycl_pool_alloc part_vals(ctx.pool(), ncand); + ggml_sycl_pool_alloc part_idx(ctx.pool(), ncand); - if (row >= nrows) return; + float * pv = part_vals.get(); + int32_t * pi = part_idx.get(); - const float * src_row = src + row * ncols; - int32_t * dst_idx_row = dst_indices + row * k; + const sycl::range<1> block_dims(split_block); - float local_vals[32]; - int local_idx[32]; + main_stream->submit([&](sycl::handler &cgh) { + sycl::local_accessor shared_vals(sycl::range<1>((split_block + 1) * k), cgh); + sycl::local_accessor shared_idx(sycl::range<1>((split_block + 1) * k), cgh); - for (int i = 0; i < k; i++) { - local_vals[i] = -FLT_MAX; - local_idx[i] = -1; - } + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>(nrows * nsplit) * block_dims, block_dims), + [=](sycl::nd_item<1> item_ct1) { + const int grp = item_ct1.get_group(0); + const int row = grp / nsplit; + const int part = grp % nsplit; + + const int begin = part * nchunk; + int end = begin + nchunk; + if (end > (int) ncols) { + end = (int) ncols; + } - for (int col = tid; col < ncols; col += block_size) { - float val = src_row[col]; + top_k_scan_merge_f32( + src + (int64_t) row * ncols, nullptr, begin, end, k, split_block, + shared_vals.get_multi_ptr().get(), + shared_idx.get_multi_ptr().get(), + pv + (size_t) grp * k, pi + (size_t) grp * k, false, item_ct1); + }); + }); - if (val > local_vals[k-1]) { - int pos = k - 1; - while (pos > 0 && val > local_vals[pos - 1]) { - pos--; - } + main_stream->submit([&](sycl::handler &cgh) { + sycl::local_accessor shared_vals(sycl::range<1>((split_block + 1) * k), cgh); + sycl::local_accessor shared_idx(sycl::range<1>((split_block + 1) * k), cgh); - for (int i = k - 1; i > pos; i--) { - local_vals[i] = local_vals[i - 1]; - local_idx[i] = local_idx[i - 1]; - } - local_vals[pos] = val; - local_idx[pos] = col; - } - } + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>(nrows) * block_dims, block_dims), + [=](sycl::nd_item<1> item_ct1) { + const int row = item_ct1.get_group(0); + const size_t off = (size_t) row * nsplit * k; + + top_k_scan_merge_f32( + pv + off, pi + off, 0, nsplit * k, k, split_block, + shared_vals.get_multi_ptr().get(), + shared_idx.get_multi_ptr().get(), + nullptr, dst_indices + (int64_t) row * k, true, item_ct1); + }); + }); - for (int i = 0; i < k; i++) { - shared_vals[tid * k + i] = local_vals[i]; - shared_idx[tid * k + i] = local_idx[i]; - } - item_ct1.barrier(sycl::access::fence_space::local_space); + return; + } - if (tid == 0) { - float final_vals[32]; - int final_idx[32]; + const int block_size = 128; - for (int i = 0; i < k; i++) { - final_vals[i] = -FLT_MAX; - final_idx[i] = -1; - } + const sycl::range<1> block_dims(block_size); + const sycl::range<1> grid_dims(nrows); - for (int t = 0; t < block_size; t++) { - for (int i = 0; i < k; i++) { - float val = shared_vals[t * k + i]; - int idx = shared_idx[t * k + i]; - - if (val > final_vals[k-1]) { - int pos = k - 1; - while (pos > 0 && val > final_vals[pos - 1]) { - pos--; - } - - for (int j = k - 1; j > pos; j--) { - final_vals[j] = final_vals[j - 1]; - final_idx[j] = final_idx[j - 1]; - } - final_vals[pos] = val; - final_idx[pos] = idx; - } - } - } + main_stream->submit([&](sycl::handler &cgh) { + sycl::local_accessor shared_vals(sycl::range<1>((block_size + 1) * k), cgh); + sycl::local_accessor shared_idx(sycl::range<1>((block_size + 1) * k), cgh); - for (int i = 0; i < k; i++) { - dst_idx_row[i] = final_idx[i]; - } + cgh.parallel_for( + sycl::nd_range<1>(grid_dims * block_dims, block_dims), + [=](sycl::nd_item<1> item_ct1) { + const int row = item_ct1.get_group(0); - if (k > 1) { - int32_t temp = dst_idx_row[0]; - dst_idx_row[0] = dst_idx_row[1]; - dst_idx_row[1] = temp; - } - } + if (row >= nrows) return; + + top_k_scan_merge_f32( + src + (int64_t) row * ncols, nullptr, 0, (int) ncols, k, block_size, + shared_vals.get_multi_ptr().get(), + shared_idx.get_multi_ptr().get(), + nullptr, dst_indices + (int64_t) row * k, true, item_ct1); }); }); } @@ -2902,7 +3042,7 @@ static void ggml_sycl_op_top_k(ggml_backend_sycl_context & ctx, ggml_tensor * ds GGML_ASSERT(k > 0 && k <= 32); GGML_ASSERT(k <= ncols); - top_k_f32_sycl(src0_dd, dst_dd, ncols, nrows, k, main_stream); + top_k_f32_sycl(ctx, src0_dd, dst_dd, ncols, nrows, k, main_stream); } inline void ggml_sycl_op_argmax(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { From 9e54e687cbc1500e2097be5b849c03bdc9de6610 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky Date: Sat, 29 Aug 2026 22:57:55 -0700 Subject: [PATCH 033/109] hexagon: support for device discovery and create sessions on demand (#27785) * hex-devices: add support for lazy session allocation and cleanup dev interfaces Co-authored-by: Marco Colombo * hex-devices: support for runtime discovery of available NPU cores Co-authored-by: Alexander Lu Co-authored-by: Ehsan Bateni * hex-devices: reject non-existing devices early during init --------- Co-authored-by: Marco Colombo Co-authored-by: Alexander Lu Co-authored-by: Ehsan Bateni --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 317 ++++++++++++++++--------- ggml/src/ggml-hexagon/htp-drv.cpp | 10 + ggml/src/ggml-hexagon/htp-drv.h | 2 + 3 files changed, 223 insertions(+), 106 deletions(-) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index e7dcdc3d5513..87e6989bd2f8 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -69,30 +69,15 @@ using u32vec = std::vector; #define GGML_HEXAGON_FENCE_SLOT_SIZE 128 struct ggml_hexagon_device_config { - int physical_idx = 0; - int virtual_idx = 0; + int physical_idx = 0; + int virtual_idx = 0; + int domain_id = 0; + std::string domain_name; std::string name; }; static ggml_hexagon_device_config opt_device_configs[GGML_HEXAGON_MAX_SESSIONS]; -static int get_domain_id(int physical_idx) { - switch (physical_idx) { - case 0: return 3; // CDSP0 (all devices) - case 1: return 4; // CDSP1 (IQ9, IQ10) - case 2: return 18; // CDSP2 (IQ10) - case 3: return 19; // CDSP3 (IQ10) - default: return CDSP_DOMAIN_ID + physical_idx; - } -} - -static std::string get_domain_name(int physical_idx) { - if (physical_idx == 0) { - return CDSP_DOMAIN_NAME; - } - return std::string("cdsp") + std::to_string(physical_idx); -} - static int opt_arch = 0; // autodetect static size_t opt_ndev = 1; static size_t opt_nhvx = 0; // use all @@ -361,7 +346,6 @@ struct ggml_hexagon_session { uint32_t session_id; uint32_t domain_id; uint64_t queue_id; - int dev_id; int phys_idx; int virt_idx; bool valid_session; @@ -376,9 +360,6 @@ struct ggml_hexagon_session { std::unordered_map> cloned_buffers; std::unordered_set sync_peers; - ggml_backend_buffer_type buffer_type = {}; - ggml_backend_buffer_type host_buffer_type = {}; - uint32_t n_threads = 0; uint32_t n_hvx = 0; uint32_t n_hmx = 0; @@ -392,12 +373,12 @@ struct ggml_hexagon_session { mutable std::unordered_set needs_repack; - ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) noexcept(false); + ggml_hexagon_session(const ggml_hexagon_device_config & config, ggml_backend_dev_t dev = nullptr) noexcept(false); ~ggml_hexagon_session() noexcept(true); const char* c_name() const { return name.c_str(); } - void allocate(int dev_id) noexcept(false); + void allocate(const ggml_hexagon_device_config & config) noexcept(false); void release() noexcept(true); void enqueue_op(const htp_opnode & node); @@ -430,14 +411,38 @@ struct ggml_hexagon_session { // ** backend buffers +struct ggml_backend_hexagon_device_context { + int dev_id; + ggml_hexagon_device_config config; + ggml_backend_dev_t dev = nullptr; + size_t max_bufsize = 0; + + ggml_backend_buffer_type buffer_type = {}; + ggml_backend_buffer_type host_buffer_type = {}; + + std::unique_ptr sess; + + ggml_backend_hexagon_device_context(int dev_id, const ggml_hexagon_device_config & config, ggml_backend_dev_t dev); + ~ggml_backend_hexagon_device_context(); + + const char * c_name() const { return config.name.c_str(); } + + ggml_hexagon_session * session() { + if (!sess) { + sess = std::make_unique(config, dev); + } + return sess.get(); + } +}; + struct ggml_backend_hexagon_buffer_type_context { - ggml_backend_hexagon_buffer_type_context(const std::string & name, ggml_hexagon_session * sess) { - this->sess = sess; - this->name = name; + ggml_backend_hexagon_buffer_type_context(const std::string & name, ggml_backend_hexagon_device_context * dev_ctx) { + this->dev_ctx = dev_ctx; + this->name = name; } - ggml_hexagon_session * sess; - std::string name; + ggml_backend_hexagon_device_context * dev_ctx; + std::string name; }; struct ggml_hexagon_rpcmem_block { @@ -576,7 +581,8 @@ struct ggml_hexagon_shared_buffer { }; static ggml_hexagon_session * ggml_backend_hexagon_buffer_get_sess(ggml_backend_buffer_t buffer) { - return static_cast(buffer->buft->context)->sess; + auto sbuf = static_cast(buffer->context); + return sbuf->sess; } static void ggml_backend_hexagon_buffer_free_buffer(ggml_backend_buffer_t buffer) { @@ -1494,24 +1500,26 @@ static const char * ggml_backend_hexagon_buffer_type_name(ggml_backend_buffer_ty static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { - auto sess = static_cast(buffer_type->context)->sess; + auto dev_ctx = static_cast(buffer_type->context)->dev_ctx; + auto sess = dev_ctx->session(); try { ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false, GGML_HEXAGON_FENCE_BUFFER_SIZE); return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_buffer_interface, sbuf, size); } catch (const std::exception & exc) { - GGML_LOG_ERROR("ggml-hex: %s failed to allocate device buffer context: %s\n", sess->c_name(), exc.what()); + GGML_LOG_ERROR("ggml-hex: %s failed to allocate device buffer context: %s\n", dev_ctx->c_name(), exc.what()); return nullptr; } } static ggml_backend_buffer_t ggml_backend_hexagon_host_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { - auto sess = static_cast(buffer_type->context)->sess; + auto dev_ctx = static_cast(buffer_type->context)->dev_ctx; + auto sess = dev_ctx->session(); try { ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false, GGML_HEXAGON_FENCE_BUFFER_SIZE); return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_host_buffer_interface, sbuf, size); } catch (const std::exception & exc) { - GGML_LOG_ERROR("ggml-hex: %s failed to allocate host buffer context: %s\n", sess->c_name(), exc.what()); + GGML_LOG_ERROR("ggml-hex: %s failed to allocate host buffer context: %s\n", dev_ctx->c_name(), exc.what()); return nullptr; } } @@ -1536,7 +1544,7 @@ static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffe static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { auto * context = static_cast(buft->context); - return context->sess->max_bufsize; + return context->dev_ctx->max_bufsize; } static bool ggml_backend_hexagon_buffer_type_is_host(ggml_backend_buffer_type_t buft) { @@ -1567,6 +1575,22 @@ static ggml_backend_buffer_type_i ggml_backend_hexagon_host_buffer_type_interfac /* .is_host = */ ggml_backend_hexagon_host_buffer_type_is_host, }; +ggml_backend_hexagon_device_context::ggml_backend_hexagon_device_context(int dev_id, const ggml_hexagon_device_config & config, ggml_backend_dev_t dev) + : dev_id(dev_id), config(config), dev(dev), max_bufsize(opt_mbuf) { + buffer_type.device = dev; + buffer_type.iface = ggml_backend_hexagon_buffer_type_interface; + buffer_type.context = new ggml_backend_hexagon_buffer_type_context(config.name, this); + + host_buffer_type.device = dev; + host_buffer_type.iface = ggml_backend_hexagon_host_buffer_type_interface; + host_buffer_type.context = new ggml_backend_hexagon_buffer_type_context(config.name + "-HOST", this); +} + +ggml_backend_hexagon_device_context::~ggml_backend_hexagon_device_context() { + delete static_cast(buffer_type.context); + delete static_cast(host_buffer_type.context); +} + static bool ggml_backend_buffer_is_hexagon(const struct ggml_backend_buffer * b) { return b->buft->iface.get_alignment == ggml_backend_hexagon_buffer_type_get_alignment; } @@ -2811,8 +2835,7 @@ static size_t ggml_hexagon_measure_max_vmem(ggml_hexagon_session *sess) { return vmem - step; // backoff to account for overhead from internal mappings } -void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { - const auto & config = opt_device_configs[dev_id]; +void ggml_hexagon_session::allocate(const ggml_hexagon_device_config & config) noexcept(false) { int phys_idx = config.physical_idx; int virt_idx = config.virtual_idx; @@ -2823,21 +2846,31 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { this->phys_idx = phys_idx; this->virt_idx = virt_idx; - this->domain_id = get_domain_id(phys_idx); + this->domain_id = config.domain_id; this->session_id = 0; - this->dev_id = dev_id; this->name = config.name; this->op_pending = 0; GGML_LOG_DEBUG("ggml-hex: %s allocating new session\n", this->name.c_str()); - domain * my_domain = htpdrv_get_domain(this->domain_id); - if (my_domain == NULL) { - GGML_LOG_ERROR("ggml-hex: unable to get domain struct for CDSP (domain_id %d)\n", this->domain_id); - throw std::runtime_error("ggml-hex: failed to get CDSP domain (see log for details)"); + if (config.domain_id < 0 || config.domain_name.empty()) { + GGML_LOG_ERROR("ggml-hex: %s: invalid physical CDSP core %d\n", config.name.c_str(), config.physical_idx); + throw std::runtime_error("ggml-hex: invalid physical CDSP core"); } - std::string dom_name = get_domain_name(phys_idx); + const std::string & dom_name = config.domain_name; + + // Enable Unsigned PD for all domains + { + struct remote_rpc_control_unsigned_module u; + u.domain = -1; + u.enable = 1; + int err = remote_session_control(DSPRPC_CONTROL_UNSIGNED_MODULE, (void *) &u, sizeof(u)); + if (err != AEE_SUCCESS) { + GGML_LOG_ERROR("ggml-hex: %s failed to enable unsigned PD : error 0x%x\n", this->c_name(), err); + throw std::runtime_error("ggml-hex: remote_session_control(unsign) failed (see log for details)"); + } + } // Create new session if virtual_idx > 0 if (virt_idx > 0) { @@ -2849,7 +2882,8 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { int err = remote_session_control(FASTRPC_RESERVE_NEW_SESSION, (void *) &n, sizeof(n)); if (err != AEE_SUCCESS) { - GGML_LOG_ERROR("ggml-hex: failed to reserve new session %d (physical %d, virtual %d) : error 0x%x\n", dev_id, phys_idx, virt_idx, err); + GGML_LOG_ERROR("ggml-hex: %s failed to reserve new session (physical %d, virtual %d) : error 0x%x\n", + this->c_name(), phys_idx, virt_idx, err); throw std::runtime_error("ggml-hex: remote_session_control(new-sess) failed (see log for details)"); } @@ -2857,10 +2891,21 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { this->session_id = n.session_id; this->domain_id = n.effective_domain_id; this->valid_session = true; + } else { + struct remote_rpc_effective_domain_id eff = {}; + eff.domain_name = const_cast(dom_name.c_str()); + eff.domain_name_len = dom_name.size(); + eff.session_id = 0; + + int err = remote_session_control(FASTRPC_GET_EFFECTIVE_DOMAIN_ID, (void *) &eff, sizeof(eff)); + if (err == AEE_SUCCESS) { + this->domain_id = eff.effective_domain_id; + } else { + GGML_LOG_DEBUG("ggml-hex: %s FASTRPC_GET_EFFECTIVE_DOMAIN_ID returned 0x%x, using domain_id %d\n", + this->name.c_str(), err, this->domain_id); + } } - // Get session URI - char session_uri[256]; { char htp_uri[256]; @@ -2877,31 +2922,18 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { int err = remote_session_control(FASTRPC_GET_URI, (void *) &u, sizeof(u)); if (err != AEE_SUCCESS) { - // fallback to single session uris - int htp_URI_domain_len = strlen(htp_uri) + MAX_DOMAIN_NAMELEN; - - snprintf(session_uri, htp_URI_domain_len, "%s%s", htp_uri, my_domain->uri); + snprintf(session_uri, sizeof(session_uri), "%s&_dom=%s&_session=%u", + htp_uri, dom_name.c_str(), this->session_id); - GGML_LOG_WARN("ggml-hex: failed to get URI for session %d (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n", dev_id, phys_idx, virt_idx, err, session_uri); - } - } - - // Enable Unsigned PD - { - struct remote_rpc_control_unsigned_module u; - u.domain = this->domain_id; - u.enable = 1; - int err = remote_session_control(DSPRPC_CONTROL_UNSIGNED_MODULE, (void *) &u, sizeof(u)); - if (err != AEE_SUCCESS) { - GGML_LOG_ERROR("ggml-hex: failed to enable unsigned PD for session %d : error 0x%x\n", dev_id, err); - throw std::runtime_error("ggml-hex: remote_session_control(unsign) failed (see log for details)"); + GGML_LOG_WARN("ggml-hex: %s failed to get URI (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n", + this->c_name(), phys_idx, virt_idx, err, session_uri); } } // Open session int err = htp_iface_open(session_uri, &this->handle); if (err != AEE_SUCCESS) { - GGML_LOG_ERROR("ggml-hex: failed to open session %d : error 0x%x\n", dev_id, err); + GGML_LOG_ERROR("ggml-hex: %s failed to open session : error 0x%x\n", this->c_name(), err); throw std::runtime_error("ggml-hex: failed to open session (see log for details)"); } @@ -2991,7 +3023,7 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { this->op_batch = new ggml_hexagon_opbatch(this, opt_opbatch, this->max_vmem); // Start dspqueue/opbatch processing - err = htp_iface_start(this->handle, dev_id, this->queue_id, opt_nhvx, opt_nhmx, this->max_vmem); + err = htp_iface_start(this->handle, this->session_id, this->queue_id, opt_nhvx, opt_nhmx, this->max_vmem); if (err != 0) { GGML_LOG_ERROR("ggml-hex: %s failed to start session: 0x%08x\n", this->c_name(), (unsigned) err); throw std::runtime_error("ggml-hex: iface start failed (see log for details)"); @@ -3054,33 +3086,23 @@ void ggml_hexagon_session::release() noexcept(true) { this->cloned_buffers.clear(); } -ggml_hexagon_session::ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) noexcept(false) { - buffer_type.device = dev; - host_buffer_type.device = dev; - +ggml_hexagon_session::ggml_hexagon_session(const ggml_hexagon_device_config & config, ggml_backend_dev_t dev) noexcept(false) { op_batch = nullptr; op_queue = nullptr; fence_seq = ((uintptr_t)this) & 0xFFFF; try { - allocate(dev_id); - - buffer_type.iface = ggml_backend_hexagon_buffer_type_interface; - buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name, this); - - host_buffer_type.iface = ggml_backend_hexagon_host_buffer_type_interface; - host_buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name + "-HOST", this); + allocate(config); } catch (const std::exception & exc) { release(); throw; } + + GGML_UNUSED(dev); } ggml_hexagon_session::~ggml_hexagon_session() noexcept(true) { release(); - - delete static_cast(buffer_type.context); - delete static_cast(host_buffer_type.context); } // ** backend interface @@ -3957,11 +3979,13 @@ static void ggml_hexagon_precompute_fused_mmnx_params( } static bool ggml_hexagon_tensor_is_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { - return t && t->buffer && t->buffer->buft == &sess->host_buffer_type; + return t && t->buffer && ggml_backend_buft_is_host(t->buffer->buft); + GGML_UNUSED(sess); } static bool ggml_hexagon_tensor_is_non_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { - return t && t->buffer && t->buffer->buft != &sess->host_buffer_type; + return t && t->buffer && !ggml_backend_buft_is_host(t->buffer->buft); + GGML_UNUSED(sess); } static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * sess, const struct ggml_tensor * dst) { @@ -5269,7 +5293,8 @@ bool ggml_backend_is_hexagon(ggml_backend_t backend) { // device interface static ggml_backend_t ggml_backend_hexagon_device_init(ggml_backend_dev_t dev, const char * params) { - auto sess = static_cast(dev->context); + auto dev_ctx = static_cast(dev->context); + auto sess = dev_ctx->session(); return new ggml_backend{ /* .guid = */ ggml_backend_hexagon_guid(), @@ -5282,8 +5307,8 @@ static ggml_backend_t ggml_backend_hexagon_device_init(ggml_backend_dev_t dev, c } static const char * ggml_backend_hexagon_device_get_name(ggml_backend_dev_t dev) { - auto sess = static_cast(dev->context); - return sess->c_name(); + auto dev_ctx = static_cast(dev->context); + return dev_ctx->c_name(); GGML_UNUSED(dev); } @@ -5321,16 +5346,16 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct } static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_buffer_type(ggml_backend_dev_t dev) { - auto sess = static_cast(dev->context); - return &sess->buffer_type; + auto dev_ctx = static_cast(dev->context); + return &dev_ctx->buffer_type; } static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_host_buffer_type(ggml_backend_dev_t dev) { if (!opt_hostbuf) { return NULL; } - auto sess = static_cast(dev->context); - return &sess->host_buffer_type; + auto dev_ctx = static_cast(dev->context); + return &dev_ctx->host_buffer_type; } static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { @@ -5421,7 +5446,8 @@ static bool ggml_hexagon_supported_fill(const struct ggml_hexagon_session * sess } static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { - auto sess = static_cast(dev->context); + auto dev_ctx = static_cast(dev->context); + auto sess = dev_ctx->session(); // reject ops that match the filter if (opt_opfilter && std::regex_match(ggml_op_desc(op), *opt_opfilter)) { @@ -5493,6 +5519,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons supp = ggml_hexagon_supported_unary(sess, op); break; default: + supp = false; break; } break; @@ -5505,6 +5532,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons supp = ggml_hexagon_supported_activations(sess, op); break; default: + supp = false; break; } break; @@ -5590,17 +5618,17 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons } static bool ggml_backend_hexagon_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - auto sess = static_cast(dev->context); + auto dev_ctx = static_cast(dev->context); // Technically we can clone hexagon buffers from any session but for some reason the output is garbled with layer-split, // tensor-split works correctly, so it needs mode debugging and investigation. For now accept only our own buffers. #if 0 bool supp = (buft->iface.get_alignment == ggml_backend_hexagon_buffer_type_get_alignment); #else - bool supp = (buft == &sess->host_buffer_type) || (buft == &sess->buffer_type); + bool supp = (buft == &dev_ctx->host_buffer_type) || (buft == &dev_ctx->buffer_type); #endif - HEX_VERBOSE("ggml-hex: %s device-supports-buft %s %s\n", sess->name.c_str(), ggml_backend_buft_name(buft), supp ? "yes" : "no"); + HEX_VERBOSE("ggml-hex: %s device-supports-buft %s %s\n", dev_ctx->c_name(), ggml_backend_buft_name(buft), supp ? "yes" : "no"); return supp; } @@ -5629,16 +5657,11 @@ ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) { GGML_LOG_INFO("ggml-hex: Hexagon Arch version v%d\n", opt_arch); - // Create devices / sessions + // Create devices for (size_t i = 0; i < opt_ndev; i++) { - devices[i].iface = ggml_backend_hexagon_device_i; - devices[i].reg = reg; - try { - devices[i].context = new ggml_hexagon_session(i, &devices[i]); - } catch (const std::exception & exc) { - GGML_LOG_ERROR("ggml-hex: failed to create device/session %zu\n", i); - devices[i].context = nullptr; - } + devices[i].iface = ggml_backend_hexagon_device_i; + devices[i].reg = reg; + devices[i].context = new ggml_backend_hexagon_device_context(i, opt_device_configs[i], &devices[i]); } } @@ -5646,10 +5669,10 @@ ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) { ggml_hexagon_registry::~ggml_hexagon_registry() { GGML_LOG_INFO("ggml-hex: releasing registry\n"); - // Release devices / sessions + // Release devices for (size_t i = 0; i < opt_ndev; i++) { - auto sess = static_cast(devices[i].context); - delete sess; + auto dev_ctx = static_cast(devices[i].context); + delete dev_ctx; } } @@ -5818,6 +5841,85 @@ template std::string vec_to_str(std::vector v) { return str; } +// Enumerate NPU (aka CDSP) domains via FASTRPC_GET_DOMAINS if supported, +// and populate domain_id and domain_name for all configured devices. +static void ggml_hexagon_discover_devices() { + std::unordered_map cdsp_map; + bool discovery_supported = false; + + system_req_payload domain_info = {}; + domain_info.id = FASTRPC_GET_DOMAINS; + domain_info.sys.domains = nullptr; + domain_info.sys.max_domains = 0; + domain_info.sys.flags = DOMAINS_LIST_FLAGS_SET_TYPE(0, FASTRPC_NSP); + + int err = remote_system_request(&domain_info); + if (err == AEE_SUCCESS && domain_info.sys.num_domains > 0) { + std::vector domains(domain_info.sys.num_domains); + domain_info.sys.domains = domains.data(); + domain_info.sys.max_domains = (int) domains.size(); + + err = remote_system_request(&domain_info); + if (err == AEE_SUCCESS) { + discovery_supported = true; + const int n_domains = std::min(domain_info.sys.num_domains, (int) domains.size()); + for (int i = 0; i < n_domains; i++) { + GGML_LOG_INFO("ggml-hex: FASTRPC_GET_DOMAINS[%d]: type %d id %d name '%s' status %d instance-id %d\n", + i, (int) domains[i].type, domains[i].id, domains[i].name, domains[i].status, domains[i].instance_id); + if (domains[i].type != FASTRPC_NSP) { + GGML_LOG_DEBUG("ggml-hex: skipping non-CDSP domain (type=%d)\n", (int) domains[i].type); + continue; + } + if (!domains[i].status) { + GGML_LOG_WARN("ggml-hex: skipping CDSP domain id=%d (status=down)\n", domains[i].id); + continue; + } + cdsp_map[domains[i].instance_id] = domains[i]; + GGML_LOG_INFO("ggml-hex: using CDSP domain: instance-id %d id %d name '%s'\n", + domains[i].instance_id, domains[i].id, domains[i].name); + } + } else { + GGML_LOG_WARN("ggml-hex: FASTRPC_GET_DOMAINS fetch failed (0x%x), using static CDSP domains\n", (unsigned) err); + } + } else if (err != AEE_SUCCESS) { + GGML_LOG_DEBUG("ggml-hex: FASTRPC_GET_DOMAINS query failed (0x%x), using static CDSP domains\n", (unsigned) err); + } + + // Populate domain IDs and names for all configured devices + for (size_t i = 0; i < opt_ndev; i++) { + auto & cfg = opt_device_configs[i]; + if (discovery_supported) { + auto it = cdsp_map.find(cfg.physical_idx); + if (it != cdsp_map.end()) { + cfg.domain_id = it->second.id; + cfg.domain_name = it->second.name; + } else { + GGML_LOG_ERROR("ggml-hex: physical CDSP core %d not found on device (%zu CDSP core(s) available)\n", + cfg.physical_idx, cdsp_map.size()); + cfg.domain_id = -1; + cfg.domain_name = ""; + } + } else { + switch (cfg.physical_idx) { + case 0: + cfg.domain_id = 3; + cfg.domain_name = CDSP_DOMAIN_NAME; + break; + case 1: + cfg.domain_id = 4; + cfg.domain_name = "cdsp1"; + break; + default: + GGML_LOG_ERROR("ggml-hex: physical CDSP core %d not supported without dynamic discovery\n", + cfg.physical_idx); + cfg.domain_id = -1; + cfg.domain_name = ""; + break; + } + } + } +} + static void ggml_hexagon_init(ggml_backend_reg * reg) { // Basic sanity checks to make sure definitions match static_assert((unsigned int) HTP_TYPE_Q4_0 == (unsigned int) GGML_TYPE_Q4_0, @@ -5983,6 +6085,9 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { } #endif + // Resolve domain info for all configured devices + ggml_hexagon_discover_devices(); + if (str_profile) { opt_pmu_evt = [&]() -> std::vector { auto v = str_to_vec(str_profile); diff --git a/ggml/src/ggml-hexagon/htp-drv.cpp b/ggml/src/ggml-hexagon/htp-drv.cpp index 4f0790801731..437e367c9d35 100644 --- a/ggml/src/ggml-hexagon/htp-drv.cpp +++ b/ggml/src/ggml-hexagon/htp-drv.cpp @@ -73,6 +73,7 @@ typedef int (*remote_handle64_close_pfn_t)(remote_handle h); typedef int (*remote_handle_control_pfn_t)(uint32_t req, void* data, uint32_t datalen); typedef int (*remote_handle64_control_pfn_t)(remote_handle64 h, uint32_t req, void* data, uint32_t datalen); typedef int (*remote_session_control_pfn_t)(uint32_t req, void *data, uint32_t datalen); +typedef int (*remote_system_request_pfn_t)(system_req_payload * req); // // Driver API pfns @@ -99,6 +100,7 @@ remote_handle64_close_pfn_t remote_handle64_close_pfn = nullptr; remote_handle_control_pfn_t remote_handle_control_pfn = nullptr; remote_handle64_control_pfn_t remote_handle64_control_pfn = nullptr; remote_session_control_pfn_t remote_session_control_pfn = nullptr; +remote_system_request_pfn_t remote_system_request_pfn = nullptr; // // Driver API @@ -206,6 +208,13 @@ HTPDRV_API int remote_session_control(uint32_t req, void * data, uint32_t datale return remote_session_control_pfn(req, data, datalen); } +HTPDRV_API int remote_system_request(system_req_payload * req) { + if (!remote_system_request_pfn) { + return AEE_EUNSUPPORTEDAPI; + } + return remote_system_request_pfn(req); +} + #ifdef _WIN32 static std::string wstr_to_str(std::wstring_view wstr) { @@ -367,6 +376,7 @@ int htpdrv_init() { dlsym(handle.get(), remote_handle64_control_pfn_t, remote_handle64_control_pfn, remote_handle64_control, false); dlsym(handle.get(), remote_session_control_pfn_t, remote_session_control_pfn, remote_session_control, false); dlsym(handle.get(), remote_handle64_close_pfn_t, remote_handle64_close_pfn, remote_handle64_close, false); + dlsym(handle.get(), remote_system_request_pfn_t, remote_system_request_pfn, remote_system_request, true); lib_cdsp_rpc_handle = std::move(handle); initialized = true; diff --git a/ggml/src/ggml-hexagon/htp-drv.h b/ggml/src/ggml-hexagon/htp-drv.h index f3cc0da75c28..8232780e7fdc 100644 --- a/ggml/src/ggml-hexagon/htp-drv.h +++ b/ggml/src/ggml-hexagon/htp-drv.h @@ -116,6 +116,8 @@ HTPDRV_API domain * htpdrv_get_domain(int domain_id); */ HTPDRV_API int htpdrv_get_arch(int domain, int * arch); +HTPDRV_API int remote_system_request(system_req_payload * req); + #ifdef __cplusplus } #endif From 2bf04151520843e9ea5694e655e7d4a537973b54 Mon Sep 17 00:00:00 2001 From: Ryan C Date: Sun, 30 Aug 2026 05:59:25 +0000 Subject: [PATCH 034/109] rpc : fix pre-rdma macOS versions (#27815) --- ggml/src/ggml-rpc/CMakeLists.txt | 6 +++++- ggml/src/ggml-rpc/transport-apple.cpp | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-rpc/CMakeLists.txt b/ggml/src/ggml-rpc/CMakeLists.txt index b2f086380d5e..af3bd0290f7c 100644 --- a/ggml/src/ggml-rpc/CMakeLists.txt +++ b/ggml/src/ggml-rpc/CMakeLists.txt @@ -34,10 +34,14 @@ if (GGML_RPC_RDMA) find_library(RDMA_LIB ${RDMA_LIB_NAME} REQUIRED) endif() target_compile_definitions(ggml-rpc PRIVATE GGML_RPC_RDMA) - target_link_libraries(ggml-rpc PRIVATE ${RDMA_LIB}) if (APPLE) + # librdma.dylib only exists on macOS 26.2 and later. Link it weakly so a build made + # where it exists still loads where it does not; checked at runtime before use. + target_link_options(ggml-rpc PRIVATE "LINKER:-weak_library,${RDMA_LIB}") target_compile_definitions(ggml-rpc PRIVATE GGML_RPC_RDMA_APPLE) target_sources(ggml-rpc PRIVATE transport-apple.cpp) + else() + target_link_libraries(ggml-rpc PRIVATE ${RDMA_LIB}) endif() message(STATUS " RDMA transport enabled (${RDMA_DESC})") else() diff --git a/ggml/src/ggml-rpc/transport-apple.cpp b/ggml/src/ggml-rpc/transport-apple.cpp index c8be77a6dcef..2cfaa5d3fd3d 100644 --- a/ggml/src/ggml-rpc/transport-apple.cpp +++ b/ggml/src/ggml-rpc/transport-apple.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -184,11 +185,28 @@ static uint8_t rdma_first_active_port(struct ibv_context * ctx, struct ibv_port_ return 0; } +// librdma.dylib is weak-linked, so its symbols are null when it is absent. Nothing may +// call one before this has returned true. +static bool rdma_library_present() { + static const bool present = [] { + void * handle = dlopen("/usr/lib/librdma.dylib", RTLD_LAZY); + if (handle == nullptr) { + return false; + } + dlclose(handle); + return true; + }(); + return present; +} + // Called before the endpoints are exchanged: pick the local device facing this // peer, create a UC QP and register the frame rings. RDMA is point-to-point, so // the device is the one whose GID equals the bootstrap connection's local // address, i.e. the one cabled to the peer. std::unique_ptr apple_rdma::probe(int fd, const uint8_t * target_gid, uint8_t * caps) { + if (!rdma_library_present()) { + return nullptr; + } int ndev = 0; ibv_device ** devs = ibv_get_device_list(&ndev); if (!devs) return nullptr; From dc7aecf70d4d80ae754c7a15d4965e8bdcb3d3fd Mon Sep 17 00:00:00 2001 From: "Alessandro de Oliveira Faria (A.K.A.CABELO)" Date: Sun, 30 Aug 2026 03:01:51 -0300 Subject: [PATCH 035/109] vendor : update cpp-httplib to 0.54.0 (#27919) * vendor : update cpp-httplib to 0.54.0 * vendor : update cpp-httplib to 0.54.0 and 0.54.1 --- scripts/sync_vendor.py | 2 +- vendor/cpp-httplib/httplib.cpp | 1133 ++++++++++++++++++++++++++------ vendor/cpp-httplib/httplib.h | 142 +++- 3 files changed, 1052 insertions(+), 225 deletions(-) diff --git a/scripts/sync_vendor.py b/scripts/sync_vendor.py index 98b9ddc8ef6b..0170b5b168f4 100755 --- a/scripts/sync_vendor.py +++ b/scripts/sync_vendor.py @@ -5,7 +5,7 @@ import sys import subprocess -HTTPLIB_VERSION = "refs/tags/v0.53.1" +HTTPLIB_VERSION = "refs/tags/v0.54.1" # used by examples/gguf-hash, these repos have no release tag, so we pin a commit XXHASH_COMMIT = "9f465f1ea932d6ad9a26cd77496311ffa544cd68" diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp index 81cdcfe3a0c8..7fd10b3939bc 100644 --- a/vendor/cpp-httplib/httplib.cpp +++ b/vendor/cpp-httplib/httplib.cpp @@ -517,7 +517,8 @@ std::string from_i_to_hex(size_t n) { return ret; } -std::string compute_etag(const FileStat &fs) { +std::string compute_etag(const FileStat &fs, + const std::string &suffix = std::string()) { if (!fs.is_file()) { return std::string(); } // If mtime cannot be determined (negative value indicates an error @@ -531,7 +532,7 @@ std::string compute_etag(const FileStat &fs) { auto size = fs.size(); return std::string("W/\"") + from_i_to_hex(mtime) + "-" + - from_i_to_hex(size) + "\""; + from_i_to_hex(size) + suffix + "\""; } // Format time_t as HTTP-date (RFC 9110 Section 5.6.7): "Sun, 06 Nov 1994 @@ -817,17 +818,14 @@ std::string websocket_accept_key(const std::string &client_key) { bool is_websocket_upgrade(const Request &req) { if (req.method != "GET") { return false; } - // Check Upgrade: websocket (case-insensitive) - auto upgrade_it = req.headers.find("Upgrade"); - if (upgrade_it == req.headers.end()) { return false; } - auto upgrade_val = case_ignore::to_lower(upgrade_it->second); - if (upgrade_val != "websocket") { return false; } + // Check Upgrade: websocket. RFC 9110 7.8 defines Upgrade as a comma-separated + // list of protocols and asks recipients to match each name + // case-insensitively, so look for the token rather than compare the whole + // field value. + if (!has_header_token(req.headers, "Upgrade", "websocket")) { return false; } - // Check Connection header contains "Upgrade" - auto connection_it = req.headers.find("Connection"); - if (connection_it == req.headers.end()) { return false; } - auto connection_val = case_ignore::to_lower(connection_it->second); - if (connection_val.find("upgrade") == std::string::npos) { return false; } + // Check Connection: Upgrade + if (!has_header_token(req.headers, "Connection", "upgrade")) { return false; } // Check Sec-WebSocket-Key is a valid base64-encoded 16-byte value (24 chars) // RFC 6455 Section 4.2.1 @@ -1221,22 +1219,14 @@ bool parse_trailers(stream_line_reader &line_reader, Headers &dest, "trailer"}; case_ignore::unordered_set declared_trailers; - auto trailer_header = get_header_value(src_headers, "Trailer", "", 0); - if (trailer_header && std::strlen(trailer_header)) { - auto len = std::strlen(trailer_header); - split(trailer_header, trailer_header + len, ',', - [&](const char *b, const char *e) { - const char *kbeg = b; - const char *kend = e; - while (kbeg < kend && (*kbeg == ' ' || *kbeg == '\t')) { - ++kbeg; - } - while (kend > kbeg && (kend[-1] == ' ' || kend[-1] == '\t')) { - --kend; - } - std::string key(kbeg, static_cast(kend - kbeg)); - if (!key.empty() && - prohibited_trailers.find(key) == prohibited_trailers.end()) { + auto trailer_header = get_combined_header_value(src_headers, "Trailer"); + if (!trailer_header.empty()) { + // split() trims each token and skips empty ones, so the name arrives ready + // to look up. + split(trailer_header.data(), trailer_header.data() + trailer_header.size(), + ',', [&](const char *b, const char *e) { + std::string key(b, e); + if (prohibited_trailers.find(key) == prohibited_trailers.end()) { declared_trailers.insert(key); } }); @@ -1340,6 +1330,55 @@ void split(const char *b, const char *e, char d, size_t m, } } +// Same contract as split(), except that a delimiter inside a quoted-string is +// not a delimiter. RFC 9110 Section 5.6.6 lets a parameter value be a +// quoted-string, and ';' and '=' are legal characters inside one. +void split_unquoted(const char *b, const char *e, char d, size_t m, + std::function fn) { + size_t i = 0; + size_t beg = 0; + size_t count = 1; + auto in_quotes = false; + + while (e ? (b + i < e) : (b[i] != '\0')) { + if (b[i] == '"') { + in_quotes = !in_quotes; + } else if (b[i] == d && !in_quotes && count < m) { + auto r = trim(b, e, beg, i); + if (r.first < r.second) { fn(&b[r.first], &b[r.second]); } + beg = i + 1; + count++; + } + i++; + } + + if (i) { + auto r = trim(b, e, beg, i); + if (r.first < r.second) { fn(&b[r.first], &b[r.second]); } + } +} + +void split_unquoted(const char *b, const char *e, char d, + std::function fn) { + return split_unquoted(b, e, d, (std::numeric_limits::max)(), + std::move(fn)); +} + +// Divide a header parameter at its first '='. RFC 9110 Section 5.6.6 makes the +// key a token, so the first '=' is the separator even when the value is a +// quoted-string carrying more of them. +void divide_param_pair(const char *b, const char *e, std::string &key, + std::string &val) { + divide( + b, static_cast(e - b), '=', + [&](const char *kb, std::size_t klen, const char *vb, std::size_t vlen) { + const auto kr = trim(kb, kb + klen, 0, klen); + key.assign(kb + kr.first, kb + kr.second); + const auto vr = trim(vb, vb + vlen, 0, vlen); + val.assign(vb + vr.first, vb + vr.second); + }); +} + bool split_find(const char *b, const char *e, char d, size_t m, std::function fn) { size_t i = 0; @@ -2424,6 +2463,36 @@ bool is_connection_error() { #endif } +// accept() failed because the process or the network stack is temporarily out +// of resources. The listening socket is still usable, so back off briefly and +// try again. +bool is_accept_resource_error() { +#ifdef _WIN32 + auto err = WSAGetLastError(); + return err == WSAEMFILE || err == WSAENOBUFS; +#else + auto err = errno; + return err == EMFILE || err == ENFILE || err == ENOBUFS || err == ENOMEM; +#endif +} + +// accept() failed for a reason that says nothing about the listening socket: +// the pending connection went away before it could be accepted, or the call +// was interrupted. Retry immediately. WSAAccept()'s own documentation omits +// WSAECONNRESET, but the accept() it wraps reports an aborted pending +// connection that way. +bool is_accept_transient_error() { +#ifdef _WIN32 + auto err = WSAGetLastError(); + return err == WSAEINTR || err == WSAEWOULDBLOCK || err == WSAECONNRESET || + err == WSAECONNABORTED; +#else + auto err = errno; + return err == EINTR || err == EAGAIN || err == EWOULDBLOCK || + err == ECONNABORTED; +#endif +} + bool bind_ip_address(socket_t sock, const std::string &host) { struct addrinfo hints; struct addrinfo *result; @@ -2728,21 +2797,16 @@ extract_media_type(const std::string &content_type, if (params) { // Parse parameters: key=value pairs separated by ';' - split(param_str.data(), param_str.data() + param_str.size(), ';', - [&](const char *b, const char *e) { - std::string key; - std::string val; - split(b, e, '=', [&](const char *b2, const char *e2) { - if (key.empty()) { - key.assign(b2, e2); - } else { - val.assign(b2, e2); - } - }); - if (!key.empty()) { - params->emplace(trim_copy(key), trim_double_quotes_copy(val)); - } - }); + split_unquoted(param_str.data(), param_str.data() + param_str.size(), ';', + [&](const char *b, const char *e) { + std::string key; + std::string val; + divide_param_pair(b, e, key, val); + if (!key.empty()) { + params->emplace(trim_copy(key), + trim_double_quotes_copy(val)); + } + }); } } @@ -2828,12 +2892,11 @@ bool parse_quality(const char *b, const char *e, std::string &token, return !invalid; } -EncodingType encoding_type(const Request &req, const Response &res) { - if (!can_compress_content_type(res.get_header_value("Content-Type"))) { - return EncodingType::None; - } +EncodingType encoding_type(const Request &req, + const std::string &content_type) { + if (!can_compress_content_type(content_type)) { return EncodingType::None; } - const auto &s = req.get_header_value("Accept-Encoding"); + auto s = get_combined_header_value(req.headers, "Accept-Encoding"); if (s.empty()) { return EncodingType::None; } // Single-pass: iterate tokens and track the best supported encoding. @@ -2885,6 +2948,10 @@ EncodingType encoding_type(const Request &req, const Response &res) { return best; } +EncodingType encoding_type(const Request &req, const Response &res) { + return encoding_type(req, res.get_header_value("Content-Type")); +} + std::unique_ptr make_compressor(EncodingType type) { #ifdef CPPHTTPLIB_ZLIB_SUPPORT if (type == EncodingType::Gzip) { @@ -3174,13 +3241,6 @@ bool zstd_decompressor::decompress(const char *data, size_t data_length, } #endif -bool contains_case_ignore(const std::string &s, const char *token) { - auto token_end = token + std::strlen(token); - return std::search(s.begin(), s.end(), token, token_end, [](char a, char b) { - return case_ignore::to_lower(a) == case_ignore::to_lower(b); - }) != s.end(); -} - // Content codings are case-insensitive (RFC 9110 8.4.1). Matching them // case-sensitively would make a response labeled e.g. "GZIP" look like an // unknown coding, and its payload would be handed back still compressed. @@ -3190,11 +3250,11 @@ bool is_zlib_encoding(const std::string &encoding) { } bool is_brotli_encoding(const std::string &encoding) { - return contains_case_ignore(encoding, "br"); + return case_ignore::equal(encoding, "br"); } bool is_zstd_encoding(const std::string &encoding) { - return contains_case_ignore(encoding, "zstd"); + return case_ignore::equal(encoding, "zstd"); } // Returns true if the content coding is one cpp-httplib is able to decompress @@ -3281,6 +3341,45 @@ size_t get_header_value_count(const Headers &headers, return headers.count(key); } +// RFC 9110 Section 5.2 and 5.3: a field that is defined as a comma-separated +// list may be sent as several field lines, and the combined field value is +// those values joined by commas in the order they were received. Callers that +// parse such a list must work on the combined value; reading only the first +// occurrence silently drops whatever the later field lines carry. +std::string get_combined_header_value(const Headers &headers, + const std::string &key) { + std::string combined; + auto rng = headers.equal_range(key); + for (auto it = rng.first; it != rng.second; ++it) { + // RFC 9110 Section 5.6.1.2: a recipient has to parse and ignore empty list + // elements, so an empty field line must not contribute a bare comma to the + // combined value. + if (it->second.empty()) { continue; } + if (!combined.empty()) { combined += ", "; } + combined += it->second; + } + return combined; +} + +bool has_header_token(const Headers &headers, const std::string &key, + const std::string &token) { + // RFC 9110 7.6.1: a comma-separated token list field such as Connection may + // carry several tokens, and RFC 9110 5.3 lets that list be split across + // several lines. Match complete tokens rather than searching the raw value, + // so that a value such as "notupgrade" is not read as the token "upgrade". + auto rng = headers.equal_range(key); + for (auto it = rng.first; it != rng.second; ++it) { + const auto &value = it->second; + if (split_find(value.data(), value.data() + value.size(), ',', + [&](const char *b, const char *e) { + return case_ignore::equal(std::string(b, e), token); + })) { + return true; + } + } + return false; +} + template typename Map::mapped_type get_multimap_value(const Map &m, const std::string &key, size_t id) { @@ -3411,19 +3510,14 @@ bool read_websocket_upgrade_response(Stream &strm, return false; } - // Verify Upgrade: websocket (case-insensitive) - auto upgrade_it = headers.find("Upgrade"); - if (upgrade_it == headers.end() || - case_ignore::to_lower(upgrade_it->second) != "websocket") { + // Verify Upgrade: websocket (a comma-separated list, matched per token) + if (!has_header_token(headers, "Upgrade", "websocket")) { upgrade.error = Error::WebSocketHandshake; return false; } - // Verify Connection header contains "Upgrade" (case-insensitive) - auto connection_it = headers.find("Connection"); - if (connection_it == headers.end() || - case_ignore::to_lower(connection_it->second).find("upgrade") == - std::string::npos) { + // Verify Connection: Upgrade + if (!has_header_token(headers, "Connection", "upgrade")) { upgrade.error = Error::WebSocketHandshake; return false; } @@ -3589,7 +3683,7 @@ bool prepare_content_receiver(T &x, int &status, bool decompress, size_t payload_max_length, bool &exceed_payload_max_length, U callback) { if (decompress) { - std::string encoding = x.get_header_value("Content-Encoding"); + auto encoding = get_combined_header_value(x.headers, "Content-Encoding"); std::unique_ptr decompressor; if (!encoding.empty()) { @@ -3769,6 +3863,7 @@ bool write_content_with_progress(Stream &strm, size_t end_offset = offset + length; size_t start_offset = offset; auto ok = true; + auto finished = false; DataSink data_sink; data_sink.write = [&](const char *d, size_t l) -> bool { @@ -3792,7 +3887,14 @@ bool write_content_with_progress(Stream &strm, data_sink.is_writable = [&]() -> bool { return strm.is_peer_alive(); }; - while (offset < end_offset && !is_shutting_down()) { + // The body is framed by `length`, so a provider that reports itself done + // early has truncated it. Record that and let the short-body check below + // fail the write, rather than calling the provider again forever. + data_sink.done = [&]() { finished = true; }; + + while (offset < end_offset && !finished && !is_shutting_down()) { + auto last_offset = offset; + if (!strm.wait_writable() || !strm.is_peer_alive()) { error = Error::Write; return false; @@ -3803,9 +3905,18 @@ bool write_content_with_progress(Stream &strm, error = Error::Write; return false; } + + // A provider that reports success without writing anything and without + // reporting itself done gets handed the same offset and length again on + // the next pass, so it would spin here for as long as the peer stays + // connected. Treat making no progress as a short body, like done() early. + if (!finished && offset == last_offset) { + error = Error::Write; + return false; + } } - if (offset < end_offset) { // exited due to is_shutting_down(), not completion + if (offset < end_offset) { // done() called early, or is_shutting_down() error = Error::Write; return false; } @@ -3866,6 +3977,67 @@ write_content_without_length(Stream &strm, // down } +// Runs a known-length content provider to completion and compresses what it +// writes into `out`. Nothing is buffered in identity form: a provider backed +// by an mmap hands the compressor a pointer straight into the mapping. +bool compress_content_provider(const ContentProvider &content_provider, + size_t length, compressor &cmp, + std::string &out) { + size_t offset = 0; + auto ok = true; + auto finished = false; + DataSink data_sink; + + auto append = [&](const char *data, size_t data_len) { + out.append(data, data_len); + return true; + }; + + data_sink.write = [&](const char *d, size_t l) -> bool { + if (!ok) { return false; } + offset += l; + if (l > 0 && !cmp.compress(d, l, false, append)) { ok = false; } + return ok; + }; + + // The body is framed by `length`, so a provider that reports itself done + // early has truncated it; the short-body check below turns that into a + // failure rather than calling the provider again forever. + data_sink.done = [&]() { finished = true; }; + + while (offset < length && !finished) { + auto prev_offset = offset; + if (!content_provider(offset, length - offset, data_sink) || !ok) { + return false; + } + // No Stream to block on here, so a provider that keeps returning true + // without writing would spin. Treat a pass that made no progress as a + // failure. + if (offset == prev_offset) { return false; } + } + + if (offset != length) { return false; } + + return cmp.compress(nullptr, 0, true, append); +} + +// Serves `m` as the response body. `set_content_provider()` clears the coding, +// so recording it has to come after; keeping both here means a third +// file-serving path cannot get that order wrong. +void set_file_content_provider(Response &res, + const std::shared_ptr &m, + const std::string &content_type, + EncodingType encoding) { + res.set_content_provider( + m->size(), content_type, + [m](size_t offset, size_t length, DataSink &sink) -> bool { + sink.write(m->data() + offset, length); + return true; + }); + + res.file_content_encoding_ = encoding; +} + template bool write_content_chunked(Stream &strm, const ContentProvider &content_provider, @@ -3876,8 +4048,10 @@ write_content_chunked(Stream &strm, const ContentProvider &content_provider, DataSink data_sink; data_sink.write = [&](const char *d, size_t l) -> bool { - if (ok) { - data_available = l > 0; + // Only done()/done_with_trailer() end a chunked body. A pass with nothing + // to hand over is ordinary (an empty buffer popped off a queue), and a + // zero-length chunk is the terminator, so it must not be emitted here. + if (ok && l > 0) { offset += l; std::string payload; @@ -4130,31 +4304,30 @@ bool parse_multipart_boundary(const std::string &content_type, auto it = params.find("boundary"); if (it == params.end()) { return false; } boundary = it->second; - return !boundary.empty(); + // RFC 2046 5.1.1 caps a boundary at 70 characters. The parser scans the body + // for "--" + boundary, so a body crafted to repeat that delimiter's leading + // bytes costs a nearly full comparison at nearly every position: the + // boundary's length multiplies the worst-case cost of scanning a body. + return !boundary.empty() && boundary.size() <= 70; } void parse_disposition_params(const std::string &s, Params ¶ms) { std::set cache; - split(s.data(), s.data() + s.size(), ';', [&](const char *b, const char *e) { - std::string kv(b, e); - if (cache.find(kv) != cache.end()) { return; } - cache.insert(kv); + split_unquoted(s.data(), s.data() + s.size(), ';', + [&](const char *b, const char *e) { + std::string kv(b, e); + if (cache.find(kv) != cache.end()) { return; } + cache.insert(kv); - std::string key; - std::string val; - split(b, e, '=', [&](const char *b2, const char *e2) { - if (key.empty()) { - key.assign(b2, e2); - } else { - val.assign(b2, e2); - } - }); + std::string key; + std::string val; + divide_param_pair(b, e, key, val); - if (!key.empty()) { - params.emplace(trim_double_quotes_copy((key)), - trim_double_quotes_copy((val))); - } - }); + if (!key.empty()) { + params.emplace(trim_double_quotes_copy(key), + trim_double_quotes_copy(val)); + } + }); } #ifdef CPPHTTPLIB_NO_EXCEPTIONS @@ -4224,12 +4397,6 @@ bool parse_accept_header(const std::string &s, // Empty string is considered valid (no preference) if (s.empty()) { return true; } - // Check for invalid patterns: leading/trailing commas or consecutive commas - if (s.front() == ',' || s.back() == ',' || - s.find(",,") != std::string::npos) { - return false; - } - struct AcceptEntry { std::string media_type; double quality; @@ -4240,16 +4407,16 @@ bool parse_accept_header(const std::string &s, int order = 0; bool has_invalid_entry = false; - // Split by comma and parse each entry + // Split by comma and parse each entry. RFC 9110 Section 5.6.1.2: a recipient + // has to parse and ignore empty list elements, so a leading, trailing or + // doubled comma must not turn a legal Accept value into 400 Bad Request. + // split() skips them, and the header length limit bounds how many a sender + // can send, so ignoring all of them cannot be used as a denial-of-service + // vector. split(s.data(), s.data() + s.size(), ',', [&](const char *b, const char *e) { std::string entry(b, e); entry = trim_copy(entry); - if (entry.empty()) { - has_invalid_entry = true; - return; - } - AcceptEntry accept_entry; accept_entry.order = order++; @@ -4314,13 +4481,25 @@ class FormDataParser { bool parse(const char *buf, size_t n, const FormDataHeader &header_callback, const ContentReceiver &content_callback) { + // Once the close delimiter has been seen the rest of the body is epilogue + // to be discarded (RFC 2046). Drop it without buffering so a large epilogue + // spread across reads is not copied in only to be erased right away. + if (state_ == 5) { return true; } + buf_append(buf, n); while (buf_size() > 0) { switch (state_) { case 0: { // Initial boundary auto pos = buf_find(dash_boundary_crlf_); - if (pos == buf_size()) { return true; } + if (pos == buf_size()) { + // Not found yet: keep only a possible partial boundary at the tail so + // that a body which never contains the boundary cannot grow the + // buffer (and get rescanned from the start) without bound. + auto keep = dash_boundary_crlf_.size() - 1; + if (buf_size() > keep) { buf_erase(buf_size() - keep); } + return true; + } buf_erase(pos + dash_boundary_crlf_.size()); state_ = 1; break; @@ -4443,18 +4622,26 @@ class FormDataParser { if (buf_start_with(crlf_)) { buf_erase(crlf_.size()); state_ = 1; + } else if (buf_start_with(dash_)) { + buf_erase(dash_.size()); + is_valid_ = true; + state_ = 5; } else { - if (dash_.size() > buf_size()) { return true; } - if (buf_start_with(dash_)) { - buf_erase(dash_.size()); - is_valid_ = true; - buf_erase(buf_size()); // Remove epilogue - } else { - return true; - } + // Only CRLF (another part follows) and "--" (close-delimiter) are + // accepted after a boundary; RFC 2046 allows transport-padding in + // between, but this parser has never supported it. Either way the + // body is already destined to be rejected, so fail now instead of + // buffering the rest of it. Both are two bytes, so the check above + // already guarantees enough buffered data to decide. + is_valid_ = false; + return false; } break; } + case 5: { // Epilogue + buf_erase(buf_size()); + break; + } } } @@ -5050,9 +5237,11 @@ bool has_framed_body(const Request &req) { } bool is_connection_persistent(const Request &req) { - auto conn = req.get_header_value("Connection"); - if (conn == "close") { return false; } - if (req.version == "HTTP/1.0" && conn != "Keep-Alive") { return false; } + if (has_header_token(req.headers, "Connection", "close")) { return false; } + if (req.version == "HTTP/1.0" && + !has_header_token(req.headers, "Connection", "keep-alive")) { + return false; + } return true; } @@ -5082,38 +5271,108 @@ class WSInit { static WSInit wsinit_; #endif +// RFC 9110 Section 11.6.1 defines a challenge list as +// WWW-Authenticate = #challenge +// challenge = auth-scheme [ 1*SP ( token68 / [ #auth-param ] ) ] +// auth-param = token BWS "=" BWS ( token / quoted-string ) +// so a server may offer several schemes, each with its own comma-separated +// auth-param list, in either order and either as separate field lines or +// packed into one. Splitting on every comma would break apart a challenge's +// own param list; splitting only on the first space would miss a Digest +// challenge that isn't first. Split on commas that aren't inside a +// quoted-string instead, then track which scheme each resulting segment +// belongs to: a segment whose text before "=" contains whitespace (or that +// has no "=" at all) starts a new challenge named by its leading token. +std::vector split_challenge_segments(const std::string &s) { + std::vector segments; + size_t start = 0; + auto in_quotes = false; + for (size_t i = 0; i < s.size(); i++) { + auto c = s[i]; + if (in_quotes) { + if (c == '\\' && i + 1 < s.size()) { + i++; + } else if (c == '"') { + in_quotes = false; + } + } else if (c == '"') { + in_quotes = true; + } else if (c == ',') { + segments.push_back(s.substr(start, i - start)); + start = i + 1; + } + } + segments.push_back(s.substr(start)); + return segments; +} + +std::string unescape_quoted_pairs(const std::string &s) { + std::string out; + out.reserve(s.size()); + for (size_t i = 0; i < s.size(); i++) { + if (s[i] == '\\' && i + 1 < s.size()) { + out += s[++i]; + } else { + out += s[i]; + } + } + return out; +} + bool parse_www_authenticate(const Response &res, std::map &auth, bool is_proxy) { auto auth_key = is_proxy ? "Proxy-Authenticate" : "WWW-Authenticate"; - if (res.has_header(auth_key)) { - thread_local auto re = - std::regex(R"~((?:(?:,\s*)?(.+?)=(?:"(.*?)"|([^,]*))))~"); - auto s = res.get_header_value(auth_key); - auto pos = s.find(' '); - if (pos != std::string::npos) { - auto type = s.substr(0, pos); - if (type == "Basic") { - return false; - } else if (type == "Digest") { - s = s.substr(pos + 1); - auto beg = std::sregex_iterator(s.begin(), s.end(), re); - for (auto i = beg; i != std::sregex_iterator(); ++i) { - const auto &m = *i; - auto key = s.substr(static_cast(m.position(1)), - static_cast(m.length(1))); - auto val = m.length(2) > 0 - ? s.substr(static_cast(m.position(2)), - static_cast(m.length(2))) - : s.substr(static_cast(m.position(3)), - static_cast(m.length(3))); - auth[std::move(key)] = std::move(val); - } - return true; + auto combined = get_combined_header_value(res.headers, auth_key); + if (combined.empty()) { return false; } + + auto found_digest = false; + auto in_digest_challenge = false; + for (const auto &raw_segment : split_challenge_segments(combined)) { + auto segment = trim_copy(raw_segment); + if (segment.empty()) { continue; } + + auto eq_pos = segment.find('='); + // BWS is allowed on both sides of "=", so the text naming the key (or, + // for the first segment of a challenge, " ") must be + // trimmed before its boundaries are inspected. + auto key_part = trim_copy( + eq_pos == std::string::npos ? segment : segment.substr(0, eq_pos)); + auto space_pos = key_part.find_last_of(" \t"); + if (space_pos != std::string::npos || eq_pos == std::string::npos) { + // "[ ]" starts a new challenge. + auto scheme_end = + space_pos == std::string::npos ? key_part.size() : space_pos; + // RFC 7616 Section 3.7: a server may offer more than one Digest + // challenge (e.g. SHA-256 and MD5); keep only the first so a nonce + // from one challenge is never paired with another's algorithm. + in_digest_challenge = + !found_digest && + case_ignore::equal(key_part.substr(0, scheme_end), "Digest"); + if (in_digest_challenge) { found_digest = true; } + if (space_pos == std::string::npos) { + // Bare scheme (or a token68), no auth-param on this segment. + continue; } + key_part = key_part.substr(space_pos + 1); } + + if (!in_digest_challenge) { continue; } + + auto val = trim_copy(segment.substr(eq_pos + 1)); + auto unquoted = trim_double_quotes_copy(val); + if (unquoted.size() != val.size()) { + unquoted = unescape_quoted_pairs(unquoted); + } + auth[std::move(key_part)] = std::move(unquoted); } - return false; + + // RFC 7616 Section 3.3 requires realm and nonce on every Digest challenge; + // make_digest_authentication_header() dereferences both unconditionally, so + // a challenge missing either can't produce a usable Authorization header. + // Treat it the same as no Digest challenge at all. + return found_digest && auth.find("realm") != auth.end() && + auth.find("nonce") != auth.end(); } class ContentProviderAdapter { @@ -5317,6 +5576,52 @@ class SSLSocketStream final : public Stream { bool readable_hint_ = false; }; +// A TLS stream for WebSocket connections, where the receive path and the +// send path (application send() plus the heartbeat ping thread) run on +// different threads. A single TLS session must never be entered +// concurrently, so every call into the session is serialized by one mutex. +// +// Unlike SSLSocketStream, the socket is kept non-blocking for the stream's +// whole lifetime and each read()/write() performs a single non-blocking TLS +// call under the lock, then waits for readiness with select() outside the +// lock. The lock is therefore held only for CPU-bound work, so a reader +// blocked waiting for data never stalls a concurrent sender. +// +// This stream is used only for wss:// connections. Plain ws:// and ordinary +// HTTP/HTTPS keep using SocketStream/SSLSocketStream unchanged. +class WebSocketSSLStream final : public Stream { +public: + WebSocketSSLStream(socket_t sock, tls::session_t session, + time_t read_timeout_sec, time_t read_timeout_usec, + time_t write_timeout_sec, time_t write_timeout_usec); + ~WebSocketSSLStream() override; + + bool is_readable() const override; + bool wait_readable() const override; + bool wait_writable() const override; + ssize_t read(char *ptr, size_t size) override; + ssize_t write(const char *ptr, size_t size) override; + void get_remote_ip_and_port(std::string &ip, int &port) const override; + void get_local_ip_and_port(std::string &ip, int &port) const override; + socket_t socket() const override; + time_t duration() const override; + void set_read_timeout(time_t sec, time_t usec = 0) override; + +private: + mutable std::mutex session_mutex_; + + socket_t sock_; + tls::session_t session_; + // WebSocket::close() shortens the read timeout from the closing thread + // while the receive thread is inside wait_readable(), so these two are read + // and written concurrently. The write timeouts are never mutated. + std::atomic read_timeout_sec_; + std::atomic read_timeout_usec_; + time_t write_timeout_sec_; + time_t write_timeout_usec_; + const std::chrono::time_point start_time_; +}; + #ifdef CPPHTTPLIB_OPENSSL_SUPPORT std::string message_digest(const std::string &s, const EVP_MD *algo) { auto context = std::unique_ptr( @@ -5893,10 +6198,15 @@ bool set_socket_opt(socket_t sock, int level, int optname, int optval) { } std::string get_bearer_token_auth(const Request &req) { - if (req.has_header("Authorization")) { - constexpr auto bearer_header_prefix_len = detail::str_len("Bearer "); - return req.get_header_value("Authorization") - .substr(bearer_header_prefix_len); + // The auth scheme is case-insensitive (RFC 9110 11.1), and a value shorter + // than the prefix carries no token. + constexpr const char bearer_prefix[] = "Bearer "; + constexpr auto bearer_prefix_len = detail::str_len(bearer_prefix); + auto value = req.get_header_value("Authorization"); + if (value.size() >= bearer_prefix_len && + detail::case_ignore::equal(value.substr(0, bearer_prefix_len), + bearer_prefix)) { + return value.substr(bearer_prefix_len); } return ""; } @@ -6018,6 +6328,7 @@ std::string to_string(const Error error) { case Error::InvalidRangeHeader: return "Invalid Range header"; case Error::UnsupportedContentEncoding: return "Unsupported Content-Encoding"; case Error::WebSocketHandshake: return "WebSocket handshake failed"; + case Error::UserCallbackException: return "User callback threw an exception"; default: break; } @@ -6137,7 +6448,19 @@ std::string decode_uri(const std::string &value) { if (value[i] == '%' && i + 2 < value.size()) { auto val = 0; if (detail::from_hex_to_i(value, i + 1, 2, val)) { - result += static_cast(val); + auto c = static_cast(val); + // Keep escapes of the reserved characters that encode_uri leaves + // literal, so decode_uri is the inverse of encode_uri and an escaped + // delimiter is not promoted into a real one (as with JS decodeURI). + if (c == ';' || c == '/' || c == '?' || c == ':' || c == '@' || + c == '&' || c == '=' || c == '+' || c == '$' || c == ',' || + c == '#') { + result += value[i]; + result += value[i + 1]; + result += value[i + 2]; + } else { + result += c; + } i += 2; } else { result += value[i]; @@ -6579,6 +6902,7 @@ void Response::set_content(const char *s, size_t n, auto rng = headers.equal_range("Content-Type"); headers.erase(rng.first, rng.second); set_header("Content-Type", content_type); + file_content_encoding_ = detail::EncodingType::None; } void Response::set_content(const std::string &s, @@ -6593,6 +6917,7 @@ void Response::set_content(std::string &&s, auto rng = headers.equal_range("Content-Type"); headers.erase(rng.first, rng.second); set_header("Content-Type", content_type); + file_content_encoding_ = detail::EncodingType::None; } void Response::set_content_provider( @@ -6603,6 +6928,7 @@ void Response::set_content_provider( if (in_length > 0) { content_provider_ = std::move(provider); } content_provider_resource_releaser_ = std::move(resource_releaser); is_chunked_content_provider_ = false; + file_content_encoding_ = detail::EncodingType::None; } void Response::set_content_provider( @@ -6613,6 +6939,7 @@ void Response::set_content_provider( content_provider_ = detail::ContentProviderAdapter(std::move(provider)); content_provider_resource_releaser_ = std::move(resource_releaser); is_chunked_content_provider_ = false; + file_content_encoding_ = detail::EncodingType::None; } void Response::set_chunked_content_provider( @@ -6623,6 +6950,7 @@ void Response::set_chunked_content_provider( content_provider_ = detail::ContentProviderAdapter(std::move(provider)); content_provider_resource_releaser_ = std::move(resource_releaser); is_chunked_content_provider_ = true; + file_content_encoding_ = detail::EncodingType::None; } void Response::set_file_content(const std::string &path, @@ -7600,6 +7928,127 @@ void SSLSocketStream::set_read_timeout(time_t sec, time_t usec) { read_timeout_usec_ = usec; } +WebSocketSSLStream::WebSocketSSLStream(socket_t sock, + tls::session_t session, + time_t read_timeout_sec, + time_t read_timeout_usec, + time_t write_timeout_sec, + time_t write_timeout_usec) + : sock_(sock), session_(session), read_timeout_sec_(read_timeout_sec), + read_timeout_usec_(read_timeout_usec), + write_timeout_sec_(write_timeout_sec), + write_timeout_usec_(write_timeout_usec), + start_time_(std::chrono::steady_clock::now()) { + // The receive and send paths run on different threads, so each TLS call is + // driven in non-blocking mode and readiness is awaited with select() + // outside the session lock. Set the socket non-blocking once here; it is + // never flipped back, so no thread races on the flag. + detail::set_nonblocking(sock_, true); +#ifdef CPPHTTPLIB_OPENSSL_SUPPORT + SSL_clear_mode(static_cast(session_), SSL_MODE_AUTO_RETRY); +#endif +} + +WebSocketSSLStream::~WebSocketSSLStream() = default; + +bool WebSocketSSLStream::is_readable() const { + std::lock_guard guard(session_mutex_); + return tls::pending(session_) > 0; +} + +bool WebSocketSSLStream::wait_readable() const { + return select_read(sock_, read_timeout_sec_, read_timeout_usec_) > 0; +} + +bool WebSocketSSLStream::wait_writable() const { + // Unlike SSLSocketStream, this deliberately does not call is_peer_closed(): + // that probe toggles the socket's blocking flag, which would race with the + // concurrent reader on a permanently non-blocking socket. + return select_write(sock_, write_timeout_sec_, write_timeout_usec_) > 0; +} + +ssize_t WebSocketSSLStream::read(char *ptr, size_t size) { + tls::TlsError err; + auto n = 1000; + while (--n >= 0) { + { + std::lock_guard guard(session_mutex_); + auto ret = tls::read(session_, ptr, size, err); + if (ret > 0) { return ret; } + if (ret == 0 || err.code == tls::ErrorCode::PeerClosed) { + error_ = Error::ConnectionClosed; + return ret; + } + } + // ret < 0. On a non-blocking socket a TLS read can stop needing either + // direction: the send path shares this session, so output it left pending + // has to be flushed before more input can be decrypted. Anything else is + // a hard error. + auto needs_readable = err.code == tls::ErrorCode::WantRead; +#ifdef _WIN32 + // On Windows a socket timeout surfaces as a syscall error, not WantRead. + needs_readable = + needs_readable || (err.code == tls::ErrorCode::SyscallError && + WSAGetLastError() == WSAETIMEDOUT); +#endif + if (!needs_readable && err.code != tls::ErrorCode::WantWrite) { return -1; } + if (!(needs_readable ? wait_readable() : wait_writable())) { + error_ = Error::Timeout; + return -1; + } + } + return -1; +} + +ssize_t WebSocketSSLStream::write(const char *ptr, size_t size) { + auto handle_size = std::min(size, (std::numeric_limits::max)()); + tls::TlsError err; + auto n = 1000; + while (--n >= 0) { + { + std::lock_guard guard(session_mutex_); + auto ret = tls::write(session_, ptr, handle_size, err); + if (ret >= 0) { return ret; } + } + // ret < 0. As in read(), either direction can be needed: a renegotiation + // or a post-handshake message must be consumed before the record goes + // out. Anything else is a hard error. + auto needs_writable = err.code == tls::ErrorCode::WantWrite; +#ifdef _WIN32 + // On Windows a socket timeout surfaces as a syscall error, not WantWrite. + needs_writable = + needs_writable || (err.code == tls::ErrorCode::SyscallError && + WSAGetLastError() == WSAETIMEDOUT); +#endif + if (!needs_writable && err.code != tls::ErrorCode::WantRead) { return -1; } + if (!(needs_writable ? wait_writable() : wait_readable())) { return -1; } + } + return -1; +} + +void WebSocketSSLStream::get_remote_ip_and_port(std::string &ip, + int &port) const { + detail::get_remote_ip_and_port(sock_, ip, port); +} + +void WebSocketSSLStream::get_local_ip_and_port(std::string &ip, + int &port) const { + detail::get_local_ip_and_port(sock_, ip, port); +} + +socket_t WebSocketSSLStream::socket() const { return sock_; } + +time_t WebSocketSSLStream::duration() const { + return std::chrono::duration_cast( + std::chrono::steady_clock::now() - start_time_) + .count(); +} + +void WebSocketSSLStream::set_read_timeout(time_t sec, time_t usec) { + read_timeout_sec_ = sec; + read_timeout_usec_ = usec; +} + } // namespace detail #endif // CPPHTTPLIB_SSL_ENABLED @@ -7687,6 +8136,57 @@ Server &Server::Options(const std::string &pattern, Handler handler) { return add_handler(options_handlers_, pattern, std::move(handler)); } +const std::set &Server::builtin_methods() { + thread_local const std::set methods{ + "GET", "HEAD", "POST", "PUT", "DELETE", + "CONNECT", "OPTIONS", "TRACE", "PATCH", "PRI"}; + return methods; +} + +Server::CustomHandlerEntry * +Server::custom_entry_for_registration(const std::string &method) { + // Built-in methods are refused for two different reasons. GET, HEAD, POST, + // PUT, DELETE, OPTIONS and PATCH are dispatched by the if/else chain in + // routing() before the custom tables are consulted, so a route registered + // for one of them could never fire. CONNECT, TRACE and PRI have no branch + // there and would be reachable, but they carry protocol-level meaning + // (tunnel setup, request echo, the HTTP/2 connection preface) that this + // library does not route. + if (!detail::fields::is_token(method) || builtin_methods().count(method)) { + output_error_log(Error::InvalidHTTPMethod, nullptr); + has_invalid_registration_ = true; + return nullptr; + } + return &custom_handlers_[method]; +} + +Server &Server::CustomRoute(const std::string &method, + const std::string &pattern, + Handler handler) { + auto *entry = custom_entry_for_registration(method); + if (!entry) { return *this; } + return add_handler(entry->handlers, pattern, std::move(handler)); +} + +Server &Server::CustomRoute(const std::string &method, + const std::string &pattern, + HandlerWithContentReader handler) { + auto *entry = custom_entry_for_registration(method); + if (!entry) { return *this; } + return add_handler(entry->handlers_for_content_reader, pattern, + std::move(handler)); +} + +const Server::CustomHandlerEntry * +Server::find_custom_entry(const std::string &method) const { + // find() alone would be correct here. The empty() check is what keeps the + // per-request cost off servers that never call CustomRoute(), which is the + // overwhelmingly common case; keep it rather than walking into the tree. + if (custom_handlers_.empty()) { return nullptr; } + auto it = custom_handlers_.find(method); + return it == custom_handlers_.end() ? nullptr : &it->second; +} + Server &Server::WebSocket(const std::string &pattern, WebSocketHandler handler) { websocket_handlers_.push_back( @@ -7898,6 +8398,21 @@ Server &Server::set_payload_max_length(size_t length) { return *this; } +Server &Server::set_static_file_compression(bool on) { + static_file_compression_ = on; + return *this; +} + +Server &Server::set_static_file_compression_min_length(size_t length) { + static_file_compression_min_length_ = length; + return *this; +} + +Server &Server::set_static_file_compression_max_length(size_t length) { + static_file_compression_max_length_ = length; + return *this; +} + Server &Server::set_websocket_max_missed_pongs(int count) { websocket_max_missed_pongs_ = count; return *this; @@ -7979,11 +8494,12 @@ bool Server::parse_request_line(const char *s, Request &req) const { if (count != 3) { return false; } } - thread_local const std::set methods{ - "GET", "HEAD", "POST", "PUT", "DELETE", - "CONNECT", "OPTIONS", "TRACE", "PATCH", "PRI"}; + // A method outside the built-in set is accepted only when a handler has been + // registered for it with CustomRoute(). + const auto &methods = builtin_methods(); - if (methods.find(req.method) == methods.end()) { + if (methods.find(req.method) == methods.end() && + !find_custom_entry(req.method)) { output_error_log(Error::InvalidHTTPMethod, &req); return false; } @@ -8044,7 +8560,8 @@ bool Server::write_response_core(Stream &strm, bool close_connection, if (need_apply_ranges) { apply_ranges(req, res, content_type, boundary); } // Prepare additional headers - if (close_connection || req.get_header_value("Connection") == "close" || + if (close_connection || + detail::has_header_token(req.headers, "Connection", "close") || 400 <= res.status) { // Don't leave connections open after errors res.set_header("Connection", "close"); } else { @@ -8352,7 +8869,29 @@ bool Server::handle_file_request(Request &req, Response &res) { res.set_header(kv.first, kv.second); } - auto etag = detail::compute_etag(stat); + auto content_type_of = [&]() { + return detail::find_content_type( + path, file_extension_and_mimetype_map_, default_file_mimetype_); + }; + + // Only the ETag needs the content type this early, and only to name + // the coding. Deciding it here would otherwise put a regex in front + // of the 304 below, which serving a file never used to pay for. + std::string content_type; + auto encoding = detail::EncodingType::None; + if (static_file_compression_) { + content_type = content_type_of(); + encoding = static_file_encoding(req, content_type, stat.size()); + } + + // The ETag names the representation actually sent, so a client that + // cached the compressed form revalidates against the compressed ETag + // and still gets a 304, while one that took identity keeps the plain + // ETag. + auto etag = detail::compute_etag( + stat, encoding == detail::EncodingType::None + ? std::string() + : std::string("-") + detail::encoding_name(encoding)); if (!etag.empty()) { res.set_header("ETag", etag); } auto mtime = stat.mtime(); @@ -8372,14 +8911,9 @@ bool Server::handle_file_request(Request &req, Response &res) { return false; } - res.set_content_provider( - mm->size(), - detail::find_content_type(path, file_extension_and_mimetype_map_, - default_file_mimetype_), - [mm](size_t offset, size_t length, DataSink &sink) -> bool { - sink.write(mm->data() + offset, length); - return true; - }); + if (!static_file_compression_) { content_type = content_type_of(); } + + detail::set_file_content_provider(res, mm, content_type, encoding); if (req.method != "HEAD" && file_request_handler_) { file_request_handler_(req, res); @@ -8403,7 +8937,8 @@ bool Server::check_if_not_modified(const Request &req, Response &res, // 2. If-Modified-Since is checked only when If-None-Match is absent if (req.has_header("If-None-Match")) { if (!etag.empty()) { - auto val = req.get_header_value("If-None-Match"); + auto val = + detail::get_combined_header_value(req.headers, "If-None-Match"); // NOTE: We use exact string matching here. This works correctly // because our server always generates weak ETags (W/"..."), and @@ -8564,16 +9099,26 @@ bool Server::listen_internal() { #endif if (sock == INVALID_SOCKET) { - if (errno == EMFILE) { - // The per-process limit of open file descriptors has been reached. - // Try to accept new connections after a short sleep. + // NOTE: Winsock reports failures through WSAGetLastError() and never + // touches the CRT errno, so the two have to be asked platform by + // platform rather than by testing errno here. + if (detail::is_accept_resource_error()) { + // The per-process descriptor limit or the network stack's buffer + // space has been reached. Try to accept new connections after a + // short sleep. std::this_thread::sleep_for(std::chrono::microseconds{1}); continue; - } else if (errno == EINTR || errno == EAGAIN) { + } else if (detail::is_accept_transient_error()) { continue; } - if (svr_sock_ != INVALID_SOCKET) { - detail::close_socket(svr_sock_); + // Take the descriptor out of svr_sock_ before closing it: a later + // stop() would otherwise shutdown()/close() a value the OS may have + // reused, and keep_alive() watches svr_sock_ to notice the server is + // gone. The exchange also settles the race with a concurrent stop(), + // since whichever side takes the descriptor closes it exactly once. + auto listen_sock = svr_sock_.exchange(INVALID_SOCKET); + if (listen_sock != INVALID_SOCKET) { + detail::close_socket(listen_sock); ret = false; output_error_log(Error::Connection, nullptr); } else { @@ -8616,7 +9161,14 @@ bool Server::routing(Request &req, Response &res, Stream &strm) { return true; } - if (detail::expect_content(req)) { + const auto *custom = find_custom_entry(req.method); + + // The second clause mirrors what expect_content() does unconditionally for + // POST/PUT/PATCH/DELETE: a content reader route fires even when the request + // carries no body. Without it a body-less PROPFIND (RFC 4918 treats one as + // `allprop`) would skip its handler and fall through to 404. + if (detail::expect_content(req) || + (custom && !custom->handlers_for_content_reader.empty())) { // Content reader handler { // Track whether the ContentReader was aborted due to the decompressed @@ -8663,6 +9215,9 @@ bool Server::routing(Request &req, Response &res, Stream &strm) { } else if (req.method == "DELETE") { dispatched = dispatch_request_for_content_reader( req, res, std::move(reader), delete_handlers_for_content_reader_); + } else if (custom) { + dispatched = dispatch_request_for_content_reader( + req, res, std::move(reader), custom->handlers_for_content_reader); } if (dispatched) { @@ -8699,6 +9254,8 @@ bool Server::routing(Request &req, Response &res, Stream &strm) { return dispatch_request(req, res, options_handlers_, strm); } else if (req.method == "PATCH") { return dispatch_request(req, res, patch_handlers_, strm); + } else if (custom) { + return dispatch_request(req, res, custom->handlers, strm); } res.status = StatusCode::BadRequest_400; @@ -8735,9 +9292,88 @@ bool Server::dispatch_request(Request &req, Response &res, return false; } +// Decides the content coding for a response served straight from a file. Both +// the ETag, which has to name the representation actually sent, and +// `apply_static_file_compression()` go through this, so the two cannot drift +// apart. +detail::EncodingType Server::static_file_encoding( + const Request &req, const std::string &content_type, size_t length) const { + if (!static_file_compression_) { return detail::EncodingType::None; } + + // Nothing to compress, and an empty file already answers with + // `Content-Length: 0`. Checked on its own so that a zero floor still cannot + // turn an empty body into a 20-byte gzip stream. + if (length == 0) { return detail::EncodingType::None; } + + // A file that already fits in a single packet gains nothing from being made + // smaller, since it still travels in that one segment, and a file of a few + // bytes comes out larger than it went in. + if (length < static_file_compression_min_length_) { + return detail::EncodingType::None; + } + + // RFC 9110 applies Range to the representation after content coding, so a + // compressed 206 would mean compressing the whole file and then slicing it. + // Serve ranges from the identity representation instead. + if (!req.ranges.empty()) { return detail::EncodingType::None; } + + if (static_file_compression_max_length_ > 0 && + length > static_file_compression_max_length_) { + return detail::EncodingType::None; + } + + return detail::encoding_type(req, content_type); +} + +// Compresses a file-backed content provider into `res.body` and takes over the +// framing headers. Returns false when the response is left untouched. +bool Server::apply_static_file_compression(const Request &req, + Response &res) const { + auto type = res.file_content_encoding_; + if (type == detail::EncodingType::None || !res.content_provider_) { + return false; + } + + auto compressor = detail::make_compressor(type); + if (!compressor) { return false; } + + output_pre_compression_log(req, res); + + std::string compressed; + if (!detail::compress_content_provider(res.content_provider_, + res.content_length_, *compressor, + compressed)) { + return false; + } + + res.body.swap(compressed); + + // The provider was consumed in full, so a resource releaser registered with + // it should hear about a success when the response goes away. + res.content_provider_success_ = true; + res.content_provider_ = nullptr; + res.content_length_ = 0; + res.file_content_encoding_ = detail::EncodingType::None; + + res.set_header("Content-Encoding", detail::encoding_name(type)); + res.set_header("Vary", "Accept-Encoding"); + res.set_header("Content-Length", std::to_string(res.body.size())); + + return true; +} + void Server::apply_ranges(const Request &req, Response &res, std::string &content_type, std::string &boundary) const { + // A known-length content provider leaves `res.body` empty, so the compressor + // at the end of this function never runs for one (issue #2545). A file-backed + // provider is fully readable right here, so compress it and answer with an + // ordinary body: `Content-Length` and HEAD keep working, and the response + // takes the same path as `set_content()` from here on. Range requests never + // get a content coding, so `Content-Range` still names identity bytes and + // none of the framing below applies. + if (apply_static_file_compression(req, res)) { return; } + if (req.ranges.size() > 1 && res.status == StatusCode::PartialContent_206) { auto it = res.headers.find("Content-Type"); if (it != res.headers.end()) { @@ -8948,12 +9584,12 @@ Server::process_request(Stream &strm, const std::string &remote_addr, return write_response(strm, close_connection, req, res); } - if (req.get_header_value("Connection") == "close") { + if (detail::has_header_token(req.headers, "Connection", "close")) { connection_closed = true; } if (req.version == "HTTP/1.0" && - req.get_header_value("Connection") != "Keep-Alive") { + !detail::has_header_token(req.headers, "Connection", "keep-alive")) { connection_closed = true; } @@ -8965,7 +9601,13 @@ Server::process_request(Stream &strm, const std::string &remote_addr, [&](const std::string &proxy) { return proxy == remote_addr; }); if (is_trusted_peer && req.has_header("X-Forwarded-For")) { - auto x_forwarded_for = req.get_header_value("X-Forwarded-For"); + // Some proxies append the address they observed as a separate + // X-Forwarded-For field line instead of extending the one the client sent + // (e.g. HAProxy's "option forwardfor"), so the whole combined value has to + // be scanned. Reading only the first occurrence would hand back the + // client-supplied, and therefore forgeable, value. + auto x_forwarded_for = + detail::get_combined_header_value(req.headers, "X-Forwarded-For"); auto derived = get_client_ip(x_forwarded_for, trusted_proxies_); req.remote_addr = derived.empty() ? remote_addr : derived; } else { @@ -8977,7 +9619,8 @@ Server::process_request(Stream &strm, const std::string &remote_addr, req.local_port = local_port; if (req.has_header("Accept")) { - const auto &accept_header = req.get_header_value("Accept"); + auto accept_header = + detail::get_combined_header_value(req.headers, "Accept"); if (!detail::parse_accept_header(accept_header, req.accept_content_types)) { connection_closed = true; res.status = StatusCode::BadRequest_400; @@ -8998,7 +9641,12 @@ Server::process_request(Stream &strm, const std::string &remote_addr, if (setup_request) { setup_request(req); } - if (req.get_header_value("Expect") == "100-continue") { + // RFC 9110 10.1.1: Expect is a comma-separated list whose value is + // case-insensitive, and a 100-continue expectation in an HTTP/1.0 request + // must be ignored. An expectation we do not recognize is left alone; the + // 417 the section allows for one is a MAY, not a requirement. + if (req.version != "HTTP/1.0" && + detail::has_header_token(req.headers, "Expect", "100-continue")) { int status = StatusCode::Continue_100; if (expect_100_continue_handler_) { status = expect_100_continue_handler_(req, res); @@ -9041,19 +9689,15 @@ Server::process_request(Stream &strm, const std::string &remote_addr, // Negotiate subprotocol std::string selected_subprotocol; if (entry.sub_protocol_selector) { - auto protocol_header = req.get_header_value("Sec-WebSocket-Protocol"); + auto protocol_header = detail::get_combined_header_value( + req.headers, "Sec-WebSocket-Protocol"); if (!protocol_header.empty()) { std::vector protocols; - std::istringstream iss(protocol_header); - std::string token; - while (std::getline(iss, token, ',')) { - // Trim whitespace - auto start = token.find_first_not_of(' '); - auto end = token.find_last_not_of(' '); - if (start != std::string::npos) { - protocols.push_back(token.substr(start, end - start + 1)); - } - } + detail::split(protocol_header.data(), + protocol_header.data() + protocol_header.size(), ',', + [&](const char *b, const char *e) { + protocols.emplace_back(b, e); + }); selected_subprotocol = entry.sub_protocol_selector(protocols); } } @@ -9081,6 +9725,24 @@ Server::process_request(Stream &strm, const std::string &remote_addr, if (websocket_upgraded) { *websocket_upgraded = true; } { +#ifdef CPPHTTPLIB_SSL_ENABLED + if (req.ssl) { + // wss: the heartbeat ping thread and the read path enter the same + // TLS session from different threads. Hand the WebSocket a stream + // that serializes every TLS call, so the shared SSLSocketStream on + // the plain HTTP/HTTPS paths stays untouched. + auto ws_strm = + std::unique_ptr(new detail::WebSocketSSLStream( + strm.socket(), const_cast(req.ssl), + CPPHTTPLIB_WEBSOCKET_READ_TIMEOUT_SECOND, 0, + write_timeout_sec_, write_timeout_usec_)); + ws::WebSocket ws(std::move(ws_strm), req, true, + websocket_ping_interval_sec_, + websocket_max_missed_pongs_); + entry.handler(req, ws); + return true; + } +#endif // Use WebSocket-specific read timeout instead of HTTP timeout strm.set_read_timeout(CPPHTTPLIB_WEBSOCKET_READ_TIMEOUT_SECOND, 0); ws::WebSocket ws(strm, req, true, websocket_ping_interval_sec_, @@ -9144,12 +9806,9 @@ Server::process_request(Stream &strm, const std::string &remote_addr, path, file_extension_and_mimetype_map_, default_file_mimetype_); } - res.set_content_provider( - mm->size(), content_type, - [mm](size_t offset, size_t length, DataSink &sink) -> bool { - sink.write(mm->data() + offset, length); - return true; - }); + detail::set_file_content_provider( + res, mm, content_type, + static_file_encoding(req, content_type, mm->size())); } } @@ -9174,7 +9833,7 @@ Server::process_request(Stream &strm, const std::string &remote_addr, // consume the next request (issue #2450). If the response has committed the // connection to close, there is no next request to protect. if (!req.body_consumed_ && detail::has_framed_body(req)) { - if (res.get_header_value("Connection") == "close") { + if (detail::has_header_token(res.headers, "Connection", "close")) { connection_closed = true; } else { int dummy_status; @@ -9190,7 +9849,7 @@ Server::process_request(Stream &strm, const std::string &remote_addr, return ret; } -bool Server::is_valid() const { return true; } +bool Server::is_valid() const { return !has_invalid_registration_; } bool Server::process_and_close_socket(socket_t sock) { std::string remote_addr; @@ -9202,15 +9861,18 @@ bool Server::process_and_close_socket(socket_t sock) { detail::get_local_ip_and_port(sock, local_addr, local_port); bool websocket_upgraded = false; - auto ret = detail::process_server_socket( - svr_sock_, sock, keep_alive_max_count_, keep_alive_timeout_sec_, - read_timeout_sec_, read_timeout_usec_, write_timeout_sec_, - write_timeout_usec_, - [&](Stream &strm, bool close_connection, bool &connection_closed) { - return process_request(strm, remote_addr, remote_port, local_addr, - local_port, close_connection, connection_closed, - nullptr, &websocket_upgraded); - }); + auto ret = serve_guarded([&]() { + return detail::process_server_socket( + svr_sock_, sock, keep_alive_max_count_, keep_alive_timeout_sec_, + read_timeout_sec_, read_timeout_usec_, write_timeout_sec_, + write_timeout_usec_, + [&](Stream &strm, bool close_connection, bool &connection_closed) { + return process_request(strm, remote_addr, remote_port, local_addr, + local_port, close_connection, + connection_closed, nullptr, + &websocket_upgraded); + }); + }); detail::drain_and_close_socket(sock); return ret; @@ -9736,7 +10398,8 @@ ClientImpl::open_stream(const std::string &method, const std::string &path, handle.body_reader_.chunked = detail::is_chunked_transfer_encoding(handle.response->headers); - auto content_encoding = handle.response->get_header_value("Content-Encoding"); + auto content_encoding = detail::get_combined_header_value( + handle.response->headers, "Content-Encoding"); if (!content_encoding.empty()) { // Same policy as prepare_content_receiver(): reject a coding we know about // but were not built with, pass an unrecognized one through as-is. @@ -9958,7 +10621,7 @@ bool ClientImpl::handle_request(Stream &strm, Request &req, if (!ret) { return false; } - if (res.get_header_value("Connection") == "close" || + if (detail::has_header_token(res.headers, "Connection", "close") || (res.version == "HTTP/1.0" && res.reason != "Connection established")) { // NOTE: this requires a not-entirely-obvious chain of calls to be correct // for this to be safe. @@ -10437,6 +11100,7 @@ ClientImpl::send_with_content_provider_and_receiver( if (content_provider) { auto ok = true; + auto finished = false; size_t offset = 0; DataSink data_sink; @@ -10460,13 +11124,27 @@ ClientImpl::send_with_content_provider_and_receiver( return ok; }; - while (ok && offset < content_length) { + // As in detail::write_content_with_progress(): the body is framed by + // content_length, so a provider that finishes early has truncated it. + // Stop and report that instead of calling the provider forever. + data_sink.done = [&]() { finished = true; }; + + while (ok && !finished && offset < content_length) { if (!content_provider(offset, content_length - offset, data_sink)) { error = Error::Canceled; output_error_log(error, &req); return nullptr; } } + + // A short body here means either the provider stopped early or the + // compressor gave up. The branch below reports a failing compressor as + // Error::Compression, so keep the two distinguishable. + if (offset < content_length) { + error = ok ? Error::Write : Error::Compression; + output_error_log(error, &req); + return nullptr; + } } else { if (!compressor->compress(body, content_length, true, [&](const char *data, size_t data_len) { @@ -10564,7 +11242,8 @@ bool ClientImpl::process_request(Stream &strm, Request &req, } // Check for Expect: 100-continue - auto expect_100_continue = req.get_header_value("Expect") == "100-continue"; + auto expect_100_continue = + detail::has_header_token(req.headers, "Expect", "100-continue"); // Send request (skip body if using Expect: 100-continue) auto write_request_success = @@ -10768,6 +11447,9 @@ ContentProviderWithoutLength ClientImpl::get_multipart_content_provider( DataSink cur_sink; auto has_data = true; cur_sink.write = sink.write; + // Forward is_writable so a provider item asking whether it may keep + // going gets the outer sink's answer rather than the default `true`. + cur_sink.is_writable = sink.is_writable; cur_sink.done = [&]() { has_data = false; }; if (!provider_items[cur_item].provider(offset - cur_start, cur_sink)) { @@ -12529,7 +13211,9 @@ SSLServer::~SSLServer() { if (ctx_) { tls::free_context(ctx_); } } -bool SSLServer::is_valid() const { return ctx_ != nullptr; } +bool SSLServer::is_valid() const { + return ctx_ != nullptr && Server::is_valid(); +} bool SSLServer::process_and_close_socket(socket_t sock) { using namespace tls; @@ -12588,16 +13272,18 @@ bool SSLServer::process_and_close_socket(socket_t sock) { int local_port = 0; detail::get_local_ip_and_port(sock, local_addr, local_port); - ret = detail::process_server_socket_ssl( - svr_sock_, session, sock, keep_alive_max_count_, keep_alive_timeout_sec_, - read_timeout_sec_, read_timeout_usec_, write_timeout_sec_, - write_timeout_usec_, - [&](Stream &strm, bool close_connection, bool &connection_closed) { - return process_request( - strm, remote_addr, remote_port, local_addr, local_port, - close_connection, connection_closed, - [&](Request &req) { req.ssl = session; }, &websocket_upgraded); - }); + ret = serve_guarded([&]() { + return detail::process_server_socket_ssl( + svr_sock_, session, sock, keep_alive_max_count_, + keep_alive_timeout_sec_, read_timeout_sec_, read_timeout_usec_, + write_timeout_sec_, write_timeout_usec_, + [&](Stream &strm, bool close_connection, bool &connection_closed) { + return process_request( + strm, remote_addr, remote_port, local_addr, local_port, + close_connection, connection_closed, + [&](Request &req) { req.ssl = session; }, &websocket_upgraded); + }); + }); return ret; } @@ -16870,6 +17556,7 @@ bool WebSocket::send_frame(Opcode op, const char *data, size_t len, } ReadResult WebSocket::read(std::string &msg) { + std::unique_lock read_lock(read_mutex_); while (!closed_) { Opcode opcode; std::string payload; @@ -16955,6 +17642,9 @@ ReadResult WebSocket::read(std::string &msg) { } // RFC 6455 Section 5.6: text frames must contain valid UTF-8 if (result == Text && !impl::is_valid_utf8(msg)) { + // close() takes the read lock to wait for the peer's Close reply, so + // it must not run while this thread still holds it. + read_lock.unlock(); close(CloseStatus::InvalidPayload, "invalid UTF-8"); return Fail; } @@ -16991,9 +17681,18 @@ void WebSocket::close(CloseStatus status, const std::string &reason) { } // RFC 6455 Section 7.1.1: after sending a Close frame, wait for the peer's - // Close response before closing the TCP connection. Use a short timeout to - // avoid hanging if the peer doesn't respond. + // Close response before closing the TCP connection. + // + // Wait only when no other thread is parsing frames. When one is, it is the + // thread positioned to see the peer's reply, and reading here would take + // bytes out of the message it is assembling. Bailing out also leaves the + // stream, including its read timeout, entirely to that thread. + std::unique_lock read_lock(read_mutex_, std::try_to_lock); + if (!read_lock.owns_lock()) { return; } + + // Use a short timeout to avoid hanging if the peer doesn't respond. strm_.set_read_timeout(CPPHTTPLIB_WEBSOCKET_CLOSE_TIMEOUT_SECOND, 0); + Opcode op; std::string resp; bool fin; @@ -17171,7 +17870,7 @@ bool WebSocketClient::create_stream(std::unique_ptr &strm, return false; } - strm = std::unique_ptr(new detail::SSLSocketStream( + strm = std::unique_ptr(new detail::WebSocketSSLStream( sock_, tls_session_, read_timeout_sec_, read_timeout_usec_, write_timeout_sec_, write_timeout_usec_)); return true; diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h index 6fc86c7c75bf..ca7c96a41a38 100644 --- a/vendor/cpp-httplib/httplib.h +++ b/vendor/cpp-httplib/httplib.h @@ -8,8 +8,8 @@ #ifndef CPPHTTPLIB_HTTPLIB_H #define CPPHTTPLIB_HTTPLIB_H -#define CPPHTTPLIB_VERSION "0.53.1" -#define CPPHTTPLIB_VERSION_NUM "0x003501" +#define CPPHTTPLIB_VERSION "0.54.1" +#define CPPHTTPLIB_VERSION_NUM "0x003601" #ifdef _WIN32 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00 @@ -134,6 +134,16 @@ #define CPPHTTPLIB_FORM_URL_ENCODED_PAYLOAD_MAX_LENGTH 8192 #endif +#ifndef CPPHTTPLIB_STATIC_FILE_COMPRESSION_MIN_LENGTH +// 1400 rather than a round number: a body that already fits in one 1500-byte +// MTU gains nothing from being made smaller. +#define CPPHTTPLIB_STATIC_FILE_COMPRESSION_MIN_LENGTH 1400 +#endif + +#ifndef CPPHTTPLIB_STATIC_FILE_COMPRESSION_MAX_LENGTH +#define CPPHTTPLIB_STATIC_FILE_COMPRESSION_MAX_LENGTH (4 * 1024 * 1024) // 4MB +#endif + #ifndef CPPHTTPLIB_RANGE_MAX_COUNT #define CPPHTTPLIB_RANGE_MAX_COUNT 1024 #endif @@ -1429,9 +1439,16 @@ class DataSink { DataSink &operator=(DataSink &&) = delete; std::function write; - std::function is_writable; - std::function done; - std::function done_with_trailer; + + // Only `write` is mandatory. The rest are defaulted so that a provider + // calling one on a writer that does not set it gets sensible behaviour + // rather than std::bad_function_call thrown from a worker thread. Capturing + // `this` is safe: DataSink is neither copyable nor movable. + std::function is_writable = []() { return true; }; + std::function done = []() {}; + std::function done_with_trailer = + [this](const Headers & /*trailer*/) { done(); }; + std::ostream os; private: @@ -1516,7 +1533,10 @@ make_file_body(const std::string &filepath) { auto to_read = (std::min)(sizeof(buf), length); f.read(buf, static_cast(to_read)); auto n = static_cast(f.gcount()); - if (n == 0) { break; } + // The file is shorter than the size make_file_body() measured, which the + // caller has already committed to as Content-Length. The body cannot be + // completed, so fail as every other error here does. + if (n == 0) { return false; } if (!sink.write(buf, n)) { return false; } length -= n; } @@ -1723,6 +1743,14 @@ struct Request { #endif }; +namespace detail { + +// Declared up here, away from the rest of the compression helpers, because +// `Response` stores one. +enum class EncodingType { None = 0, Gzip, Brotli, Zstd }; + +} // namespace detail + struct Response { std::string version; int status = -1; @@ -1788,6 +1816,11 @@ struct Response { bool content_provider_success_ = false; std::string file_content_path_; std::string file_content_content_type_; + + // Content coding chosen for a file-backed content provider, decided once + // where the file is opened so that the ETag and the body cannot disagree. + // `EncodingType::None` for every other kind of response. + detail::EncodingType file_content_encoding_ = detail::EncodingType::None; }; enum class Error { @@ -1827,6 +1860,7 @@ enum class Error { InvalidRangeHeader, UnsupportedContentEncoding, WebSocketHandshake, + UserCallbackException, // For internal use only SSLPeerCouldBeClosed_, @@ -2020,6 +2054,10 @@ class RegexMatcher final : public MatcherBase { int close_socket(socket_t sock) noexcept; +bool is_accept_resource_error(); + +bool is_accept_transient_error(); + ssize_t write_headers(Stream &strm, const Headers &headers); bool set_socket_opt_time(socket_t sock, int level, int optname, time_t sec, @@ -2107,6 +2145,17 @@ class Server { Server &Delete(const std::string &pattern, HandlerWithContentReader handler); Server &Options(const std::string &pattern, Handler handler); + // Register a handler for an HTTP method outside the built-in set (e.g. the + // WebDAV methods from RFC 4918). Registering a method here is what makes the + // server accept it; an unregistered method is still rejected with 400. + // `method` must be a valid HTTP method token and must not be one of the + // built-in methods, which have their own registration functions above. A + // rejected registration makes is_valid() return false, so listen() fails. + Server &CustomRoute(const std::string &method, const std::string &pattern, + Handler handler); + Server &CustomRoute(const std::string &method, const std::string &pattern, + HandlerWithContentReader handler); + Server &WebSocket(const std::string &pattern, WebSocketHandler handler); Server &WebSocket(const std::string &pattern, WebSocketHandler handler, SubProtocolSelector sub_protocol_selector); @@ -2174,6 +2223,10 @@ class Server { Server &set_payload_max_length(size_t length); + Server &set_static_file_compression(bool on); + Server &set_static_file_compression_min_length(size_t length); + Server &set_static_file_compression_max_length(size_t length); + Server &set_websocket_ping_interval(time_t sec); template Server &set_websocket_ping_interval( @@ -2202,6 +2255,35 @@ class Server { const std::function &setup_request, bool *websocket_upgraded = nullptr); + // Runs the per-connection serving loop and stops an exception thrown by a + // user callback from escaping the worker thread. + // + // process_request() wraps only routing() in a try/catch. Content providers, + // the post-routing, error, logging and expect-100 handlers and WebSocket + // handlers all run outside it, and the task queue calls the job without a + // catch, so an exception from any of those would terminate the process. + // + // No 500 is possible here: by the time a content provider runs, the status + // line and headers are already on the wire. Report it through the error + // logger and drop the connection, which is what the peer observes either + // way. Other connections are unaffected. + template bool serve_guarded(Serve &&serve) const { +#ifdef CPPHTTPLIB_NO_EXCEPTIONS + return serve(); +#else + try { + return serve(); + } catch (...) { + // The error logger is a user callback too, so it must not be able to + // throw the guard back open. + try { + output_error_log(Error::UserCallbackException, nullptr); + } catch (...) {} + return false; + } +#endif + } + std::atomic svr_sock_{INVALID_SOCKET}; std::vector trusted_proxies_; @@ -2215,6 +2297,11 @@ class Server { time_t idle_interval_sec_ = CPPHTTPLIB_IDLE_INTERVAL_SECOND; time_t idle_interval_usec_ = CPPHTTPLIB_IDLE_INTERVAL_USECOND; size_t payload_max_length_ = CPPHTTPLIB_PAYLOAD_MAX_LENGTH; + bool static_file_compression_ = false; + size_t static_file_compression_min_length_ = + CPPHTTPLIB_STATIC_FILE_COMPRESSION_MIN_LENGTH; + size_t static_file_compression_max_length_ = + CPPHTTPLIB_STATIC_FILE_COMPRESSION_MAX_LENGTH; time_t websocket_ping_interval_sec_ = CPPHTTPLIB_WEBSOCKET_PING_INTERVAL_SECOND; int websocket_max_missed_pongs_ = CPPHTTPLIB_WEBSOCKET_MAX_MISSED_PONGS; @@ -2226,9 +2313,21 @@ class Server { std::vector, HandlerWithContentReader>>; + // Both handler tables for one custom method live in a single entry, so that + // routing() needs only one map lookup per request to reach either of them. + struct CustomHandlerEntry { + Handlers handlers; + HandlersForContentReader handlers_for_content_reader; + }; + using CustomHandlers = std::map; + static std::unique_ptr make_matcher(const std::string &pattern); + static const std::set &builtin_methods(); + CustomHandlerEntry *custom_entry_for_registration(const std::string &method); + const CustomHandlerEntry *find_custom_entry(const std::string &method) const; + template Server &add_handler( std::vector, H>> &handlers, @@ -2259,6 +2358,10 @@ class Server { const HandlersForContentReader &handlers) const; bool parse_request_line(const char *s, Request &req) const; + detail::EncodingType static_file_encoding(const Request &req, + const std::string &content_type, + size_t length) const; + bool apply_static_file_compression(const Request &req, Response &res) const; void apply_ranges(const Request &req, Response &res, std::string &content_type, std::string &boundary) const; bool write_response(Stream &strm, bool close_connection, Request &req, @@ -2292,6 +2395,10 @@ class Server { std::atomic is_running_{false}; std::atomic is_decommissioned{false}; + // Set when CustomRoute() refuses a registration. Written before listen(), + // read by is_valid() on the same thread, so it needs no synchronization. + bool has_invalid_registration_ = false; + struct MountPointEntry { std::string mount_point; std::string base_dir; @@ -2313,6 +2420,7 @@ class Server { Handlers delete_handlers_; HandlersForContentReader delete_handlers_for_content_reader_; Handlers options_handlers_; + CustomHandlers custom_handlers_; struct WebSocketHandlerEntry { std::unique_ptr matcher; @@ -3500,6 +3608,16 @@ void split(const char *b, const char *e, char d, void split(const char *b, const char *e, char d, size_t m, std::function fn); +bool split_find(const char *b, const char *e, char d, + std::function fn); + +bool has_header_token(const Headers &headers, const std::string &key, + const std::string &token); + +std::string websocket_accept_key(const std::string &client_key); + +bool is_websocket_upgrade(const Request &req); + bool process_client_socket( socket_t sock, time_t read_timeout_sec, time_t read_timeout_usec, time_t write_timeout_sec, time_t write_timeout_usec, @@ -3520,6 +3638,9 @@ socket_t create_client_socket(const std::string &host, const std::string &ip, const char *get_header_value(const Headers &headers, const std::string &key, const char *def, size_t id); +std::string get_combined_header_value(const Headers &headers, + const std::string &key); + std::string params_to_query_str(const Params ¶ms); void parse_query_text(const char *data, std::size_t size, Params ¶ms); @@ -3534,11 +3655,13 @@ bool parse_range_header(const std::string &s, Ranges &ranges); bool parse_accept_header(const std::string &s, std::vector &content_types); +void parse_disposition_params(const std::string &s, Params ¶ms); + ssize_t send_socket(socket_t sock, const void *ptr, size_t size, int flags); ssize_t read_socket(socket_t sock, void *ptr, size_t size, int flags); -enum class EncodingType { None = 0, Gzip, Brotli, Zstd }; +EncodingType encoding_type(const Request &req, const std::string &content_type); EncodingType encoding_type(const Request &req, const Response &res); @@ -4318,6 +4441,11 @@ class WebSocket { int unacked_pings_ = 0; std::atomic closed_{false}; std::mutex write_mutex_; + // Owned by whichever thread is parsing frames off strm_. Only one thread + // may do so: read_websocket_frame() reads a payload until it has the whole + // declared length, so a second parser stealing bytes silently corrupts the + // message the first one is assembling. + std::mutex read_mutex_; std::thread ping_thread_; std::mutex ping_mutex_; std::condition_variable ping_cv_; From b8b743c3c1707251d2be304557b0740078543151 Mon Sep 17 00:00:00 2001 From: Daya Adianto Date: Sun, 30 Aug 2026 06:02:22 +0000 Subject: [PATCH 036/109] metal : Add fa-vec tuning for M3 Pro (#27963) Related issue: #27668 --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 212 ++++++++++++++++++++++ 1 file changed, 212 insertions(+) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 4abdafb48f30..9345892aeae8 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -720,6 +720,218 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 128, 128, 2, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 128, 128, 2, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 192, 192, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 192, 128, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 256, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 128, 128, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 192, 192, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 128, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 128, 128, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 128, 128, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 256, 256, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 96, 96, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 128, 128, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 128, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 128, 128, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 128, 128, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 256, 256, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } }, { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } }, From 093adb242e6d205d06979a390d4f4f690dd87bf1 Mon Sep 17 00:00:00 2001 From: Nils Gladitz Date: Sun, 30 Aug 2026 08:06:29 +0200 Subject: [PATCH 037/109] metal: add fa-vec tunings for M3 Ultra (#27999) --- ggml/src/ggml-metal/ggml-metal-tuning.cpp | 198 ++++++++++++++++++++++ 1 file changed, 198 insertions(+) diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 9345892aeae8..9114ab7424de 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -1024,6 +1024,204 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } }, { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 64, 64, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 96, 96, 1, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 128, 128, 1, 1 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 192, 192, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 192, 128, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 256, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 512, 512, 3, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_F16, 512, 512, 3, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 192, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 320, 256, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 320, 256, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 192, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 192, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 192, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 128, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 128, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 32, 32, 3, 4 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 96, 96, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 96, 96, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 128, 128, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 256, 256, 1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 256, 256, 2, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 256, 256, 2, 4 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 320, 256, 1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 32, 32, 3, 4 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 96, 96, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 96, 96, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 128, 128, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 256, 256, 2, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 256, 256, 2, 4 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 320, 256, 1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 320, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 320, 256, 1, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 512, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q5_1, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 128, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 128, 128, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 128, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 128, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 128, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 128, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 128, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 192, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 512, 512, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } }, { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, From 742347b2e717dcebd432437ddf5d088dc4fd9232 Mon Sep 17 00:00:00 2001 From: Ryan C Date: Sun, 30 Aug 2026 06:16:26 +0000 Subject: [PATCH 038/109] rpc: fix apple rdma error spew on teardown (#27908) --- ggml/src/ggml-rpc/transport-apple.cpp | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/ggml/src/ggml-rpc/transport-apple.cpp b/ggml/src/ggml-rpc/transport-apple.cpp index 2cfaa5d3fd3d..b1934175b1dd 100644 --- a/ggml/src/ggml-rpc/transport-apple.cpp +++ b/ggml/src/ggml-rpc/transport-apple.cpp @@ -115,16 +115,9 @@ struct apple_rdma::impl { ~impl() { broken = true; - // the QP must be destroyed before the memory it can still write to is - // deregistered and freed: ERR only starts flushing the posted WQEs - if (qp) { - struct ibv_qp_attr a = {}; - a.qp_state = IBV_QPS_ERR; - ibv_modify_qp(qp, &a, IBV_QP_STATE); - struct ibv_wc wc[RDMA_NBUF * 2]; - while (ibv_poll_cq(cq, RDMA_NBUF * 2, wc) > 0) {} - ibv_destroy_qp(qp); - } + // destroy the QP first: it can still write to the rings until it is gone. + // no IBV_QPS_ERR before it - Apple's provider then fails every region unmap. + if (qp) ibv_destroy_qp(qp); if (send_mr) ibv_dereg_mr(send_mr); if (recv_mr) ibv_dereg_mr(recv_mr); free(send_mem); From 73f56d105bb6b5aeb37d0c7dcc6a7d58c2f7974a Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sun, 30 Aug 2026 09:17:47 +0300 Subject: [PATCH 039/109] ggml : add ggml_backend_op_alloc_size_may_expand, use it in RPC (#27960) some backends (Metal, SYCL, WebGPU) require additional memory for fleeting data for certain ops, which is reflected in their get_alloc_size implementations. add ggml_backend_op_alloc_size_may_expand() to the backend utils, listing these ops, and assert in ggml_backend_buft_get_alloc_size that a backend expanding the alloc size of a compute op only does so for ops listed in the helper. use the helper in the RPC backend to decide whether to query the remote server for the actual alloc size, instead of a hardcoded list. Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/include/ggml-backend.h | 4 ++++ ggml/src/ggml-backend.cpp | 23 +++++++++++++++++++++++ ggml/src/ggml-rpc/ggml-rpc.cpp | 6 +++--- 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e35..27375bd0a51e 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -424,6 +424,10 @@ extern "C" { // Compare the output of two backends GGML_API bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t backend2, struct ggml_cgraph * graph, ggml_backend_eval_callback callback, void * user_data, struct ggml_tensor const * const * test_nodes, size_t num_test_nodes); + // returns true for ops that may require additional memory for fleeting data on some backends, + // i.e. the backend's get_alloc_size may return more than ggml_nbytes for the output tensor + GGML_API bool ggml_backend_op_alloc_size_may_expand(enum ggml_op op); + // Tensor initialization GGML_API enum ggml_status ggml_backend_tensor_alloc(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, void * addr); GGML_API enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 78eb10dfe992..fec7d7c92bf3 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -65,6 +65,13 @@ size_t ggml_backend_buft_get_alloc_size(ggml_backend_buffer_type_t buft, const s if (buft->iface.get_alloc_size) { size_t size = buft->iface.get_alloc_size(buft, tensor); assert(size >= ggml_nbytes(tensor)); + + // [TAG_ALLOC_SIZE_EXPAND] + // if you hit this assert, update ggml_backend_op_alloc_size_may_expand() accordingly + GGML_ASSERT(size <= ggml_nbytes(tensor) || + ggml_op_is_empty(tensor->op) || + ggml_backend_op_alloc_size_may_expand(tensor->op)); + return size; } return ggml_nbytes(tensor); @@ -2101,6 +2108,22 @@ ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched, // utils +// [TAG_ALLOC_SIZE_EXPAND] +// returns true for ops that may require additional memory for fleeting data on some backends, +// i.e. the backend's get_alloc_size may return more than ggml_nbytes for the output tensor +bool ggml_backend_op_alloc_size_may_expand(enum ggml_op op) { + switch (op) { + case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_CUMSUM: + case GGML_OP_ARGSORT: + case GGML_OP_TOP_K: + return true; + default: + return false; + } +} + enum ggml_status ggml_backend_view_init(struct ggml_tensor * tensor) { GGML_ASSERT(tensor); GGML_ASSERT(tensor->buffer == NULL); diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 9aa5883d80de..58a8a030cfa5 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -826,10 +826,10 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty // See comments in init_tensor. rpc_get |= ggml_is_quantized(tensor->type) && (tensor->ne[0] % 512 != 0) && (tensor->view_src == nullptr); - // ops that require additional memory for fleeting data on certain backends + // [TAG_ALLOC_SIZE_EXPAND] + // ops that may require additional memory for fleeting data on certain backends // ref: https://github.com/ggml-org/llama.cpp/pull/15966 - rpc_get |= tensor->op == GGML_OP_FLASH_ATTN_EXT; - rpc_get |= tensor->op == GGML_OP_MUL_MAT_ID; + rpc_get |= ggml_backend_op_alloc_size_may_expand(tensor->op); if (rpc_get) { ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context; From bebc9350ecc42a31ad119da1513998386671cf5b Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sun, 30 Aug 2026 09:18:10 +0300 Subject: [PATCH 040/109] common: rename --tensor-read-lazy to --lazy-mode, add -lzm shorthand (#27969) Rename the --tensor-read-lazy CLI argument to --lazy-mode, to match the internal lazy_mode parameter, and add a -lzm shorthand. Sync the READMEs. Assisted-by: pi:llama.cpp/Qwen3.8-27B --- common/arg.cpp | 4 ++-- tools/cli/README.md | 2 +- tools/completion/README.md | 2 +- tools/llama-bench/README.md | 2 +- tools/llama-bench/llama-bench.cpp | 6 +++--- tools/server/README.md | 2 +- 6 files changed, 9 insertions(+), 9 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 4469612cd5b2..79405b59e076 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2729,7 +2729,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } ).set_env("LLAMA_ARG_LOAD_MODE")); add_opt(common_arg( - {"--tensor-read-lazy"}, "MODE", + {"-lzm", "--lazy-mode"}, "MODE", "on-demand reading of certain tensors, for example per-layer embeddings (default: auto)\n" "- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)\n" "- auto: on, but only for tensors larger than 4 GiB\n" @@ -2740,7 +2740,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex else if (value == "off") { params.lazy_mode = LLAMA_LAZY_MODE_OFF; } else { throw std::invalid_argument("invalid value"); } } - ).set_env("LLAMA_ARG_TENSOR_READ_LAZY")); + ).set_env("LLAMA_ARG_LAZY_MODE")); add_opt(common_arg( {"--numa"}, "TYPE", "attempt optimizations that help on some NUMA systems\n" diff --git a/tools/cli/README.md b/tools/cli/README.md index 163ee4fbaf0a..b874d020730d 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -59,7 +59,7 @@ | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | -| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | +| `-lzm, --lazy-mode MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_LAZY_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/completion/README.md b/tools/completion/README.md index 0cd86bac70fc..145be77e31c9 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -142,7 +142,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | -| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | +| `-lzm, --lazy-mode MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_LAZY_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | diff --git a/tools/llama-bench/README.md b/tools/llama-bench/README.md index a1404d2e33e3..8adc56514d10 100644 --- a/tools/llama-bench/README.md +++ b/tools/llama-bench/README.md @@ -67,7 +67,7 @@ test parameters: -nkvo, --no-kv-offload <0|1> (default: 0) -fa, --flash-attn (default: auto) -dev, --device (default: auto) - --tensor-read-lazy (default: auto) + -lzm, --lazy-mode (default: auto) -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode) -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode) -embd, --embeddings <0|1> (default: 0) diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 1b4bbde4d4ff..1fff21f701e2 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -280,7 +280,7 @@ static const char * lazy_mode_str(llama_lazy_mode mode) { case LLAMA_LAZY_MODE_ON: return "on"; default: - GGML_ABORT("invalid tensor read lazy mode"); + GGML_ABORT("invalid lazy mode"); } } @@ -475,7 +475,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); - printf(" --tensor-read-lazy (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str()); + printf(" -lzm, --lazy-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str()); printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); @@ -802,7 +802,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { break; } params.load_mode.insert(params.load_mode.end(), modes.begin(), modes.end()); - } else if (arg == "--tensor-read-lazy") { + } else if (arg == "-lzm" || arg == "--lazy-mode") { if (++i >= argc) { invalid_param = true; break; diff --git a/tools/server/README.md b/tools/server/README.md index 3c2228f34322..c6e907ba9199 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -76,7 +76,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | -| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | +| `-lzm, --lazy-mode MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_LAZY_MODE) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | From 816c7bdc60e688ddb196a2a8be09ef13ce8c5d43 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Mon, 13 Jul 2026 11:09:53 +0000 Subject: [PATCH 041/109] experiment: dequant-once FA scratch for all KV quant types (q4_0/q4_1/q5_0/q5_1/iq4_nl) Evidence branch only - NOT for upstream. Extends the q8_0 dequant-once FA path to every KV-eligible quant type via per-type fused dequant+transpose shaders, plus a prefill-only fa_kv_ok gate for iq4_nl (no native coopmat1 path) and a GGML_VK_FA_DEQUANT env toggle. Correctness: dequant-once == native FA bit-exact for q4_0/q4_1/q5_0/q5_1; iq4_nl matches CPU. Finding: prefill is quant-type independent (all dequant to identical f16 scratch); iq4_nl is a poor KV type (ppl ~2x q4_0 at equal bits). Retained as gating evidence. Assisted-by: Claude Opus 4.8 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 9 +++++++- .../vulkan-shaders/dequant_iq4_nl.comp | 21 +++++++++++++++---- .../vulkan-shaders/dequant_q4_0.comp | 12 +++++++++++ .../vulkan-shaders/dequant_q4_1.comp | 11 ++++++++++ .../vulkan-shaders/dequant_q5_0.comp | 11 ++++++++++ .../vulkan-shaders/dequant_q5_1.comp | 11 ++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 2 +- 7 files changed, 71 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8fbb1359f406..301a603e0650 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -5449,6 +5449,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q4_0], "dequant_q4_0_transpose", dequant_q4_0_transpose_len, dequant_q4_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q4_1], "dequant_q4_1_transpose", dequant_q4_1_transpose_len, dequant_q4_1_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q5_0], "dequant_q5_0_transpose", dequant_q5_0_transpose_len, dequant_q5_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q5_1], "dequant_q5_1_transpose", dequant_q5_1_transpose_len, dequant_q5_1_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q8_0], "dequant_q8_0_transpose", dequant_q8_0_transpose_len, dequant_q8_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); @@ -5466,6 +5470,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ3_S], "dequant_iq3_s", dequant_iq3_s_len, dequant_iq3_s_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ4_XS], "dequant_iq4_xs", dequant_iq4_xs_len, dequant_iq4_xs_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ4_NL], "dequant_iq4_nl", dequant_iq4_nl_len, dequant_iq4_nl_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_IQ4_NL], "dequant_iq4_nl_transpose", dequant_iq4_nl_transpose_len, dequant_iq4_nl_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_MXFP4], "dequant_mxfp4", dequant_mxfp4_len, dequant_mxfp4_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_NVFP4], "dequant_nvfp4", dequant_nvfp4_len, dequant_nvfp4_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); @@ -10930,7 +10935,9 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx }; const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32; const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32; - const bool use_dequant_kv = k_quant && v_quant && neq1 >= 64 && + static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); + const bool fa_dequant_off = fa_dequant_env && fa_dequant_env[0] == '0'; + const bool use_dequant_kv = !fa_dequant_off && k_quant && v_quant && neq1 >= 64 && is_dense_kv_cache(k) && is_dense_kv_cache(v) && (uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && (uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp index 8f7833eab2e7..befcdc1c5100 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp @@ -10,8 +10,6 @@ layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; void main() { const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; - init_iq_shmem(gl_WorkGroupSize); - const uint tid = gl_LocalInvocationID.x % 64; const uint il = tid/32; const uint ir = tid%32; @@ -21,12 +19,27 @@ void main() { } const uint q_idx = 8*il; + +#ifdef DEQUANT_TRANSPOSE + // Fused dequant+transpose for FA quant-KV: source physical is [HS, NH, KV, NS] + // (p.M=HS, p.K=NH, p.stride_a=KV); write to per-head-contiguous dest [HS, KV, NH, NS] so the + // f16 FA reads KV coalesced. An iq4_nl block = 32 consecutive HS elements at fixed (head,kv) -> + // 32 contiguous dest positions (intra-block nibble offsets unchanged). + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + q_idx; +#else const uint b_idx = 1024*i + 32*ir + q_idx; +#endif const float d = float(data_a[ib].d); [[unroll]] for (uint l = 0; l < 8; ++l) { - data_b[b_idx + l + 0] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] & 0xF]); - data_b[b_idx + l + 16] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] >> 4]); + data_b[b_idx + l + 0] = D_TYPE(d * float(kvalues_iq4nl_const[data_a[ib].qs[q_idx + l] & 0xF])); + data_b[b_idx + l + 16] = D_TYPE(d * float(kvalues_iq4nl_const[data_a[ib].qs[q_idx + l] >> 4])); } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp index b92b292135b4..51a6c89a60c8 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp @@ -19,7 +19,19 @@ void main() { } const uint q_idx = 8*il; + +#ifdef DEQUANT_TRANSPOSE + // fused dequant+transpose for FA quant-KV: per-head-contiguous f16 scratch (see dequant_q8_0.comp) + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + q_idx; +#else const uint b_idx = 1024*i + 32*ir + q_idx; +#endif const float d = float(data_a[ib].d); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp index 6b63cbe5833b..76f2d958cbab 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp @@ -18,7 +18,18 @@ void main() { return; } +#ifdef DEQUANT_TRANSPOSE + // fused dequant+transpose for FA quant-KV: per-head-contiguous f16 scratch (see dequant_q8_0.comp) + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + 8*il; +#else const uint b_idx = 1024*i + 32*ir + 8*il; +#endif const float d = float(data_a[ib].d); const float m = float(data_a[ib].m); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp index f1b0bac87271..1402fa4292d8 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp @@ -18,7 +18,18 @@ void main() { return; } +#ifdef DEQUANT_TRANSPOSE + // fused dequant+transpose for FA quant-KV: per-head-contiguous f16 scratch (see dequant_q8_0.comp) + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + 8*il; +#else const uint b_idx = 1024*i + 32*ir + 8*il; +#endif const float d = float(data_a[ib].d); const uint qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp index c495b31f1754..1fd6e2552af7 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp @@ -18,7 +18,18 @@ void main() { return; } +#ifdef DEQUANT_TRANSPOSE + // fused dequant+transpose for FA quant-KV: per-head-contiguous f16 scratch (see dequant_q8_0.comp) + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + 8*il; +#else const uint b_idx = 1024*i + 32*ir + 8*il; +#endif const float d = float(data_a[ib].d); const float m = float(data_a[ib].m); 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 d375c2d12771..a26e8b7c0b79 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -781,7 +781,7 @@ void process_shaders() { string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}})); } // Fused dequant+transpose variant for FA quant-KV (per-head-contiguous f16 scratch). - if (tname == "q8_0") { + if (tname == "q8_0" || tname == "iq4_nl" || tname == "q4_0" || tname == "q4_1" || tname == "q5_0" || tname == "q5_1") { string_to_spv("dequant_" + tname + "_transpose", "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}, {"DEQUANT_TRANSPOSE", "1"}})); } From 0aac213ef7b46f075f1b68655be514046e4174ce Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 14 Jul 2026 14:09:02 +0000 Subject: [PATCH 042/109] vulkan : gate FA dequant-once scratch on device-local capacity The dequant-once path materializes a per-layer f16 K/V scratch (~2 KiB/token). On discrete devices that resident footprint can push the working set past free VRAM, at which point the driver silently pages device-local memory: measured ~15x prefill regression on an 8 GB card (RTX 3070, driver 591.86) at long context, with no error reported. Integrated/UMA devices have no separate device pool to overflow and are unaffected. Gate the path on this process's device-local usage against the physical heap size, keeping a conservative reserve for memory not observable in-process. heapBudget is deliberately not used as the signal: ggml_backend_vk_get_device_memory computes heapBudget - heapUsage in unsigned arithmetic, which wraps to a huge value exactly when the device is oversubscribed. The allocation cannot gate itself either - on WDDM vkAllocateMemory only fails at roughly physical heap size, which is above the free-VRAM level where paging begins, so a successful allocation is not evidence of a resident fit. Also fix the scratch size check: K and V are bound as one storage buffer, so their sum must fit maxStorageBufferRange, not each half independently. GGML_VK_FA_DEQUANT=0 forces the path off and =1 skips the capacity check; GGML_VK_FA_DEQUANT_RESERVE_MB overrides the reserve. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 87 ++++++++++++++++++++++++++-- 1 file changed, 83 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 301a603e0650..7e339dc4bb8c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2389,6 +2389,10 @@ struct ggml_backend_vk_context { ggml_vk_garbage_collector gc; size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k, prealloc_size_add_rms_partials, prealloc_size_add_rms_partials_offset; vk_buffer prealloc_x, prealloc_y, prealloc_split_k, prealloc_add_rms_partials, sync_staging; + // memoized capacity decision for the FA dequant-once scratch, see ggml_vk_fa_dequant_scratch_fits + uint64_t fa_dequant_gate_sz; + bool fa_dequant_gate_fits; + bool fa_dequant_gate_logged; vk::Fence fence, almost_ready_fence; bool submit_pending {}; bool almost_ready_fence_pending {}; @@ -7695,6 +7699,9 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { ctx->prealloc_size_x = 0; ctx->prealloc_size_y = 0; ctx->prealloc_size_split_k = 0; + ctx->fa_dequant_gate_sz = 0; + ctx->fa_dequant_gate_fits = false; + ctx->fa_dequant_gate_logged = false; // Fixed size of 1KB, for deterministic behavior ctx->prealloc_size_add_rms_partials = 1024; @@ -10867,6 +10874,69 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co return supported; } +// Capacity gate for the dequant-once f16 K/V scratch. On discrete devices the scratch can push the +// working set past free VRAM, and the driver then silently pages device-local memory (~15x prefill +// regression measured on an 8 GB card at long context). UMA has no separate pool to overflow. +// +// Gates our own device-local usage against the physical heap size, less a reserve for memory not +// visible in-process. heapBudget is deliberately not the signal: ggml_backend_vk_get_device_memory +// computes heapBudget - heapUsage unsigned, which wraps when the device is oversubscribed. Nor can +// the allocation gate itself: on WDDM vkAllocateMemory only fails near physical heap size, above +// the free-VRAM level where paging starts. The reserve is necessarily conservative because other +// processes' VRAM use is invisible to us. GGML_VK_FA_DEQUANT=0/1 forces the path off/on; +// GGML_VK_FA_DEQUANT_RESERVE_MB overrides the reserve. +static bool ggml_vk_fa_dequant_scratch_fits(ggml_backend_vk_context * ctx, uint64_t scratch_sz) { + const vk_device& device = ctx->device; + + if (device->uma) { + return true; + } + + // Decided once per scratch size, so the decision cannot flip between layers as usage grows. + if (ctx->fa_dequant_gate_sz == scratch_sz) { + return ctx->fa_dequant_gate_fits; + } + + static const uint64_t reserve = [] { + const char * env = getenv("GGML_VK_FA_DEQUANT_RESERVE_MB"); + return (uint64_t)(env ? atoi(env) : 1024) * 1024 * 1024; + }(); + + bool fits = false; + + // Without VK_EXT_memory_budget our usage is unknowable, so leave the path disabled. + if (vk_instance.device_supports_membudget[device->idx]) { + vk::PhysicalDeviceMemoryBudgetPropertiesEXT budgetprops; + vk::PhysicalDeviceMemoryProperties2 memprops = {}; + memprops.pNext = &budgetprops; + device->physical_device.getMemoryProperties2(&memprops); + + uint64_t heap_size = 0; + uint64_t heap_used = 0; + for (uint32_t i = 0; i < memprops.memoryProperties.memoryHeapCount; ++i) { + const vk::MemoryHeap & heap = memprops.memoryProperties.memoryHeaps[i]; + if (heap.flags & vk::MemoryHeapFlagBits::eDeviceLocal) { + heap_size += heap.size; + heap_used += budgetprops.heapUsage[i]; + } + } + // heap_used already covers scratch allocated on a previous ubatch, so counting scratch_sz + // in full is conservative by up to the current scratch size. + fits = heap_size > reserve && heap_used + scratch_sz + reserve <= heap_size; + } + + if (!fits && !ctx->fa_dequant_gate_logged) { + ctx->fa_dequant_gate_logged = true; + GGML_LOG_INFO("ggml_vulkan: flash attention dequant-once disabled: %llu MiB K/V scratch does not fit " + "device-local memory with a %llu MiB reserve. Set GGML_VK_FA_DEQUANT=1 to force it on.\n", + (unsigned long long)(scratch_sz >> 20), (unsigned long long)(reserve >> 20)); + } + + ctx->fa_dequant_gate_sz = scratch_sz; + ctx->fa_dequant_gate_fits = fits; + return fits; +} + static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; @@ -10926,7 +10996,14 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16; - // dequant K/V once into an f16 scratch, reordered KV layout so FA can read without a stride + // For prefill with quantized K/V, dequantize+transpose K/V once into a per-head-contiguous + // f16 scratch and run the f16 FA path, instead of the coopmat1 shader re-dequantizing the + // whole KV inside every Q-workgroup. The KV-cache view reaching FA is [0,2,1,3]-permuted but + // dense, so we require dense allocation (not ggml_is_contiguous) and block-contiguous dim0, + // and only engage where a fused dequant-transpose shader exists. Prefill only + // (n_rows >= 64); measured neutral at shallow depth and up to ~2x at long context. + // The scratch is bound as a single storage buffer holding K and V back to back, so the SUM of + // the two must fit maxStorageBufferRange, not each half independently. auto is_dense_kv_cache = [](const ggml_tensor * t) { return t->nb[0] == ggml_type_size(t->type) && t->nb[2] == ggml_row_size(t->type, t->ne[0]) && @@ -10935,19 +11012,21 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx }; const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32; const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32; + const uint64_t kv_f16_sz = ((uint64_t)ggml_nelements(k) + (uint64_t)ggml_nelements(v)) * sizeof(ggml_fp16_t); static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); const bool fa_dequant_off = fa_dequant_env && fa_dequant_env[0] == '0'; + const bool fa_dequant_on = fa_dequant_env && fa_dequant_env[0] == '1'; const bool use_dequant_kv = !fa_dequant_off && k_quant && v_quant && neq1 >= 64 && is_dense_kv_cache(k) && is_dense_kv_cache(v) && - (uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && - (uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && + kv_f16_sz <= ctx->device->properties.limits.maxStorageBufferRange && ctx->device->pipeline_dequant_transpose[k->type] != nullptr && ctx->device->pipeline_dequant_transpose[v->type] != nullptr && // coopmat2 path does not benefit from the f16 scratch !ctx->device->coopmat2 && // Intel Xe1 regresses, see PR 25494 (ctx->device->vendor_id != VK_VENDOR_ID_INTEL || - (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)); + (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)) && + (fa_dequant_on || ggml_vk_fa_dequant_scratch_fits(ctx, kv_f16_sz)); 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; From c8230e7826ba531cbdba57eb0f36279b6df4e056 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 28 Jul 2026 10:21:44 +0000 Subject: [PATCH 043/109] vulkan: contiguize strided f16 KV for FA prefill (GGML_VK_FA_KV_CONTIG, env-gated) The KV-cache view reaching FA has head-interleaved rows ([HS, NH, KV] physically), and the cm1 shader's direct-from-global coopMatLoads run ~2x slower on that stride than on per-head-contiguous K/V: a 16x16 tile touches 16 distant cache lines instead of 4. Measured on Strix Halo (RADV gfx1151), hd128/GQA8/kv10240/nb2048 f16: 29.8ms contiguous vs 63.1ms dense-permuted (the model layout; matches the in-model 59.9ms from the perf logger, where FLASH_ATTN_EXT was 72.6% of the graph at pp2048@d8192). GGML_VK_FA_KV_CONTIG=1 extends the dequant-once FA scratch to f16 K/V: dequant_f16_transpose.comp is a pure strided copy ([HS,NH,KV] -> [HS,KV,NH], same push-constant ABI and dispatch as the quant transpose shaders), engaged only when the rows are actually strided, prefill only (neq1 >= 64). FA op 63.1 -> 30.5ms (2.07x) incl. copy cost. Model-level (Qwen3-Coder-30B Q6_K_XL, ub2048, f16 KV, r=3, vs ROCm 571d0d5 nowmma): pp8192 877 -> 1199 t/s (ROCm 1216, was -28% now parity); pp2048@d4096/8192/16384: 776/513/300 -> 1120/850/580. Shallow prefill unchanged-to-better (pp2048 1542 -> 1633). Not the fix: shmem staging on AMD (loses on occupancy, 29.8 -> 54.4ms contiguous), bigger-tile/GQA-packed streaming (1-KV-head L2-resident probe runs identical -> kernel is issue-bound, not bandwidth-bound). test-backend-ops -o FLASH_ATTN_EXT green with the flag off and on (pre-existing iq4_nl+sinks failures unchanged). Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 14 +++++++- .../vulkan-shaders/dequant_f16_transpose.comp | 32 +++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 4 +++ 3 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 7e339dc4bb8c..cb4ad08f0d93 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -5459,6 +5459,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q5_1], "dequant_q5_1_transpose", dequant_q5_1_transpose_len, dequant_q5_1_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q8_0], "dequant_q8_0_transpose", dequant_q8_0_transpose_len, dequant_q8_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_F16], "dequant_f16_transpose", dequant_f16_transpose_len, dequant_f16_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 8, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ2_0], "dequant_tq2_0", dequant_tq2_0_len, dequant_tq2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); @@ -11016,7 +11017,18 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); const bool fa_dequant_off = fa_dequant_env && fa_dequant_env[0] == '0'; const bool fa_dequant_on = fa_dequant_env && fa_dequant_env[0] == '1'; - const bool use_dequant_kv = !fa_dequant_off && k_quant && v_quant && neq1 >= 64 && + // EXPERIMENT (GGML_VK_FA_KV_CONTIG=1): run the same contiguize pass for f16 K/V. The + // KV-cache view reaching FA is head-interleaved ([HS, NH, KV] physically), and the cm1 + // direct-from-global coopMatLoads run ~2-5x slower on those strided rows than on + // per-head-contiguous K/V. Copy K/V once into the scratch instead (dequant_f16_transpose + // is a pure strided copy). Only engages when the rows are actually strided. + static const char * fa_kv_contig_env = getenv("GGML_VK_FA_KV_CONTIG"); + const bool fa_kv_contig = fa_kv_contig_env && fa_kv_contig_env[0] == '1'; + const bool kv_f16_strided = k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 && + (k->nb[1] != (uint64_t)HSK * sizeof(ggml_fp16_t) || + v->nb[1] != (uint64_t)HSV * sizeof(ggml_fp16_t)) && + (HSK % 8) == 0 && (HSV % 8) == 0; + const bool use_dequant_kv = !fa_dequant_off && ((k_quant && v_quant) || (fa_kv_contig && kv_f16_strided)) && neq1 >= 64 && is_dense_kv_cache(k) && is_dense_kv_cache(v) && kv_f16_sz <= ctx->device->properties.limits.maxStorageBufferRange && ctx->device->pipeline_dequant_transpose[k->type] != nullptr && diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose.comp new file mode 100644 index 000000000000..dda53d749acc --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose.comp @@ -0,0 +1,32 @@ +#version 450 + +#include "dequant_head.glsl" + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {f16vec4 data_a[];}; +layout (binding = 1) writeonly buffer D {f16vec4 data_b[];}; + +// Strided-copy counterpart of the fused dequant+transpose shaders for FA f16 KV: +// source physical is [HS, NH, KV, NS] (p.M=HS, p.K=NH, p.stride_a=KV); write to +// per-head-contiguous dest [HS, KV, NH, NS] so the f16 FA reads KV coalesced. +// HS stays innermost in both layouts, so each invocation moves 8 HS-contiguous +// elements (two f16vec4) requiring HS % 8 == 0 (enforced by the host gate). +void main() { + const uint i = gl_GlobalInvocationID.x; + const uint e0 = i * 8; + if (e0 >= p.nel) { + return; + } + + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint dst = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH); + + data_b[dst / 4 ] = data_a[e0 / 4 ]; + data_b[dst / 4 + 1] = data_a[e0 / 4 + 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 a26e8b7c0b79..752380f0b7af 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -784,6 +784,10 @@ void process_shaders() { if (tname == "q8_0" || tname == "iq4_nl" || tname == "q4_0" || tname == "q4_1" || tname == "q5_0" || tname == "q5_1") { string_to_spv("dequant_" + tname + "_transpose", "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}, {"DEQUANT_TRANSPOSE", "1"}})); } + // Strided-copy counterpart for f16 KV (contiguize the head-interleaved cache layout). + if (tname == "f16") { + string_to_spv("dequant_f16_transpose", "dequant_f16_transpose.comp", {}); + } shader = (tname == "f32" || tname == "f16" || tname == "bf16") ? "get_rows.comp" : "get_rows_quant.comp"; From 40aaa496f9cc78b691078d70655c7ab7d0694f13 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 28 Jul 2026 10:21:56 +0000 Subject: [PATCH 044/109] tests: dense-permuted K/V option + Strix FA prefill perf/probe cases test_flash_attn_ext always built K/V as sparse views (physical dim1 doubled, half viewed), which can never satisfy the contiguize path's ggml_is_contiguously_allocated gate - so no permuted correctness case exercised it. Add a kv_view parameter (default true = unchanged) and dense-permuted eval cases matching the real KV-cache layout, including ALiBi and logit-softcap variants; all pass vs CPU with GGML_VK_FA_KV_CONTIG=1. Perf additions: Qwen3-Coder-30B prefill-at-depth shapes (hd128, 4 KV heads, GQA 8, kv up to 10240, nb 512/2048), the dense-permuted variant (model layout), and the probe set used to establish that the contiguous cm1 kernel is issue-bound: 32-distinct-KV-head MALL-spill (flat), 1-KV-head L2-resident (flat), mask=0 (-5.5%), f16 acc (-3%). Co-Authored-By: Claude Fable 5 --- tests/test-backend-ops.cpp | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4a7a0623174c..ba3764045bf5 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7091,7 +7091,7 @@ struct test_flash_attn_ext : public test_case { const ggml_type type_K; const ggml_type type_V; std::array permute; - const bool kv_view; // create K/V as views of a larger buffer (like a KV cache) + const bool kv_view; // create K/V as views of a larger buffer (like a KV cache); false = dense permuted like the model KV cache const bool v_is_view_of_k; std::string vars() override { @@ -9958,6 +9958,12 @@ static std::vector> make_test_cases_eval() { } } + // dense-permuted K/V (model KV-cache layout, engages the f16 contiguize path at nb>=64) + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 1024, 128, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(96, 96, 8, {4, 1}, 512, 80, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {8, 1}, 512, 75, true, false, 8.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 512, 96, true, false, 0, 30.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + // mixed quant and Q1_0 test cases test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); @@ -10376,6 +10382,23 @@ static std::vector> make_test_cases_perf() { // Qwen3-VL-8B https://github.com/ggml-org/llama.cpp/issues/17012 test_cases.emplace_back(new test_flash_attn_ext(72, 72, 16, {1, 1}, 5776, 5776, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // Qwen3-Coder-30B-A3B prefill at depth: hd128, 4 KV heads, GQA 8, ub 2048 + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 2048, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 6144, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 10240, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 10240, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // MALL-spill probe: same FLOPs, 32 distinct KV heads (no GQA) -> K/V footprint 8x (168MB > 32MB MALL) + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 32, {1, 1}, 10240, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // KV-cache layout probe: same shape, K/V strided token-major like the real cache (heads interleaved) + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 10240, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3})); + // Same, dense-permuted (exact model KV-cache layout; eligible for the f16 contiguize path) + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 10240, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + // L2-residence probe: 1 KV head x GQA 32 (K/V 5.2MB fits L2) - distinguishes cache-BW-bound from issue-bound + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, {32, 1}, 10240, 2048, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // cost-partition probes: no mask; f16 accumulate + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 10240, 2048, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 10240, 2048, true, false, 0, 0, GGML_PREC_DEFAULT, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)); From e7a0e5808afdde57f4d6b1274ffe4acd0d4f48f6 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 28 Jul 2026 13:25:49 +0000 Subject: [PATCH 045/109] vulkan: enable the f16 KV contiguize pass by default (GGML_VK_FA_KV_CONTIG=0 opts out) Flip 442d7df from opt-in to default-on, matching the quant dequant-once path (GGML_VK_FA_DEQUANT) convention. The pass still self-gates: f16 K/V only, prefill only (neq1 >= 64), only when rows are actually strided, dense allocation, and the shared scratch-capacity check. Validated on Strix Halo (RADV gfx1151): FLASH_ATTN_EXT suite green with default env (dense-permuted cases exercise the pass) and with the opt-out; model-level pp2048@d8192 with no FA env matches the explicit GGML_VK_FA_KV_CONTIG=1 validation run (846.6 vs 847.9 t/s). Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index cb4ad08f0d93..3a02d44e85a5 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11017,13 +11017,14 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); const bool fa_dequant_off = fa_dequant_env && fa_dequant_env[0] == '0'; const bool fa_dequant_on = fa_dequant_env && fa_dequant_env[0] == '1'; - // EXPERIMENT (GGML_VK_FA_KV_CONTIG=1): run the same contiguize pass for f16 K/V. The - // KV-cache view reaching FA is head-interleaved ([HS, NH, KV] physically), and the cm1 + // Contiguize pass for f16 K/V (GGML_VK_FA_KV_CONTIG=0 opts out). The KV-cache view + // reaching FA is head-interleaved ([HS, NH, KV] physically), and the cm1 // direct-from-global coopMatLoads run ~2-5x slower on those strided rows than on // per-head-contiguous K/V. Copy K/V once into the scratch instead (dequant_f16_transpose - // is a pure strided copy). Only engages when the rows are actually strided. + // is a pure strided copy). Only engages when the rows are actually strided, and shares + // the quant path's prefill/allocation/scratch-capacity gates below. static const char * fa_kv_contig_env = getenv("GGML_VK_FA_KV_CONTIG"); - const bool fa_kv_contig = fa_kv_contig_env && fa_kv_contig_env[0] == '1'; + const bool fa_kv_contig = !(fa_kv_contig_env && fa_kv_contig_env[0] == '0'); const bool kv_f16_strided = k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 && (k->nb[1] != (uint64_t)HSK * sizeof(ggml_fp16_t) || v->nb[1] != (uint64_t)HSV * sizeof(ggml_fp16_t)) && From bf54e452d2aa9a51873b17adc545ad6c48e226a3 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 28 Jul 2026 23:40:27 +0000 Subject: [PATCH 046/109] vulkan: single source of truth for native FA K/V types + non-native hard-gate Rebased adaptation of the iq4_nl routing fix: upstream 8161641 made iq4_nl a native FA type, so the original motivation (iq4_nl had no native shader and silently read garbage outside the dequant-once path) no longer applies to any currently-admitted type. Keep the machinery as hardening: ggml_vk_fa_kv_native() is the one list, supports_op mirrors every hard condition of the dispatch-time dequant gate for any future non-native type, and dispatch asserts the invariant instead of falling back to a garbage-reading shader. Native list synced with upstream (iq4_nl in, q1_0 out to match current admission). Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 80 +++++++++++++++++++++------- 1 file changed, 60 insertions(+), 20 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 3a02d44e85a5..ea9023040561 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -10875,6 +10875,27 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co return supported; } +// K/V types the FA shaders can read directly (scalar/coopmat1 select the dequant code via the +// FaTypeK/FaTypeV spec constants; a type outside this list silently reads garbage). Types that +// are FA-supported but not listed here (iq4_nl) are only correct through the dequant-once +// scratch path, so supports_op and the dispatch-time gate must agree on when that path runs. +static bool ggml_vk_fa_kv_native(ggml_type t, bool coopmat2) { + switch (t) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_IQ4_NL: // native FA support since upstream 8161641 + return true; + default: + return false; + } +} + // Capacity gate for the dequant-once f16 K/V scratch. On discrete devices the scratch can push the // working set past free VRAM, and the driver then silently pages device-local memory (~15x prefill // regression measured on an 8 GB card at long context). UMA has no separate pool to overflow. @@ -11029,17 +11050,27 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx (k->nb[1] != (uint64_t)HSK * sizeof(ggml_fp16_t) || v->nb[1] != (uint64_t)HSV * sizeof(ggml_fp16_t)) && (HSK % 8) == 0 && (HSV % 8) == 0; - const bool use_dequant_kv = !fa_dequant_off && ((k_quant && v_quant) || (fa_kv_contig && kv_f16_strided)) && neq1 >= 64 && + // A K/V type the FA shaders cannot read directly (iq4_nl) is only correct through the + // dequant path. supports_op only admits such types when the hard conditions below hold, + // and the VRAM heuristic must not veto them (there is no fallback), so force the path on. + const bool kv_needs_dequant = !ggml_vk_fa_kv_native(k->type, ctx->device->coopmat2) || + !ggml_vk_fa_kv_native(v->type, ctx->device->coopmat2); + const bool use_dequant_kv = !fa_dequant_off && + ((k_quant && v_quant) || kv_needs_dequant || (fa_kv_contig && kv_f16_strided)) && neq1 >= 64 && is_dense_kv_cache(k) && is_dense_kv_cache(v) && kv_f16_sz <= ctx->device->properties.limits.maxStorageBufferRange && ctx->device->pipeline_dequant_transpose[k->type] != nullptr && ctx->device->pipeline_dequant_transpose[v->type] != nullptr && - // coopmat2 path does not benefit from the f16 scratch - !ctx->device->coopmat2 && + // coopmat2 reads its native types directly; non-native still needs the scratch + (kv_needs_dequant || !ctx->device->coopmat2) && // Intel Xe1 regresses, see PR 25494 - (ctx->device->vendor_id != VK_VENDOR_ID_INTEL || + (kv_needs_dequant || + ctx->device->vendor_id != VK_VENDOR_ID_INTEL || (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)) && - (fa_dequant_on || ggml_vk_fa_dequant_scratch_fits(ctx, kv_f16_sz)); + (fa_dequant_on || kv_needs_dequant || ggml_vk_fa_dequant_scratch_fits(ctx, kv_f16_sz)); + // If this fires, supports_op admitted a non-native K/V type the gate then rejected; the + // native shader would return garbage rather than fail, so abort instead. + GGML_ASSERT(use_dequant_kv || !kv_needs_dequant); 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; @@ -18592,25 +18623,34 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) { return false; } - auto fa_kv_ok = [](ggml_type t) { - switch (t) { - case GGML_TYPE_F32: - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q4_0: - case GGML_TYPE_IQ4_NL: - return true; - default: - return false; - } + auto fa_kv_ok = [&](ggml_type t) { + // ggml_vk_fa_kv_native is the single source of truth; on this base every + // admitted type is native, and the non-native hard-gate below is dormant + // hardening for any future type routed through the dequant-once scratch. + return ggml_vk_fa_kv_native(t, coopmat2); }; if (!fa_kv_ok(op->src[1]->type) || !fa_kv_ok(op->src[2]->type)) { return false; } + if (!ggml_vk_fa_kv_native(op->src[1]->type, coopmat2) || !ggml_vk_fa_kv_native(op->src[2]->type, coopmat2)) { + // Only correct through the dequant-once scratch path; admit only when every + // hard condition of the dispatch-time gate holds, so dispatch can never fall + // back to the native shader (it reads garbage for these types, not an error). + const ggml_tensor * k = op->src[1]; + const ggml_tensor * v = op->src[2]; + static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); + const bool fa_dequant_off = fa_dequant_env && fa_dequant_env[0] == '0'; + const uint64_t kv_f16_sz = ((uint64_t)ggml_nelements(k) + (uint64_t)ggml_nelements(v)) * sizeof(ggml_fp16_t); + if (fa_dequant_off || + op->src[0]->ne[1] < 64 || + device->pipeline_dequant_transpose[k->type] == nullptr || + device->pipeline_dequant_transpose[v->type] == nullptr || + k->nb[0] != ggml_type_size(k->type) || v->nb[0] != ggml_type_size(v->type) || + !ggml_is_contiguously_allocated(k) || !ggml_is_contiguously_allocated(v) || + kv_f16_sz > device->properties.limits.maxStorageBufferRange) { + return false; + } + } if ((op->src[1]->type == GGML_TYPE_BF16) != (op->src[2]->type == GGML_TYPE_BF16)) { return false; } From 66e73489fc71228f4e56a2a6c7a50b4f2549188f Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 30 Jul 2026 01:21:49 +0000 Subject: [PATCH 047/109] vulkan: hoist the coopmat1 FA P-fragment load out of the hsv_tile loop The GEMM2 P load depends only on bc_chunk, but sat inside the hsv_tile loop, so all four fragments were re-read from shared memory once per tile (2 tiles at HSV=128, 4 at HSV=256). Load them once into a coopmat array before the loop. Psh is not written again until the next KV block, so the fragments stay valid across the barriers inside it. Measured on gfx1151 (RADV), Qwen3-Coder-30B-A3B, pp2048 f16 KV: d8192 +6.9%, d16384 +8.1%, d32768 +9.2%. VGPRs and LDS unchanged, 0 spilled. FLASH_ATTN_EXT suite 5105/5105. Assisted-by: Claude Fable 5 --- .../ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) 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 057ed739aa8d..d30816b515f8 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -423,6 +423,13 @@ void main() { const uint num_hsv_tiles = (HSV + MatBc * row_split - 1) / (MatBc * row_split); // round up + // Psh is not written again until the next KV block, so the P fragments are the same + // for every hsv_tile. Load them once instead of re-reading LDS per tile. + coopmat PMat[Bc / MatBc]; + [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { + coopMatLoad(PMat[bc_chunk], Psh, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); + } + // 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; @@ -476,8 +483,6 @@ void main() { if (hsv_offset < HSV_pad) { [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { - coopMatLoad(KMat, Psh, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); - if (SHMEM_STAGING == 0) { if (!USE_DECODE_V && !KV_bounds_check) { // F16/BF16 values can be loaded directly from global memory @@ -493,7 +498,7 @@ void main() { coopMatLoad(QMat, kvsh, v_tile_offset, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); } - PVMat = coopMatMulAdd(KMat, QMat, PVMat); + PVMat = coopMatMulAdd(PMat[bc_chunk], QMat, PVMat); } // Store PVMat to pvsh and load into Of From 6a6896cdd78234a6dbe90f45c794be0c9faf74a6 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 28 Jul 2026 23:51:32 +0000 Subject: [PATCH 048/109] tests: dense-permuted iq4_nl FA cases (dequant-once route coverage) iq4_nl has no native Vulkan FA shader; the dequant-once scratch path is its only route (1abdd92). The sweep's iq4_nl cases are all sparse-view, which that path correctly rejects, so iq4_nl had zero passing FA coverage. Add model-layout (dense [0,2,1,3]-permuted) cases at prefill batch size: iq4_nl/iq4_nl with sinks off and on, hd72+GQA with sinks, and mixed K=iq4_nl/V=f16. Validated on RADV gfx1151 at 146fb73: FLASH_ATTN_EXT 4765/4765 with default env and with GGML_VK_FA_KV_CONTIG=0/1; 4761/4761 with GGML_VK_FA_DEQUANT=0 (new cases correctly report unsupported); full test-backend-ops suite 15538/15538. Co-Authored-By: Claude Fable 5 --- tests/test-backend-ops.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index ba3764045bf5..437073855e4e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9964,6 +9964,13 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {8, 1}, 512, 75, true, false, 8.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 512, 96, true, false, 0, 30.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + // dense-permuted iq4_nl K/V at prefill batch sizes: iq4_nl has no native FA shader, so these + // exercise the only supported route (the dequant-once path), incl. sinks and mixed-with-f16 + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(72, 72, 4, {4, 1}, 113, 75, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_IQ4_NL, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + // mixed quant and Q1_0 test cases test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); From 306eecf3ef25a494301c377071a9c4bd85398bc9 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 30 Jul 2026 01:23:34 +0000 Subject: [PATCH 049/109] vulkan: store coopmat1 FA Psh query-major so the GEMM2 A load vectorizes Psh held P as [kv][query]. The GEMM2 UseA load therefore had to request ColumnMajor, and RADV only attaches an alignment hint to the internal column-major case, which for UseA means the RowMajor request. The load was emitted as 16 separate 16-bit shared reads per fragment. Store P as [query][kv] instead and request RowMajor. The producer now writes four scalar components rather than one vec4; those go to disjoint bytes, so there is no read-modify-write and no partially written vec4. Also updates the host shared-memory estimator, which mirrored the old stride and would otherwise disagree with the shader. Perf-neutral on gfx1151 (within 1% at hd128 and hd256), but it shrinks Psh: LDS 16384 -> 15360 B and code size 13896 -> 13596 at hd128, VGPRs unchanged at 96 with 0 spilled. Kept because it is free and removes a scalar shared-memory access pattern. FLASH_ATTN_EXT suite 5105/5105. Assisted-by: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 4 ++-- .../vulkan-shaders/flash_attn_cm1.comp | 17 ++++++++++++----- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8fbb1359f406..c3036751295f 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -10837,8 +10837,8 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co const uint32_t qstride = hsk_pad / 4 + 2; const uint32_t Qf = Br * qstride * f16vec4; - const uint32_t psh_stride = Br / 4 + 2; - const uint32_t Psh = Bc * psh_stride * f16vec4; + const uint32_t psh_stride = Bc / 4 + 2; + const uint32_t Psh = Br * psh_stride * f16vec4; const uint32_t sfshstride = (hsk <= 128) ? (Br + 8) : Br; const uint32_t sfsh = Bc * sfshstride * acctype; 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 d30816b515f8..c812657973e2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -44,8 +44,11 @@ shared float tmpsh[row_split]; const uint32_t qstride = HSK_pad / 4 + 2; shared FLOAT_TYPEV4 Qf[Br * qstride]; -const uint psh_stride = Br / 4 + 2; -shared FLOAT_TYPEV4 Psh[Bc * psh_stride]; +// P is stored query-major with KV contiguous so the GEMM2 UseA load can be RowMajor: +// RADV flips the requested layout for UseA, and only the resulting internal column-major +// case gets an alignment hint, which is what lets the 16 element loads vectorize. +const uint psh_stride = Bc / 4 + 2; +shared FLOAT_TYPEV4 Psh[Br * psh_stride]; // Avoid padding for hsk==256 to make it fit in 48KB shmem. const uint32_t sfshstride = (HSK <= 128) ? (Br / 4 + 2) : Br / 4; @@ -382,15 +385,19 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { const uint row = tile_row(r); + const uint pcol_vec = col / 4; + const uint pcol_comp = col % 4; if (KV_bounds_check && j * Bc + col >= KV) { - Psh[col * psh_stride + row / 4] = FLOAT_TYPEV4(0.0f); + [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { + Psh[(row + vec_idx) * psh_stride + pcol_vec][pcol_comp] = FLOAT_TYPE(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)); [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { Lf[r + vec_idx] += Pf[vec_idx]; + Psh[(row + vec_idx) * psh_stride + pcol_vec][pcol_comp] = Pf[vec_idx]; } - Psh[col * psh_stride + row / 4] = Pf; } } } @@ -427,7 +434,7 @@ void main() { // for every hsv_tile. Load them once instead of re-reading LDS per tile. coopmat PMat[Bc / MatBc]; [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { - coopMatLoad(PMat[bc_chunk], Psh, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad(PMat[bc_chunk], Psh, bc_chunk * (MatBc / 4), psh_stride, gl_CooperativeMatrixLayoutRowMajor); } // Each subgroup handles HSV/4 columns From 48407ef0f148b0805efb58aa695e018aa0db0ca3 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Fri, 31 Jul 2026 01:25:46 +0000 Subject: [PATCH 050/109] vulkan: bill the FA K/V contiguize pass on its own perf-logger line The contiguize/dequant pass is dispatched inside the FLASH_ATTN_EXT node handler, and the perf logger writes one timestamp per graph node, so its cost was charged to FLASH_ATTN_EXT with no way to separate the two. That made the copy invisible and the kernel look correspondingly slower. Adds a sub-node timestamp: a handler can close an interval mid-node and have it logged under its own name. Measured on Coder-30B UD-Q6_K_XL, pp2048/ub2048 at d32768, f16 KV: graph total 4612.3 ms FLASH_ATTN_EXT 3490.6 ms 75.68% of graph FA_KV_CONTIGUIZE 33.1 ms 0.72% of graph (0.95% of FA) so the copy is under 1% of the graph and FA is 75.7% of it at that depth. Bench throughput is unchanged with the instrumentation compiled in (441.09 vs 442.14 t/s), since the marks are only emitted when the logger is on in per-op mode. Two latent bugs in the query-pool handling fall out of this and are fixed here: - The pool is created with n_nodes+100 slots but only the first n_nodes+1 were reset each graph, so anything using the headroom would read stale results. - The results buffer was sized n_nodes+1 while getQueryPoolResults was asked for query_idx entries. Equal today, but it is an overflow waiting for the first caller that writes an extra timestamp. Sub-op intervals log no flops, since they move bytes rather than doing math; attributing the node's flop count to them would corrupt the GFLOPS column for both halves. --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 50 +++++++++++++++++++++++++--- 1 file changed, 46 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index ea9023040561..528c112645a7 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2356,6 +2356,13 @@ class vk_perf_logger { timings[name].push_back(time); } + // Log a sub-node interval under a caller-supplied name. Used for work a node's handler + // dispatches before the op itself (e.g. the FA K/V contiguize/dequant pass), which would + // otherwise be billed to the op and invisible. No flops: these move bytes, not math. + void log_timing_named(const char *name, uint64_t time) { + timings[std::string(name)].push_back(time); + } + void log_timing(const std::vector &nodes, const std::vector &names, uint64_t time) { uint64_t total_flops = 0; std::string name; @@ -2450,6 +2457,9 @@ struct ggml_backend_vk_context { std::vector query_fusion_node_count; std::vector query_nodes; std::vector query_node_idx; + // non-null => this query slot closes a sub-node interval logged under this literal name, + // not a graph node. See ggml_vk_perf_mark_subop. + std::vector query_sub_names; int32_t num_queries {}; int32_t query_idx {}; }; @@ -10959,6 +10969,22 @@ static bool ggml_vk_fa_dequant_scratch_fits(ggml_backend_vk_context * ctx, uint6 return fits; } +// Close a timestamp interval mid-node so work a handler dispatches before its op is billed +// separately instead of being folded into the op's own time. `name` must be a string literal +// (stored by pointer). No-op unless the perf logger is on in per-op mode. +static void ggml_vk_perf_mark_subop(ggml_backend_vk_context * ctx, vk_context& subctx, const char * name) { + if (!vk_perf_logger_enabled || vk_perf_logger_concurrent || ctx->query_pool == VK_NULL_HANDLE) { + return; + } + if (ctx->query_idx >= (int)ctx->num_queries) { + return; // pool headroom exhausted; drop the mark rather than overflow + } + ctx->query_nodes[ctx->query_idx] = nullptr; + ctx->query_fusion_names[ctx->query_idx] = nullptr; + ctx->query_sub_names[ctx->query_idx] = name; + subctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++); +} + static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; @@ -11270,6 +11296,11 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx ggml_vk_sync_buffers(ctx, subctx); k_buf = k_dst; v_buf = v_dst; + // Bill the K/V contiguize/dequant pass on its own line. Without this it is charged to + // FLASH_ATTN_EXT, which makes the copy invisible and the kernel look slower than it is. + ggml_vk_perf_mark_subop(ctx, subctx, kv_needs_dequant || (k_quant && v_quant) + ? "FA_KV_DEQUANT (sub-op)" + : "FA_KV_CONTIGUIZE (sub-op)"); } uint32_t mask_n_head_log2 = ((sinks != nullptr) << 24) | n_head_log2; @@ -17512,9 +17543,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->query_fusion_node_count.resize(ctx->num_queries); ctx->query_nodes.resize(ctx->num_queries); ctx->query_node_idx.resize(ctx->num_queries); + ctx->query_sub_names.resize(ctx->num_queries); } - ctx->device->device.resetQueryPool(ctx->query_pool, 0, cgraph->n_nodes+1); + // Reset the whole pool, not just n_nodes+1: sub-op marks consume slots past that. + ctx->device->device.resetQueryPool(ctx->query_pool, 0, ctx->num_queries); std::fill(ctx->query_fusion_names.begin(), ctx->query_fusion_names.end(), nullptr); std::fill(ctx->query_fusion_node_count.begin(), ctx->query_fusion_node_count.end(), 0); std::fill(ctx->query_nodes.begin(), ctx->query_nodes.end(), nullptr); @@ -17843,6 +17876,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg // track a single node/fusion for the current query ctx->query_nodes[ctx->query_idx] = cgraph->nodes[i]; ctx->query_fusion_names[ctx->query_idx] = fusion_string; + ctx->query_sub_names[ctx->query_idx] = nullptr; compute_ctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++); ggml_vk_sync_buffers(ctx, compute_ctx); } else { @@ -17884,14 +17918,22 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->compute_ctx.reset(); // Get the results and pass them to the logger - std::vector timestamps(cgraph->n_nodes + 1); - VK_CHECK(ctx->device->device.getQueryPoolResults(ctx->query_pool, 0, ctx->query_idx, (cgraph->n_nodes + 1)*sizeof(uint64_t), timestamps.data(), sizeof(uint64_t), vk::QueryResultFlagBits::e64 | vk::QueryResultFlagBits::eWait), "get timestamp results", ctx->device); + // Sized to the pool, not n_nodes+1: sub-op marks push query_idx past the node count. + std::vector timestamps(ctx->num_queries); + VK_CHECK(ctx->device->device.getQueryPoolResults(ctx->query_pool, 0, ctx->query_idx, ctx->num_queries*sizeof(uint64_t), timestamps.data(), sizeof(uint64_t), vk::QueryResultFlagBits::e64 | vk::QueryResultFlagBits::eWait), "get timestamp results", ctx->device); if (!vk_perf_logger_concurrent) { // Log each op separately for (int i = 1; i < ctx->query_idx; i++) { + const uint64_t dt = uint64_t((timestamps[i] - timestamps[i-1]) * ctx->device->properties.limits.timestampPeriod); + if (ctx->query_sub_names[i] != nullptr) { + // sub-node interval (e.g. the FA K/V contiguize pass) - billed separately so + // it is not silently folded into the op that dispatched it + ctx->perf_logger->log_timing_named(ctx->query_sub_names[i], dt); + continue; + } auto node = ctx->query_nodes[i]; auto name = ctx->query_fusion_names[i]; - ctx->perf_logger->log_timing(node, name, uint64_t((timestamps[i] - timestamps[i-1]) * ctx->device->properties.limits.timestampPeriod)); + ctx->perf_logger->log_timing(node, name, dt); } } else { // Log each group of nodes From d7124e9bfba03dc7fff7365ed77899110966a55d Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 30 Jul 2026 01:25:28 +0000 Subject: [PATCH 051/109] vulkan: pin a 32-wide subgroup for coopmat1 FA where narrowing is free get_fa_tuning_params_coopmat1 took device->subgroup_size unconditionally, so on a 64-wide device the coopmat1 FA path always ran wave64. The sibling scalar path already does AMD-specific wave selection; the coopmat1 path never got the equivalent. Narrowing is free exactly when it does not add an iteration to the O-accumulation loop, which runs ceil((HSV/4) / threads_per_rowgroup) per row. On a 64-wide device the test reduces to hsv <= 128. Above it the narrow subgroup issues 1.5x to 1.8x the instructions for the same SIMD passes and hd256 measures 6 to 18 percent slower, so the rule declines. Gated behind GGML_VK_FA_WAVE32 (=1 rule, =2 forces the pin regardless of head size, diagnostic only). Off by default. Measured on gfx1151 (RADV), model-level pp2048, Qwen3-Coder-30B-A3B: d0 +2.5%, d8192 +8.4%, d16384 +10.1%, d32768 +11.3%. Op-level across head sizes, largest where wave64 wastes the most lanes: hsv=64 +12.3%, hsv=96 +7.4%, hsv=128 +6.5%, hsv=256 correctly declined. Full test-backend-ops suite 15884/15884 at =0, =1 and =2; the pin was confirmed to engage in-band via pipeline VGPR statistics rather than inferred from the env var being set. Assisted-by: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 33 ++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c3036751295f..c6b86355c129 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3851,8 +3851,41 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat1(const vk_device& device result.block_cols = coopmat_block_cols * num_subgroups; result.row_split = num_subgroups; result.subgroup_size = device->subgroup_size; + + // Pin a 32-wide subgroup only where narrowing is free. The shader derives cols_per_iter, + // threads_per_rowgroup and every strided load loop from gl_WorkGroupSize.x, and + // workgroup_size is num_subgroups * subgroup_size, so threads_per_rowgroup always equals the + // real subgroup size. Halving the subgroup halves the workgroup, and the per-lane O state + // grows as d_per_thread = ceil((HSV/4) / threads_per_rowgroup). Pin only when that count is + // unchanged. Above that point the narrow subgroup issues roughly 1.5x to 1.8x the + // instructions for the same number of SIMD passes, which loses on an issue-bound kernel: + // hd256 measures 6 to 18 percent slower. The test depends on HSV only; HSK does not enter + // d_per_thread. On a 64-wide device it reduces exactly to hsv <= 128. + // =1 applies the rule; =2 forces the pin regardless of head size, for measuring the + // configurations the rule rejects. Diagnostic only. + static const int fa_wave32 = [] { + const char * e = getenv("GGML_VK_FA_WAVE32"); + return e ? atoi(e) : 0; + }(); + if (fa_wave32 != 0 && + device->subgroup_size_control && + 32 < device->subgroup_size && // narrow only, never widen + device->subgroup_min_size <= 32 && 32 <= device->subgroup_max_size && + (result.block_cols % 32) == 0 && // cols_per_thread stays >= 1 + (result.block_cols * result.block_rows / 4) >= num_subgroups * 32 && // mask_cache != 0 + (fa_wave32 == 2 || + CEIL_DIV(hsv / 4, 32u) == CEIL_DIV(hsv / 4, device->subgroup_size))) { + result.subgroup_size = 32; + } + result.workgroup_size = num_subgroups * result.subgroup_size; + // threads_per_rowgroup == the real subgroup size is load-bearing in three places: + // the subgroupMax row reduction, the subgroupAdd of Lf, and tmpsh[gl_SubgroupID], which is + // sized by row_split and would be written out of bounds if gl_NumSubgroups exceeded it. + GGML_ASSERT(result.workgroup_size == result.row_split * result.subgroup_size); + GGML_ASSERT(result.block_cols % result.subgroup_size == 0); + 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); From c4444f9b1e3977fb358b328366e822dc796dbe22 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 9 Aug 2026 10:43:34 +0000 Subject: [PATCH 052/109] vulkan: scale the FA MMQ dot product in fp32 before narrowing acc is an int32 sum of dotPacked4x8EXT results. With q8_0 both operands are full int8, so its bound is d_per_step*4*127*127, which overflows f16 when ACC_TYPE is f16 (GGML_PREC_DEFAULT): the score goes +inf, softmax is destroyed and the output comes back -FLT_MAX. Nibble types bound at ~30480 and stay in range, which is why only q8_0 tripped it. Apply the scales in fp32, then narrow. Identical arithmetic when ACC_TYPE is float, so the f32acc path is untouched. This is an upstream bug, not a fork regression, and it was fixed here once before - 61e77f4 carried it and the rebase onto b10133 dropped it. Only the scalar shader has MMQ, so the failing shape is hsk=128 + q8_0 K + prec=def + nb=1 + nr23=[1,1]; GQA>1 and nb>1 route to coopmat1 and escape. test-backend-ops -o FLASH_ATTN_EXT on gfx1151, quiet box: 13295/13295 twice, up from 13257/13295. All 38 failures were type_K=q8_0 prec=def. Both ggml_flash_attn_ext call sites in the tree force GGML_PREC_F32, so no model path reaches this - it is a landmine for the next person touching precision, not a live bug. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp index 0c1b6d0673e9..18a37add9bf8 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp @@ -420,7 +420,11 @@ void main() { acc += dotPacked4x8EXT(Qf[qib].qs[qiqs + d], k_quants[d]); } - Sf[r][c] += ACC_TYPE(acc) * ACC_TYPE(Qf[qib].ds.x) * k_dm.x; + // scale in fp32 before narrowing: acc is an int32 sum of dotPacked4x8EXT + // results, bounded by d_per_step*4*127*127 with q8_0 on both sides, which + // overflows f16 when ACC_TYPE is f16 (GGML_PREC_DEFAULT). Identical + // arithmetic when ACC_TYPE is float. + Sf[r][c] += ACC_TYPE(float(acc) * float(Qf[qib].ds.x) * float(k_dm.x)); if ((d_tid * (HSK_per_thread / 4) + d_block) % 8 == 0) { Sf[r][c] += k_dot_correction(qib, k_dm); } From a94369dcf57399fb242c938867d0ac3fce769f19 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 30 Aug 2026 10:17:20 +0000 Subject: [PATCH 053/109] vulkan: enable the coopmat1 FA wave32 narrowing rule by default Split from the fork's flag-flip commit (f7d804ee7): GGML_VK_FA_WAVE32 now defaults to 1 (apply the HSV-based narrowing rule); =0 opts out and =2 keeps its diagnostic force meaning. The subgroup_size_control device guard is unchanged, so devices without a 32-wide subgroup mode are unaffected. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c6b86355c129..e86fca82379a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3861,11 +3861,11 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat1(const vk_device& device // instructions for the same number of SIMD passes, which loses on an issue-bound kernel: // hd256 measures 6 to 18 percent slower. The test depends on HSV only; HSK does not enter // d_per_thread. On a 64-wide device it reduces exactly to hsv <= 128. - // =1 applies the rule; =2 forces the pin regardless of head size, for measuring the - // configurations the rule rejects. Diagnostic only. + // On by default (=1, applies the rule); =0 disables. =2 forces the pin regardless of + // head size, for measuring the configurations the rule rejects. Diagnostic only. static const int fa_wave32 = [] { const char * e = getenv("GGML_VK_FA_WAVE32"); - return e ? atoi(e) : 0; + return e ? atoi(e) : 1; }(); if (fa_wave32 != 0 && device->subgroup_size_control && From 41d17f1c767558b207116d59acaf04a340aaae5e Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 14 Jul 2026 05:53:12 +0000 Subject: [PATCH 054/109] vulkan: mul_mat_id per-expert-n tile selection (Stage 2a, env-gated) Select the matmul_id tile by expected per-expert token count (nei1*nei0/n_expert) instead of aggregate nei1, which always picked the widest tile and left most N-lanes empty at MoE prefill (~16 rows/expert vs BN=64). Gated on GGML_VK_MMID_SMALLN=1, default off pending cross-model validation. Was null (+0.7%) standalone pre-row-lists: smaller tiles meant more workgroups each re-paying the per-WG id scan. With the Stage 1 row-list prepass the scan is gone and the occupancy win materializes. Strix Halo (RADV gfx1151), Qwen3.6-35B-A3B UD-Q5_K_XL, clean window, pp512 4-way (rowlists x smalln): 914.7 / 937.9 / 1005.3 / 1063.6 t/s (combined +16.3% vs baseline). MUL_MAT_ID q5_K 2.33->4.43 TFLOPS (1.91x), q6_K 1.99->3.57 (1.79x) cumulative. Hot pipelines verified via probe: matmul_id_*_f32 (y_f32 path, quantize_y does not engage for these quants); smalln flips tile _m -> _s. 790/790 MUL_MAT_ID both rowlists configs. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8fbb1359f406..8a84c4a88733 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -10255,7 +10255,18 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& const ggml_type effective_src1_type = quantize_y ? GGML_TYPE_Q8_1 : (y_f32_kernel ? GGML_TYPE_F32 : src1->type); - const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, nei1, qx_needs_dequant ? f16_type : src0->type, effective_src1_type)); + // EXPERIMENT (GGML_VK_MMID_SMALLN=1): select the matmul tile by the EXPECTED PER-EXPERT + // token count rather than the whole batch. With E experts and nei0 active per token, each + // expert sees ~nei1*nei0/E rows; selecting by aggregate nei1 always picks the widest tile + // and leaves most N-lanes empty at MoE prefill (measured: MUL_MAT_ID at ~20% of dense + // matmul efficiency, ~78% of MoE prefill time on Qwen3.6-35B-A3B). + uint32_t n_for_tile = (uint32_t)nei1; + static const char * mmid_smalln_env = getenv("GGML_VK_MMID_SMALLN"); + if (mmid_smalln_env && atoi(mmid_smalln_env) != 0 && ne02 > 1) { + n_for_tile = std::max(1u, (uint32_t)((nei1 * nei0 + ne02 - 1) / ne02)); + } + + const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, n_for_tile, qx_needs_dequant ? f16_type : src0->type, effective_src1_type)); // Coopmat2 MUL_MAT_ID BK specialization constants in ggml_vk_load_shaders are at most 64. const uint32_t y_staged_row_stride = ctx->device->coopmat2 && !quantize_y ? ggml_vk_align_size(ne10, 64) : ne10; const bool y_needs_k_padding = ne10 != y_staged_row_stride; @@ -10264,10 +10275,9 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& // Not implemented GGML_ASSERT(y_needs_reformat || !qy_needs_dequant); // NOLINT - const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && nei1 > 8; - vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type); + vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, n_for_tile, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type); if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) { pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline); From 7f0c243b8111d86dc55d80c1ea41307851c7cc06 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 14 Jul 2026 06:04:18 +0000 Subject: [PATCH 055/109] vulkan: mul_mat_id small-tile shape probes (env-gated) Two env-gated overrides of the mmid small-tile config (KHR coopmat branch only, scoped to the mul_mat_id quant pipelines; dense untouched): - GGML_VK_MMID_TILE16: BN/WN 32->16. NEGATIVE (-3.8% e2e): experts with n_e>16 split into two column tiles and re-stream their full weight matrix; A-traffic scales with sum(ceil(n_e/BN)), so BN must not drop below the mean per-expert n. Kept as a documented dead end. - GGML_VK_MMID_BM64: BM 32->64, BLOCK_SIZE 64->128 (two warps). +1.3% e2e on top of rowlists+smalln: same A-traffic, half the ir-tiles so half the per-expert B re-reads, larger WGs hide latency. Strix Halo, Qwen3.6-35B-A3B UD-Q5_K_XL pp512, clean window, all on top of RL+SN control 1065.4 +/- 3.3: BM64 1078.9 +/- 3.5, TILE16 1024.6, BM64+TILE16 1007.5. Cumulative vs pre-Stage-1 baseline: 914.7 -> 1078.9 (+17.9%). 790/790 MUL_MAT_ID for all configs. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8a84c4a88733..5e40aee445ab 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4848,6 +4848,34 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } #endif + // EXPERIMENT (GGML_VK_MMID_TILE16=1): narrow the matmul_id small tile to BN=16. + // MoE prefill leaves ~nei1*nei0/n_expert rows per expert (~16 at ub512/256E/8a), + // so even the 32-wide small tile runs half empty. Only meaningful stacked on + // GGML_VK_MMID_SMALLN=1 (routes mmid to the small tile) + the row-list prepass. + // Shadows the s-tile config for the mmid quant pipelines only; dense unaffected. + auto s_warptile_mmq_id16 = s_warptile_mmq; + auto s_mmq_wg_denoms_id16 = s_mmq_wg_denoms; + { + const char * tile16_env = getenv("GGML_VK_MMID_TILE16"); + if (tile16_env && atoi(tile16_env) != 0) { + s_warptile_mmq_id16[2] = 16; // BN + s_warptile_mmq_id16[5] = 16; // WN + s_mmq_wg_denoms_id16[1] = 16; + } + // GGML_VK_MMID_BM64=1: taller small tile (BM 32->64, two warps). Same + // A-traffic (ic-tile count unchanged), halves per-expert B re-reads + // (ir-tile count), larger WGs to hide latency. + const char * bm64_env = getenv("GGML_VK_MMID_BM64"); + if (bm64_env && atoi(bm64_env) != 0) { + s_warptile_mmq_id16[0] = 2 * mul_mat_subgroup_size; // BLOCK_SIZE + s_warptile_mmq_id16[1] = 64; // BM + s_mmq_wg_denoms_id16[0] = 64; + } + } + { + const auto &s_warptile_mmq = s_warptile_mmq_id16; + const auto &s_mmq_wg_denoms = s_mmq_wg_denoms_id16; + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); From afedbabda1a814745df5052fbdb0f09e32b668f9 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 14 Jul 2026 08:42:07 +0000 Subject: [PATCH 056/109] vulkan: mul_mat_id taller medium tile probe (GGML_VK_MMID_M128, env-gated) Extends the BM64 idea to the medium tile: GGML_VK_MMID_M128=1 raises the mmid m-tile to BM=128 / BLOCK_SIZE=4*subgroup (four warps), same A-traffic, half the ir-tiles and per-expert B re-reads. The medium tile is what the per-expert-n heuristic selects at n~64 (e.g. ub2048 on 256-expert/top-8 models, or ub1024 at 128 experts). Strix Halo, Qwen3.6-35B-A3B UD-Q5_K_XL pp2048, drain-verified window: ub2048 stack 1039.9 +/- 1.1 -> +M128 1091.6 +/- 2.6 (+5.0%); inert in the s-tile regime (ub1024 1137.2 vs 1147.9 ref). ub1024 remains the throughput sweet spot for this model. 790/790 MUL_MAT_ID. Also measured this session, NOT kept: caching counts+row lists across the three expert matmuls of a layer (they share one ids tensor) was correctness-clean but perf-null (1066.3 vs 1070.9 baseline) - the prepass dispatches and barriers are already free on this queue; reverted rather than carry the invalidation surface. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 5e40aee445ab..43bb141c6291 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4855,6 +4855,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // Shadows the s-tile config for the mmid quant pipelines only; dense unaffected. auto s_warptile_mmq_id16 = s_warptile_mmq; auto s_mmq_wg_denoms_id16 = s_mmq_wg_denoms; + auto m_warptile_mmq_id128 = m_warptile_mmq; + auto m_mmq_wg_denoms_id128 = m_mmq_wg_denoms; { const char * tile16_env = getenv("GGML_VK_MMID_TILE16"); if (tile16_env && atoi(tile16_env) != 0) { @@ -4871,10 +4873,20 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { s_warptile_mmq_id16[1] = 64; // BM s_mmq_wg_denoms_id16[0] = 64; } + // GGML_VK_MMID_M128=1: same idea for the medium tile (BM 64->128, four + // warps) — the tile the per-expert-n heuristic picks at n~64 (e.g. ub2048). + const char * m128_env = getenv("GGML_VK_MMID_M128"); + if (m128_env && atoi(m128_env) != 0) { + m_warptile_mmq_id128[0] = 4 * mul_mat_subgroup_size; // BLOCK_SIZE + m_warptile_mmq_id128[1] = 128; // BM + m_mmq_wg_denoms_id128[0] = 128; + } } { const auto &s_warptile_mmq = s_warptile_mmq_id16; const auto &s_mmq_wg_denoms = s_mmq_wg_denoms_id16; + const auto &m_warptile_mmq = m_warptile_mmq_id128; + const auto &m_mmq_wg_denoms = m_mmq_wg_denoms_id128; CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); From 7722eb4e90c5940bf47b87fab9fe88221c47fcf7 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 14 Jul 2026 11:20:31 +0000 Subject: [PATCH 057/109] vulkan: mmid wave32 probe (GGML_VK_MMID_WAVE32, env-gated) Force required subgroup size 32 on the KHR-coopmat mmid quant pipelines (pipeline_dequant_mul_mat_mat_id[*] only; dense untouched). Hypothesis: RDNA3.5 WMMA is wave32-native, so RADV may lower KHR_coopmat better at wave32 than at the reported default of 64. Mechanics: the cm1 path in mul_mm.comp derives its warp grid from the real subgroup (warp_i = gl_SubgroupID, tiw = gl_SubgroupInvocationID) while NUM_WARPS = BLOCK_SIZE/WARP (spec constants) sizes coopmat_stage[] and ballots_sh[] and warp_r/warp_c assume NUM_WARPS == (BM/WM)*(BN/WN). Forcing sg32 with WARP=64 would over-run those shared arrays and leave warp_c outside the tile, so the gate (a) sets WARP=32 in shadowed copies of the s/m/l mmid warptiles (composes after the BM64/M128 shadows) and (b) halves WM (or WN, keeping WM>=TM, WN>=TN) until the doubled subgroup count exactly tiles BM x BN again, asserting both invariants. BLOCK_SIZE is kept, so workgroup shape, load loops and shmem match the wave64 stack; per-lane accumulator footprint is also unchanged (half the lanes per subgroup, half the (WM/TM)*(WN/TN) fragments). The required size is passed via a scoped CREATE_MM redefinition adding a trailing required_subgroup_size arg (fp16-branch pattern), gated on subgroup_size_control covering 32 since ggml_vk_create_pipeline_func silently drops the required size otherwise. Correctness (test-backend-ops -o MUL_MAT_ID -b Vulkan0): 790/790 plain, 790/790 WAVE32=1, 790/790 WAVE32+SMALLN+BM64. Perf, Qwen3.6-35B-A3B-UD-Q5_K_XL, fa=1 b/ub=512 ctk/ctv=q8_0 pp512 r=3, atomic window, canary clean: stack (SMALLN+BM64), wave64 default: 1076.61 +/- 3.71 t/s stack + WAVE32: 1106.73 +/- 2.52 t/s (+2.8%) Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 58 ++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 43bb141c6291..3c8d42f2b823 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4857,6 +4857,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { auto s_mmq_wg_denoms_id16 = s_mmq_wg_denoms; auto m_warptile_mmq_id128 = m_warptile_mmq; auto m_mmq_wg_denoms_id128 = m_mmq_wg_denoms; + auto l_warptile_mmq_idw = l_warptile_mmq; + uint32_t mmid_req_sgs = 0; { const char * tile16_env = getenv("GGML_VK_MMID_TILE16"); if (tile16_env && atoi(tile16_env) != 0) { @@ -4881,12 +4883,68 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { m_warptile_mmq_id128[1] = 128; // BM m_mmq_wg_denoms_id128[0] = 128; } + // GGML_VK_MMID_WAVE32=1: force required subgroup size 32 on the mmid + // quant coopmat pipelines (RDNA3.x WMMA is wave32-native; probe whether + // RADV lowers KHR_coopmat better at wave32). The cm1 path derives its + // warp grid from the real subgroup (warp_i = gl_SubgroupID, tiw = + // gl_SubgroupInvocationID) and sizes shared arrays (coopmat_stage, + // ballots_sh) with NUM_WARPS = BLOCK_SIZE / WARP, so the WARP spec + // constant must equal the forced size, and the coverage invariant + // NUM_WARPS == (BM/WM)*(BN/WN) must be restored. BLOCK_SIZE is kept + // (same workgroup shape, load loops and shmem as the wave64 stack), so + // the subgroup count doubles and WM (or WN) is halved until the warp + // grid exactly tiles BM x BN again. Per-lane coopmat accumulator + // footprint is unchanged: half the lanes per subgroup, half the + // (WM/TM)*(WN/TN) fragments per subgroup. Runs after the BM64/M128 + // gates so it composes with the probe stack. Applies only when the + // driver honors a required size (subgroup_size_control covering 32); + // otherwise WARP=32 with a real subgroup of 64 would corrupt tiling. + const char * wave32_env = getenv("GGML_VK_MMID_WAVE32"); + if (wave32_env && atoi(wave32_env) != 0 && device->subgroup_size_control && + device->subgroup_min_size <= 32 && 32 <= device->subgroup_max_size) { + mmid_req_sgs = 32; + auto wave32_tile = [](std::vector &w) { + // {BLOCK_SIZE, BM, BN, BK, WM, WN, WMITER, TM, TN, TK, WARP} + w[10] = 32; // WARP: must match the forced subgroup size + for (int guard = 0; guard < 4 && w[0] / w[10] != (w[1] / w[4]) * (w[2] / w[5]); ++guard) { + if (w[4] >= w[5] && w[4] > w[7]) { + w[4] /= 2; // halve WM, keeping WM >= TM + } else { + w[5] /= 2; // halve WN + } + } + GGML_ASSERT(w[0] / w[10] == (w[1] / w[4]) * (w[2] / w[5])); // NUM_WARPS == (BM/WM)*(BN/WN) + GGML_ASSERT(w[4] >= w[7] && w[5] >= w[8]); // WM >= TM, WN >= TN + }; + wave32_tile(s_warptile_mmq_id16); + wave32_tile(m_warptile_mmq_id128); + wave32_tile(l_warptile_mmq_idw); + } } { const auto &s_warptile_mmq = s_warptile_mmq_id16; const auto &s_mmq_wg_denoms = s_mmq_wg_denoms_id16; const auto &m_warptile_mmq = m_warptile_mmq_id128; const auto &m_mmq_wg_denoms = m_mmq_wg_denoms_id128; + const auto &l_warptile_mmq = l_warptile_mmq_idw; + + // Same expansion as CREATE_MM above, plus a trailing required subgroup + // size (0 = driver default) for the GGML_VK_MMID_WAVE32 probe. Scoped to + // the mmid quant pipelines below; dense pipelines are untouched. +#undef CREATE_MM +#define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, true, mmid_req_sgs); \ CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); From 098685079c02c4c788dcd80164be6adefb748767 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 14 Jul 2026 11:29:58 +0000 Subject: [PATCH 058/109] vulkan: mul_mat_id f16-B probe (GGML_VK_MMID_F16B, env-gated) Convert contiguous f32 activations (B) to f16 for quantized MUL_MAT_ID on KHR_coopmat devices and run the matmul_id_subgroup__f16 kernels instead of the f32-B ones. Halves B bytes and buf_b shared memory. The _f16 SPIR-V already exists for every quant type; this adds a parallel env-gated pipeline array (pipeline_dequant_mul_mat_mat_id_f16b), extends the getter to return it for src1=F16 on non-coopmat2 devices, relaxes the src1-type assert, and forces the existing y_non_contig convert-to- prealloc_y plumbing (same as coopmat2). Default OFF, zero behavior change when unset. Measured on Radeon 8060S (RADV gfx1151), Qwen3.6-35B-A3B-UD-Q5_K_XL, -fa 1 -b 512 -ub 512 -ctk q8_0 -ctv q8_0 -p 512 -r 3, stacked on GGML_VK_MMID_SMALLN=1 GGML_VK_MMID_BM64=1: stack (f32 B, canary): pp512 1075.04 +/- 8.66 t/s stack + F16B: pp512 1100.66 +/- 2.37 t/s (+2.4%) test-backend-ops test -o MUL_MAT_ID -b Vulkan0: 790/790 with and without GGML_VK_MMID_F16B=1. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 72 +++++++++++++++++++++++++++- 1 file changed, 70 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 3c8d42f2b823..72c66a88dd0d 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -923,6 +923,10 @@ struct vk_device_struct { vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id[GGML_TYPE_COUNT]; vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_COUNT]; + // EXPERIMENT (GGML_VK_MMID_F16B=1): f16-B mul_mat_id pipelines on KHR_coopmat + // devices (upstream only builds f32-B there). Populated only when the env flag + // is set; empty otherwise. + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_COUNT]; vk_pipeline pipeline_matmul_split_k_reduce; vk_pipeline pipeline_quantize_q8_1_x4; @@ -4212,6 +4216,18 @@ struct CompileTask { uint32_t required_subgroup_size; }; +// EXPERIMENT (GGML_VK_MMID_F16B=1): convert the contiguous f32 activations (B) of +// quantized MUL_MAT_ID to f16 and run the f16-B matmul_id kernels instead of the +// f32-B ones. Halves B bytes and buf_b shared memory (better occupancy) at the +// f32->f16 rounding cost upstream already accepts on the coopmat2 path. +static bool ggml_vk_mmid_f16b_enabled() { + static const bool enabled = [] { + const char * env = getenv("GGML_VK_MMID_F16B"); + return env != nullptr && atoi(env) != 0; + }(); + return enabled; +} + static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { VK_LOG_DEBUG("ggml_vk_load_shaders(" << device->name << ")"); @@ -4978,6 +4994,37 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); } + + // EXPERIMENT (GGML_VK_MMID_F16B=1): f16-B variants of the quant matmul_id + // pipelines. The _f16 SPIR-V exists for every quant type; upstream just never + // instantiates it in the KHR_coopmat branch. Same warptiles as the f32-B lines + // (including the SMALLN/BM64/M128 tile experiments shadowed above). + if (ggml_vk_mmid_f16b_enabled()) { + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q6_K], matmul_id_subgroup_q6_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ1_S], matmul_id_subgroup_iq1_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ1_M], matmul_id_subgroup_iq1_m_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ2_XXS], matmul_id_subgroup_iq2_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ2_XS], matmul_id_subgroup_iq2_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ2_S], matmul_id_subgroup_iq2_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ3_XXS], matmul_id_subgroup_iq3_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + } + } + #undef CREATE_MM2 #undef CREATE_MM } else @@ -8057,7 +8104,8 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_co return pipelines; } - GGML_ASSERT(src1_type == GGML_TYPE_F32 || (ctx->device->coopmat2 && src1_type == GGML_TYPE_F16)); + GGML_ASSERT(src1_type == GGML_TYPE_F32 || + ((ctx->device->coopmat2 || ggml_vk_mmid_f16b_enabled()) && src1_type == GGML_TYPE_F16)); switch (src0_type) { case GGML_TYPE_Q1_0: @@ -8089,7 +8137,11 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_co return nullptr; } - vk_matmul_pipeline2& mmp = ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type]; + // GGML_VK_MMID_F16B: on KHR_coopmat devices the f16-B mmid pipelines live in a + // parallel array (coopmat2's main array already holds f16-B pipelines). + vk_matmul_pipeline2& mmp = (src1_type == GGML_TYPE_F16 && !ctx->device->coopmat2) + ? ctx->device->pipeline_dequant_mul_mat_mat_id_f16b[src0_type] + : ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type]; // XXX TODO 'prec' is not actually allowed in mul_mat_id. bool prefer_fp16acc = ctx->device->fp16 /*&& prec == GGML_PREC_DEFAULT*/; bool support_fp16acc = !mmp.f16acc->is_empty(); @@ -10325,7 +10377,23 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& #else const bool y_decode_vector_staging = false; #endif + // EXPERIMENT (GGML_VK_MMID_F16B=1): route quantized MUL_MAT_ID through the f16-B + // kernels. Treating contiguous f32 B as y_non_contig reuses the existing + // convert-to-prealloc_y plumbing (to_fp16_vk_1), exactly like coopmat2 does. + // Gated on coopmat_support because the f16b pipelines are only created there. + const bool mmid_f16b = ggml_vk_mmid_f16b_enabled() && + ctx->device->coopmat_support && !ctx->device->coopmat2 && + ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32; + if (mmid_f16b) { + static bool mmid_f16b_logged = false; + if (!mmid_f16b_logged) { + mmid_f16b_logged = true; + fprintf(stderr, "ggml_vulkan: MUL_MAT_ID f16-B path engaged (GGML_VK_MMID_F16B)\n"); + } + } + const bool y_non_contig = y_decode_vector_staging || + mmid_f16b || (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) || !ggml_vk_dim01_contiguous(src1); From 633ab28933dd30deee5e9b8c3d987cc3470c00dd Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 26 Jul 2026 10:53:57 +0000 Subject: [PATCH 059/109] vulkan: guard mmid f16-B path on pipeline existence (Q2_0 fallback) --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 72c66a88dd0d..cddb7da8d694 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -10383,7 +10383,11 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& // Gated on coopmat_support because the f16b pipelines are only created there. const bool mmid_f16b = ggml_vk_mmid_f16b_enabled() && ctx->device->coopmat_support && !ctx->device->coopmat2 && - ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32; + ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && + // only take the f16-B path if a pipeline exists for this src0 type + // (e.g. Q2_0 has none); otherwise fall through to the normal f32-B path. + !(ctx->device->pipeline_dequant_mul_mat_mat_id_f16b[src0->type].f16acc->is_empty() && + ctx->device->pipeline_dequant_mul_mat_mat_id_f16b[src0->type].f32acc->is_empty()); if (mmid_f16b) { static bool mmid_f16b_logged = false; if (!mmid_f16b_logged) { From bc235d8e9508ea29b07d93f9d9c2ab50eb4d47f4 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 6 Aug 2026 22:30:38 +0000 Subject: [PATCH 060/109] tests: cover MMQ tile boundaries in MUL_MAT and MUL_MAT_ID The MMQ J config is chosen from n, so a wave-partitioning error in one J tile only shows up when n sits on that tile's boundary. The neighbouring n picks a different J and passes, which hides it. The existing quantized cases stop at n=129 and the general MUL_MAT set jumps 64 -> 4096, so nothing lands on 256 or 512 and the whole class went untested. Sweep n over 255/256/257/511/512/513 for q8_0, q4_0, q4_K, q5_K and q6_K, in both MUL_MAT and MUL_MAT_ID. On an RDNA3.5 build that runs the J128 kernel with 16 wave32 waves over a 128-row tile these fail for every n that selects J128 and pass at n=513, which selects J112. A build predating that config passes the whole sweep. Co-Authored-By: Claude Opus 5 (cherry picked from commit b54cd8a8add4ddb2710eaf256b9fb2ba2be4b383) --- tests/test-backend-ops.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4a7a0623174c..337190f80b3c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9393,6 +9393,19 @@ static std::vector> make_test_cases_eval() { int k = 256; test_cases.emplace_back(new test_mul_mat_id(type_a, type_b, n_mats, n_used, b, m, n, k)); } + // MMQ tile-boundary cases. The MMQ J config is picked from n, and a wave-partitioning error in + // one J tile only shows up when n sits on that tile's boundary: the neighbouring n selects a + // different J and passes, which hides it. The general cases above stop at 129 and the MUL_MAT + // set jumps 64 -> 4096, so no existing case lands on 256 or 512. + // MMQ tile-boundary sweep: n on and either side of the 256 / 512 J boundaries. + for (ggml_type type_a : {GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K}) { + for (int64_t n : {255, 256, 257, 511, 512, 513}) { + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 1024, n, 256, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, 2, false, 1024, n, 256)); + } + } + + } } } From 3df1fd5f3152315991a716e7d8441293188029ab Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sat, 8 Aug 2026 04:24:38 +0000 Subject: [PATCH 061/109] vulkan: create coopmat2 mul_mat_id pipelines with the real param count 30d8bb02b raised mul_mat_id_param_count to 6 for the fused MUL epilogue and gave every mul_mat_id shader a binding 5 (FusedScale), including the coopmat2 one, which binds it purely for descriptor-layout parity. The coopmat2 pipeline creation block was left passing a literal 5, so two things go wrong there: - the pipelines are created with 5 descriptors while the shader declares 6 - PARAMCOUNT == mul_mat_id_param_count doubles as the "this is mul_mat_id" argument to ggml_vk_mul_mm_cm2_spec, so it went false and every coopmat2 mul_mat_id pipeline was specialized as a plain matmul, dropping the trailing spec constant Only reachable where device->coopmat2 is true. gfx1151 does not take that path and the v0.5 release predates the constant bump, so neither is affected. Not validated on hardware - no coopmat2 device here. The block does compile: built with the pinned shaderc (GL_NV_cooperative_matrix2 supported). The two OCP FP4 sites stay behind GL_EXT_float_e2m1, which that glslc does not support. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 56 ++++++++++++++-------------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index cddb7da8d694..fb59e4163c2f 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4740,43 +4740,43 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { GGML_ASSERT(device->subgroup_ballot); - CREATE_MM2(pipeline_matmul_id_f16, matmul_id_subgroup_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, 5) + CREATE_MM2(pipeline_matmul_id_f16, matmul_id_subgroup_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count) #if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) if (device->coopmat_bf16_support) { - CREATE_MM(pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, 5) + CREATE_MM(pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count) } #endif - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K], matmul_id_subgroup_q6_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_S], matmul_id_subgroup_iq1_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_M], matmul_id_subgroup_iq1_m_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XXS], matmul_id_subgroup_iq2_xxs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XS], matmul_id_subgroup_iq2_xs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_S], matmul_id_subgroup_iq2_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_XXS], matmul_id_subgroup_iq3_xxs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K], matmul_id_subgroup_q6_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_S], matmul_id_subgroup_iq1_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_M], matmul_id_subgroup_iq1_m_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XXS], matmul_id_subgroup_iq2_xxs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XS], matmul_id_subgroup_iq2_xs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_S], matmul_id_subgroup_iq2_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_XXS], matmul_id_subgroup_iq3_xxs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) #if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT) if (device->ocp_fp4) { - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16_ocp, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16_ocp, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16_ocp, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16_ocp, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) } else #endif { - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) - CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) } #undef CREATE_MM #undef CREATE_MM2 From 266166ff1231b2bfda5a0ace4ae7313e99b4a92b Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 30 Aug 2026 10:20:41 +0000 Subject: [PATCH 062/109] vulkan: enable the Strix mmid tile gates by default Split from the fork's flag-flip commit (f7d804ee7): GGML_VK_MMID_F16B, BM64, M128, WAVE32 and SMALLN now default on with =0 opt-out. The wave32 gate keeps its subgroup_size_control device guard, so devices without a 32-wide subgroup mode are unaffected. GGML_VK_MMID_TILE16 stays opt-in (documented negative on gfx1151). Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index fb59e4163c2f..371461e974c5 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4223,7 +4223,7 @@ struct CompileTask { static bool ggml_vk_mmid_f16b_enabled() { static const bool enabled = [] { const char * env = getenv("GGML_VK_MMID_F16B"); - return env != nullptr && atoi(env) != 0; + return env == nullptr || atoi(env) != 0; // on by default; =0 disables }(); return enabled; } @@ -4886,7 +4886,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // A-traffic (ic-tile count unchanged), halves per-expert B re-reads // (ir-tile count), larger WGs to hide latency. const char * bm64_env = getenv("GGML_VK_MMID_BM64"); - if (bm64_env && atoi(bm64_env) != 0) { + if (!bm64_env || atoi(bm64_env) != 0) { // on by default; =0 disables s_warptile_mmq_id16[0] = 2 * mul_mat_subgroup_size; // BLOCK_SIZE s_warptile_mmq_id16[1] = 64; // BM s_mmq_wg_denoms_id16[0] = 64; @@ -4894,7 +4894,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // GGML_VK_MMID_M128=1: same idea for the medium tile (BM 64->128, four // warps) — the tile the per-expert-n heuristic picks at n~64 (e.g. ub2048). const char * m128_env = getenv("GGML_VK_MMID_M128"); - if (m128_env && atoi(m128_env) != 0) { + if (!m128_env || atoi(m128_env) != 0) { // on by default; =0 disables m_warptile_mmq_id128[0] = 4 * mul_mat_subgroup_size; // BLOCK_SIZE m_warptile_mmq_id128[1] = 128; // BM m_mmq_wg_denoms_id128[0] = 128; @@ -4916,7 +4916,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // driver honors a required size (subgroup_size_control covering 32); // otherwise WARP=32 with a real subgroup of 64 would corrupt tiling. const char * wave32_env = getenv("GGML_VK_MMID_WAVE32"); - if (wave32_env && atoi(wave32_env) != 0 && device->subgroup_size_control && + if ((!wave32_env || atoi(wave32_env) != 0) && device->subgroup_size_control && // on by default; =0 disables device->subgroup_min_size <= 32 && 32 <= device->subgroup_max_size) { mmid_req_sgs = 32; auto wave32_tile = [](std::vector &w) { @@ -10432,7 +10432,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& // matmul efficiency, ~78% of MoE prefill time on Qwen3.6-35B-A3B). uint32_t n_for_tile = (uint32_t)nei1; static const char * mmid_smalln_env = getenv("GGML_VK_MMID_SMALLN"); - if (mmid_smalln_env && atoi(mmid_smalln_env) != 0 && ne02 > 1) { + if (!(mmid_smalln_env && atoi(mmid_smalln_env) == 0) && ne02 > 1) { // on by default; =0 disables n_for_tile = std::max(1u, (uint32_t)((nei1 * nei0 + ne02 - 1) / ne02)); } From fcf2f389d28b3de1795640dce724682c9e34e01c Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 16 Aug 2026 11:49:48 +0000 Subject: [PATCH 063/109] vulkan : optional f16 B operand for quantized MUL_MAT on coopmat1 The _f16 SPIR-V and pipeline_dequant_mul_mat_mat_f16 already exist, but are only populated in the coopmat2 branch. GGML_VK_DENSE_F16B populates them for coopmat1 too and routes B through the existing convert-to-prealloc_y path. Off by default: it helps large dense models and costs a little elsewhere. gfx1151, pp2048: Qwen3.8-27B and Qwen3-32B (hidden 5120) +5 to +7% for both q6_K and q8_0 weights, Qwen2.5-7B -1.2%, Qwen3-Coder-30B MoE -0.5%. Decode is untouched, ne1==1 does not reach this path. Numerically identical: mul_mm stages B into shared FLOAT_TYPE either way, so the f32-B kernel already rounds B to f16. Wikitext PPL matches to 4 dp. Assisted-by: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 48 ++++++++++++++++++++++++++-- 1 file changed, 45 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8fbb1359f406..e053bcf530ee 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4212,6 +4212,14 @@ struct CompileTask { uint32_t required_subgroup_size; }; +// GGML_VK_DENSE_F16B=1: same idea as GGML_VK_MMID_F16B but for plain MUL_MAT. Halves the B +// bytes, which keeps the activations inside the LLC at large ubatch and moves their row stride +// off the 1-of-16 channel pattern. Off by default, it is a loss when B already fits. +static bool ggml_vk_dense_f16b_enabled() { + static const bool enabled = getenv("GGML_VK_DENSE_F16B") != nullptr; + return enabled; +} + static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { VK_LOG_DEBUG("ggml_vk_load_shaders(" << device->name << ")"); @@ -4816,6 +4824,21 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K], matmul_q6_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + // f16-B variants of the same quant pipelines. The _f16 SPIR-V is already built for + // every type; upstream only instantiates it in the coopmat2 branch. + if (ggml_vk_dense_f16b_enabled()) { + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0], matmul_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1], matmul_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0], matmul_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1], matmul_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0], matmul_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K], matmul_q2_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K], matmul_q6_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + } CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_S], matmul_iq1_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_M], matmul_iq1_m_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); @@ -7793,7 +7816,8 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte return pipelines; } - if (src1_type != GGML_TYPE_F32 && !ctx->device->coopmat2) { + if (src1_type != GGML_TYPE_F32 && !ctx->device->coopmat2 && + !(src1_type == GGML_TYPE_F16 && ggml_vk_dense_f16b_enabled())) { return nullptr; } @@ -7832,7 +7856,9 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte return prec == GGML_PREC_DEFAULT ? ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f32acc; } if (ctx->device->coopmat_support) { - return (ctx->device->fp16 && ctx->device->coopmat_acc_f16_support && prec == GGML_PREC_DEFAULT) ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; + vk_matmul_pipeline2 & p = (src1_type == GGML_TYPE_F16) ? ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type] + : ctx->device->pipeline_dequant_mul_mat_mat[src0_type]; + return (ctx->device->fp16 && ctx->device->coopmat_acc_f16_support && prec == GGML_PREC_DEFAULT) ? p.f16acc : p.f32acc; } return (ctx->device->fp16 && prec == GGML_PREC_DEFAULT) ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; } @@ -9285,7 +9311,23 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub // Reformat and convert to fp16 if non-contiguous, or for coopmat2 for better perf const bool x_non_contig = (ctx->device->coopmat2 && src0->type == GGML_TYPE_F32) || !ggml_vk_dim01_contiguous(src0); - const bool y_non_contig = (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || + // Route quantized MUL_MAT through the f16-B kernels. Treating contiguous f32 B as + // y_non_contig reuses the convert-to-prealloc_y plumbing, like coopmat2 does. + const bool dense_f16b = ggml_vk_dense_f16b_enabled() && + ctx->device->coopmat_support && !ctx->device->coopmat2 && + ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && + !(ctx->device->pipeline_dequant_mul_mat_mat_f16[src0->type].f16acc->is_empty() && + ctx->device->pipeline_dequant_mul_mat_mat_f16[src0->type].f32acc->is_empty()); + if (dense_f16b) { + static bool dense_f16b_logged = false; + if (!dense_f16b_logged) { + dense_f16b_logged = true; + fprintf(stderr, "ggml_vulkan: MUL_MAT f16-B path engaged (GGML_VK_DENSE_F16B)\n"); + } + } + + const bool y_non_contig = dense_f16b || + (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) || !ggml_vk_dim01_contiguous(src1); From 74b3efaf451345941c1fc8b77f8d1292cb9afb29 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 16 Aug 2026 14:34:15 +0000 Subject: [PATCH 064/109] vulkan : add auto mode to GGML_VK_DENSE_F16B =auto restricts the f16-B path to ne10 == 5120, the only reduction width measured to gain, so it cannot fire on the widths that lose. =1 keeps the old all-shapes behaviour as a manual override. gfx1151, Qwen3.8-27B UD-Q6_K_XL pp2048, auto vs off: +5.8 / +6.0 / +5.9 / +5.3% at ub 256/512/1024/2048, which is 97% of the all-shapes win at ub256 and 82-86% above it. Qwen3-Coder-30B MoE is untouched, the gate never fires. The width equality is a stopgap until a per-shape predicate is derived; it will silently do nothing for a dense model of another width. Assisted-by: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index e053bcf530ee..cdda787b5f6c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4212,13 +4212,21 @@ struct CompileTask { uint32_t required_subgroup_size; }; -// GGML_VK_DENSE_F16B=1: same idea as GGML_VK_MMID_F16B but for plain MUL_MAT. Halves the B -// bytes, which keeps the activations inside the LLC at large ubatch and moves their row stride -// off the 1-of-16 channel pattern. Off by default, it is a loss when B already fits. -static bool ggml_vk_dense_f16b_enabled() { - static const bool enabled = getenv("GGML_VK_DENSE_F16B") != nullptr; - return enabled; -} +// GGML_VK_DENSE_F16B: same idea as GGML_VK_MMID_F16B but for plain MUL_MAT. Halves the B bytes +// moved. Numerically identical: mul_mm stages B into shared FLOAT_TYPE either way, so the f32-B +// kernel already rounds B to f16. Helps wide dense models, costs ~1% on narrow ones. +// 0 = off, 1 = all quantized dense matmuls, 2 = auto (only the K we have positive data for) +static int ggml_vk_dense_f16b_mode() { + static const int mode = [] { + const char * e = getenv("GGML_VK_DENSE_F16B"); + if (e == nullptr) return 0; + if (e[0] == 'a') return 2; + return atoi(e) != 0 ? 1 : 0; + }(); + return mode; +} + +static bool ggml_vk_dense_f16b_enabled() { return ggml_vk_dense_f16b_mode() != 0; } static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { VK_LOG_DEBUG("ggml_vk_load_shaders(" << device->name << ")"); @@ -9313,7 +9321,10 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub !ggml_vk_dim01_contiguous(src0); // Route quantized MUL_MAT through the f16-B kernels. Treating contiguous f32 B as // y_non_contig reuses the convert-to-prealloc_y plumbing, like coopmat2 does. + // auto mode restricts to ne10 == 5120: the only width measured to gain. Narrow on purpose, + // so widths measured as losses (3584 dense, 2048 MoE) cannot trigger it. const bool dense_f16b = ggml_vk_dense_f16b_enabled() && + (ggml_vk_dense_f16b_mode() == 1 || ne10 == 5120) && ctx->device->coopmat_support && !ctx->device->coopmat2 && ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && !(ctx->device->pipeline_dequant_mul_mat_mat_f16[src0->type].f16acc->is_empty() && From bc3638c1c5b76b3e9a892efb7ca383ae6da40182 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 18 Aug 2026 01:12:00 +0000 Subject: [PATCH 065/109] vulkan: run the quantised dense coopmat pipelines at wave32 RDNA3.x WMMA is wave32-native, so a wave64 subgroup issues each coopmat op as two halves. GGML_VK_MMID_WAVE32 already exploits this for mul_mat_id and leaves the dense pipelines at the driver default; this gives dense the same treatment, gated on measurement rather than a flag. BLOCK_SIZE is kept, so the subgroup count doubles and WM (or WN) halves until the warp grid tiles BM x BN again. Only the quantised tiles are retiled: on gfx1151 a standalone MUL_MAT microbench at both dense FFN shapes reads q6_K +5.2..+10.8%, q8_0 +5.4..+8.4%, q4_K +0.7..+9.1%, q4_0 -1.5..+1.8%, while f16 reads -6.7..+6.4% and bf16 ~0 - the float paths are bandwidth-bound on the weight stream, not issue-bound. The win tracks inline dequant instruction count (q6_K 3907 -> 3433 instructions, identical 192 VGPRs and 8 subgroups/SIMD). The required subgroup size is now the tile's own WARP for every dense coopmat pipeline. The cm1 shaders derive their warp grid from gl_SubgroupID and size shared arrays as NUM_WARPS = BLOCK_SIZE / WARP, so WARP and the real subgroup must agree; leaving that to the driver made the agreement incidental. Scoped to AMD coopmat1 on a wave64 default; other vendors keep the driver default. GGML_VK_DENSE_WAVE32=0 disables, =2 also retiles the float tiles. Qwen3-32B Q6_K_XL pp2048 +7.2% at ub256 / +3.9% at ub2048, Qwen3.8-27B +5.3% / +4.8%. PPL unchanged: 6.9496 +/- 0.24246 in both arms, all 20 per-chunk values identical, since the retile changes which warp owns an output sub-tile and not the K-reduction order within an element. Assisted-by: Claude Opus 5 (cherry picked from commit 448994e9637405610ced3e0ead02c7b6fa688314) --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 83 ++++++++++++++++++++++++++-- 1 file changed, 77 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index cdda787b5f6c..3c86701863aa 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4784,20 +4784,91 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { #endif // defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) #if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) if (device->coopmat_support) { + // Deterministic subgroup sizing for the dense coopmat pipelines. Two parts: + // + // (1) The required subgroup size is the tile's own WARP, not the driver's choice. The cm1 + // shaders derive their warp grid from the real subgroup (warp_i = gl_SubgroupID) and + // size shared arrays with NUM_WARPS = BLOCK_SIZE / WARP, so WARP and the actual + // subgroup must agree - leaving that to the driver makes the agreement incidental. + // + // (2) The QUANTISED dense tiles run at wave32. RDNA3.x WMMA is wave32-native, so a wave64 + // subgroup issues each coopmat op as two halves. BLOCK_SIZE is kept, so the subgroup + // count doubles and WM (or WN) halves until the warp grid tiles BM x BN again: + // NUM_WARPS == (BM/WM)*(BN/WN). A tile that cannot be retiled is left at wave64. + // + // The FLOAT tiles are deliberately excluded. Measured on gfx1151 with a standalone + // MUL_MAT microbench at both dense FFN shapes (m=25600 k=5120, m=5120 k=25600): + // q6_K +5.2..+10.8%, q8_0 +5.4..+8.4%, q4_K +0.7..+9.1%, q4_0 -1.5..+1.8%, but + // f16 -6.7..+6.4% and bf16 ~0. The win tracks inline dequant instruction count + // (q6_K 3907 -> 3433 instructions at wave32, identical 192 VGPRs and 8 subgroups/SIMD), + // so it lands on the issue-bound quantised kernels and not on the float ones, which are + // bandwidth-bound on the weight stream. + // + // Scoped to AMD coopmat1 on a wave64 default; other vendors keep the driver default, + // since none of the above is validated there. + // GGML_VK_DENSE_WAVE32=0 disables, =2 additionally retiles the float tiles (probe). + const bool dense_sgs_scope = + device->vendor_id == VK_VENDOR_ID_AMD && + device->driver_id != vk::DriverId::eAmdProprietary && + device->subgroup_size_control; + const bool dense_wave32_possible = + dense_sgs_scope && + device->subgroup_min_size <= 32 && 32 <= device->subgroup_max_size && + device->subgroup_size > 32; + const char * dense_wave32_env = getenv("GGML_VK_DENSE_WAVE32"); + const int dense_wave32 = dense_wave32_env ? atoi(dense_wave32_env) : 1; + + if (dense_wave32_possible && dense_wave32 != 0) { + auto wave32_tile = [](std::vector & w) -> bool { + // {BLOCK_SIZE, BM, BN, BK, WM, WN, WMITER, TM, TN, TK, WARP} + std::vector t = w; + t[10] = 32; + for (int guard = 0; guard < 4 && t[0] / t[10] != (t[1] / t[4]) * (t[2] / t[5]); ++guard) { + if (t[4] >= t[5] && t[4] > t[7]) { + t[4] /= 2; // halve WM, keeping WM >= TM + } else { + t[5] /= 2; // halve WN + } + } + if (t[0] / t[10] != (t[1] / t[4]) * (t[2] / t[5]) || t[4] < t[7] || t[5] < t[8]) { + return false; + } + w = t; + return true; + }; + wave32_tile(l_warptile_mmq); + wave32_tile(m_warptile_mmq); + wave32_tile(s_warptile_mmq); + if (dense_wave32 >= 2) { + wave32_tile(l_warptile); + wave32_tile(m_warptile); + wave32_tile(s_warptile); + } + } + + // WARP -> required subgroup size, or 0 where the device cannot honor one. + auto dense_req_sgs = [dense_sgs_scope, &device](const std::vector & w) -> uint32_t { + const uint32_t warp = w[10]; + if (!dense_sgs_scope || warp < device->subgroup_min_size || warp > device->subgroup_max_size) { + return 0; + } + return warp; + }; + // Create 6 variants, {s,m,l}x{unaligned,aligned} #define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ if (device->mul_mat ## ID ## _l[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true, dense_req_sgs(l_ ## WARPTILE)); \ if (device->mul_mat ## ID ## _m[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, true, dense_req_sgs(m_ ## WARPTILE)); \ if (device->mul_mat ## ID ## _s[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, true, dense_req_sgs(s_ ## WARPTILE)); \ if (device->mul_mat ## ID ## _l[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, true, dense_req_sgs(l_ ## WARPTILE)); \ if (device->mul_mat ## ID ## _m[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, true, dense_req_sgs(m_ ## WARPTILE)); \ if (device->mul_mat ## ID ## _s[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, true, dense_req_sgs(s_ ## WARPTILE)); \ // Create 2 variants, {f16,f32} accumulator #define CREATE_MM2(TYPE, PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ From beea09e32e3f6d71deb016d3ea9c861e44f84e53 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Tue, 18 Aug 2026 02:37:00 +0000 Subject: [PATCH 066/109] vulkan: tune the coopmat matmul LDS pad per path on RADV buf_a/buf_b are FLOAT_TYPEV2 (4 B), so the shared-memory stride in elements is the stride in LDS banks, and RDNA has 32: SHMEM_STRIDE = BK/2 + pad reaches 32/gcd(BK/2+pad, 32) banks. The pad was a single constant for every non-Intel device and 4 is not the optimum on gfx1151. The stride must stay even, or the 8/16-byte ds_read_b64/b128 loads lose alignment - every odd pad measures -53% in a standalone MUL_MAT sweep. Among even pads the quantised path tracks bank spread: pad 2 (stride 18, 16 banks) is +13% mean over pad 4 (stride 20, 8 banks), and pad 0 (stride 16, 2 banks) -31%. Per weight type at m=25600 n=2048 k=5120: q4_0 +32%, q8_0 +22%, q4_K +10%, q6_K +1%. The win is largest where the dequant is cheapest, i.e. where the kernel is least ALU-bound, which is the complement of what wave32 helps. The float path is left at pad 4 on measurement, not theory. f32/f16 shaders `#define BK 32` regardless of the host spec constant, so they run the same stride as the quantised ones, yet pad 2 measures -18% on f16 and -14% on f32. Same stride, same bank pattern, opposite preference; the unmodelled difference is BK_STEP (4 float, 2 quant). The discriminator is the host warptile's BK, which labels the pipeline reliably even where the shader overrides the value. ggml_vk_matmul_shmem_support now derives its bank_conflict_offset from the same helper, so the shared-memory budget stays aligned with the pad actually pushed. Qwen3.8-27B UD-Q6_K_XL pp2048 +8.3% at ub256 / +6.6% at ub2048, Qwen3-32B +4.1% / +3.6%, Qwen3-Coder-30B-A3B Q4_K_XL +5.8% at both, all f16 KV and on top of the wave32 change. Decode unchanged (tg128 8.52 both arms). PPL unchanged: 6.9496 +/- 0.24246 in both arms with all 20 per-chunk values identical, since the pad moves addresses only, not arithmetic or its order. Assisted-by: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 58 ++++++++++++++++++++++++---- 1 file changed, 50 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 3c86701863aa..87da54a4f86e 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3978,6 +3978,44 @@ static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& s }; } +// LDS bank spread for the coopmat matmul tiles. buf_a/buf_b are FLOAT_TYPEV2 (4 B), so the shared +// stride in elements IS the stride in LDS banks, and RDNA has 32 of them: +// SHMEM_STRIDE = BK/2 + pad reaches 32/gcd(BK/2+pad, 32) banks. +// +// The stride must stay EVEN or the 8/16-byte ds_read_b64/b128 loads lose alignment - every odd pad +// measured -53% on gfx1151. Among even pads the quantised path (BK=32) tracks bank spread: +// pad 2 (stride 18, 16 banks) is +13% mean over pad 4 (stride 20, 8 banks) in a standalone MUL_MAT +// sweep, and pad 0 (stride 16, 2 banks) is -31%. +// +// The float path does NOT follow that rule and the bank model does NOT explain why. f32/f16 +// shaders `#define BK 32` regardless of the host spec constant, so they run the SAME stride as +// the quantised ones - yet pad 2 measures -18% on f16 and -14% on f32. Same stride, same bank +// pattern, opposite preference; the remaining difference is BK_STEP (4 on the float path, 2 on +// the quant path), i.e. the access pattern inside the loop, which is not modelled here. +// So this ships the MEASURED optimum per path, not the theory: pad 2 quantised, pad 4 float. +// The discriminator is the host warptile's BK (16 float / 32 quant), which is a reliable label +// for which pipeline is being built even though the shader overrides the value for f32/f16. +// GGML_VK_SHMEM_PAD=N overrides both, for probing. +static uint32_t ggml_vk_coopmat_shmem_pad(const vk_device& device, uint32_t bk) { + if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && + device->driver_id == vk::DriverId::eIntelProprietaryWindows) { + return 0; + } + static const int env_pad = [] { + const char * e = getenv("GGML_VK_SHMEM_PAD"); + return e ? atoi(e) : -1; + }(); + if (env_pad >= 0) { + return (uint32_t) env_pad; + } + const bool amd_radv = device->vendor_id == VK_VENDOR_ID_AMD && + device->driver_id != vk::DriverId::eAmdProprietary; + if (device->coopmat_support && amd_radv && bk >= 32) { + return 2; + } + return 4; +} + static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector& warptile, bool mul_mat_id, ggml_type src0_type) { uint32_t lut_size = 0; @@ -4017,10 +4055,11 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec } // Needs to be kept up to date on shader changes - // Needs to stay aligned with ggml_vk_mul_mm_spec. - const bool intel_shmem_stride_pad_zero = device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && - device->driver_id == vk::DriverId::eIntelProprietaryWindows; - const uint32_t bank_conflict_offset = intel_shmem_stride_pad_zero ? 0 : (device->coopmat_support ? 8 : 1); + // Shared-memory budget: BM*(BK + 2*pad)*type_size, so the offset is 2x the pad that + // ggml_vk_mul_mm_spec will actually push. Both call ggml_vk_coopmat_shmem_pad to stay aligned. + const uint32_t bank_conflict_offset = device->coopmat_support + ? 2 * ggml_vk_coopmat_shmem_pad(device, warptile[3]) + : 1; const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float); const uint32_t warps = warptile[0] / warptile[10]; @@ -4661,10 +4700,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { auto const &ggml_vk_mul_mm_spec = [&device](std::vector spec, bool aligned) { spec.push_back(aligned ? 1u : 0u); // constantID=11: ALIGNED - if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && - device->driver_id == vk::DriverId::eIntelProprietaryWindows) { - spec.push_back(0u); // constantID=12: SHMEM_STRIDE_PAD = 0 - spec.push_back(1u); // constantID=13: APPLY_SLM_A_RESHAPE = true + const uint32_t bk = spec[3]; + const uint32_t pad = ggml_vk_coopmat_shmem_pad(device, bk); + const bool intel_slm = device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && + device->driver_id == vk::DriverId::eIntelProprietaryWindows; + if (intel_slm || pad != 4) { + spec.push_back(pad); // constantID=12: SHMEM_STRIDE_PAD + spec.push_back(intel_slm ? 1u : 0u); // constantID=13: APPLY_SLM_A_RESHAPE } return spec; }; From 73608d23157f5fa60e4f147bd290d12b8255f3b2 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sat, 22 Aug 2026 08:43:42 +0000 Subject: [PATCH 067/109] vulkan: gate the RADV coopmat pad 2 on driver >= 25.3 baf6360be set the quantised-path LDS pad to 2 for the 16-bank spread, but the coopmat path passes SHMEM_STRIDE to coopMatLoad as its Stride operand and VUID-RuntimeSpirv-OpCooperativeMatrixLoadKHR-08986 requires pointer and stride to be 16 B aligned for the 16x16 f16 tiles. Stride bytes are (BK/2 + pad) * 4, so only pad % 4 == 0 is in contract and the 16-bank stride (18 elements, 72 B) never is. A downstream report bisected a pp512 collapse on stock Mesa to that commit, and it reproduces exactly. The mechanism is codegen luck, not a driver bug: RADV 25.2 lowers coopMatLoad with ds_read_b128, entitled by the contract, so the misaligned odd rows pay runtime splits; RADV 25.3+ lowers to ds_read_b64, for which the 72 B stride is always aligned, and the extra bank spread wins. Both codegens are pad-invariant per driver (ISA-verified), so the damage is runtime address patterns and invisible on the driver the pad was tuned on. pp512, pad 2 vs pad 4 on gfx1151 (Radeon 8060S), Qwen3.6-35B-A3B UD-Q4_K_XL unless noted: RADV 25.2.8 (Ubuntu): 597 vs 1426 (dense Qwen3.8-27B: 107 vs 333) RADV 25.3.0: 1511 vs 1433 RADV 25.3.6: 1509 vs 1420 RADV 26.0.8: 1524 vs 1436 RADV 26.1.8: 1522 vs 1433 RADV 26.2.1: 1540 vs 1444 RADV 26.3-dev: 1548 vs 1447 So pad 2 stays only where it is measured to win: DriverId eMesaRadv with driverVersion >= 25.3.0 (RADV fills it from the Mesa version). This also narrows the old "AMD and not proprietary" test, which would have matched AMDVLK where pad 2 was never measured. Everything else gets the spec-aligned pad 4, bit-identical pipelines to 25c45fea1. tg128 is unchanged everywhere; GGML_VK_SHMEM_PAD still overrides both paths for probing. On <= 25.2 the workaround GGML_VK_SHMEM_PAD=4 remains valid but is no longer needed. Assisted-by: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 34 +++++++++++++++------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 87da54a4f86e..1d1d0d4c7ec0 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3982,20 +3982,22 @@ static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& s // stride in elements IS the stride in LDS banks, and RDNA has 32 of them: // SHMEM_STRIDE = BK/2 + pad reaches 32/gcd(BK/2+pad, 32) banks. // -// The stride must stay EVEN or the 8/16-byte ds_read_b64/b128 loads lose alignment - every odd pad -// measured -53% on gfx1151. Among even pads the quantised path (BK=32) tracks bank spread: -// pad 2 (stride 18, 16 banks) is +13% mean over pad 4 (stride 20, 8 banks) in a standalone MUL_MAT -// sweep, and pad 0 (stride 16, 2 banks) is -31%. +// The pad is NOT free to chase bank spread: the coopmat path hands SHMEM_STRIDE to coopMatLoad +// as its Stride operand, and VUID-RuntimeSpirv-OpCooperativeMatrixLoadKHR-08986 requires the +// pointer and stride to be 16 B aligned for the 16x16 f16 tiles. Stride bytes are +// (BK/2 + pad) * 4, so only pad % 4 == 0 is in contract, and the 16-bank stride 18 (72 B) +// never is. // -// The float path does NOT follow that rule and the bank model does NOT explain why. f32/f16 -// shaders `#define BK 32` regardless of the host spec constant, so they run the SAME stride as -// the quantised ones - yet pad 2 measures -18% on f16 and -14% on f32. Same stride, same bank -// pattern, opposite preference; the remaining difference is BK_STEP (4 on the float path, 2 on -// the quant path), i.e. the access pattern inside the loop, which is not modelled here. -// So this ships the MEASURED optimum per path, not the theory: pad 2 quantised, pad 4 float. -// The discriminator is the host warptile's BK (16 float / 32 quant), which is a reliable label -// for which pipeline is being built even though the shader overrides the value for f32/f16. -// GGML_VK_SHMEM_PAD=N overrides both, for probing. +// What a driver does with the out-of-contract stride is codegen luck, ISA-verified on gfx1151: +// RADV <= 25.2 lowers coopMatLoad with ds_read_b128 (the 16 B the contract guarantees), so the +// misaligned rows pay runtime splits: Qwen3.6-35B pp512 597 vs 1426 t/s, Qwen3.8-27B 107 vs +// 333. RADV >= 25.3 lowers to ds_read_b64, for which 72 B is always aligned, and the extra +// bank spread nets +5-7% pp512 (measured on 25.3.0, 25.3.6, 26.0.8, 26.1.8, 26.2.1, 26.3-dev). +// The codegen is pad-invariant per driver, so the effect is runtime address patterns, not +// instruction selection. +// +// So: pad 2 for the quant tiles (BK 32) only on RADV >= 25.3.0; the spec-aligned pad 4 +// everywhere else. GGML_VK_SHMEM_PAD=N still overrides both paths, for probing. static uint32_t ggml_vk_coopmat_shmem_pad(const vk_device& device, uint32_t bk) { if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && device->driver_id == vk::DriverId::eIntelProprietaryWindows) { @@ -4008,9 +4010,9 @@ static uint32_t ggml_vk_coopmat_shmem_pad(const vk_device& device, uint32_t bk) if (env_pad >= 0) { return (uint32_t) env_pad; } - const bool amd_radv = device->vendor_id == VK_VENDOR_ID_AMD && - device->driver_id != vk::DriverId::eAmdProprietary; - if (device->coopmat_support && amd_radv && bk >= 32) { + if (device->coopmat_support && bk >= 32 && + device->driver_id == vk::DriverId::eMesaRadv && + device->properties.driverVersion >= VK_MAKE_API_VERSION(0, 25, 3, 0)) { return 2; } return 4; From b9e43c51524fe9c605995745f08219aa30110fc3 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 6 Aug 2026 01:34:10 +0000 Subject: [PATCH 068/109] vulkan: four env-gated Strix Halo prefill fixes for delta-net MoE All default OFF, so the same binary A/Bs each change. GGML_VK_CONCAT_TRANSPOSE: delta-net does ggml_transpose() into a dim-0 ggml_concat(), which the generic concat reads fully de-coalesced. Route that shape through a 32x32 shared-memory tile transpose. CONCAT 11877 -> 957 us/op. GGML_VK_MMID_SCALE_EPILOGUE: apply the following MUL's per-(expert,token) broadcast scale as mul_mat_id writes out, removing a 134 MB write plus read back. Prefill only; the existing fusion is gated to mat-vec. Not implemented in the coopmat2 shader, so it is refused there. GGML_VK_FUSE_UNARY_MUL: silu(x)*y is two nodes in the delta-net path; run it as the existing swiglu split. 750 -> 443 us/op. GGML_VK_MMID_WG256: the RADV tuning gives the dense large tile 256 threads on a 128x128 tile but left the mul_mat_id variants at 128. Qwen3.6-35B-A3B UD-Q4_K_XL, gfx1151, pp2048 at ub2048: 1223.75 -> 1733.64 t/s. test-backend-ops CONCAT/MUL_MAT_ID/UNARY/MUL pass; generated text is unchanged with each flag on, and the disabled paths are bit-identical to before. Co-Authored-By: Claude Opus 5 (cherry picked from commit 6f7a49e661d93fa42df528bdc88bc788eaf0ad2c) --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 182 ++++++++++++++++-- .../vulkan-shaders/concat_transpose.comp | 43 +++++ .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 20 +- .../vulkan-shaders/mul_mm_cm2.comp | 3 + .../ggml-vulkan/vulkan-shaders/mul_mmq.comp | 8 +- .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + 6 files changed, 241 insertions(+), 16 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/concat_transpose.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8fbb1359f406..bfe1c5a93ed0 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -963,6 +963,7 @@ struct vk_device_struct { vk_pipeline pipeline_add_id_f32; vk_pipeline pipeline_concat_i8, pipeline_concat_i16, pipeline_concat_i32, pipeline_concat_i64; + vk_pipeline pipeline_concat_transpose_i32; vk_pipeline pipeline_upscale_nearest_f32, pipeline_upscale_bilinear_f32, pipeline_upscale_bicubic_f32, pipeline_upscale_bilinear_antialias_f32; vk_pipeline pipeline_scale_f32; vk_pipeline pipeline_log[2]; @@ -1349,6 +1350,7 @@ struct vk_mat_mat_id_push_constants { uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; uint32_t n_experts; uint32_t hoist_row_ids; + uint32_t fusion_flags; }; struct vk_mat_vec_id_push_constants { uint32_t ncols; @@ -2284,6 +2286,20 @@ class vk_perf_logger { if (node->op == GGML_OP_UNARY) { return fusion_str + ggml_unary_op_name(ggml_get_unary_op(node)); } + if (node->op == GGML_OP_MUL && getenv("GGML_VK_PERF_SHAPES")) { + std::string name = "MUL "; + name += "dst(" + std::to_string(node->ne[0]) + "," + std::to_string(node->ne[1]) + "," + + std::to_string(node->ne[2]) + ") b(" + std::to_string(node->src[1]->ne[0]) + "," + + std::to_string(node->src[1]->ne[1]) + "," + std::to_string(node->src[1]->ne[2]) + ")"; + name += std::string(" a=") + ggml_op_name(node->src[0]->op); + if (node->src[0]->op == GGML_OP_UNARY) { name += std::string(":") + ggml_unary_op_name(ggml_get_unary_op(node->src[0])); } + name += std::string(" b=") + ggml_op_name(node->src[1]->op); + if (node->src[1]->op == GGML_OP_UNARY) { name += std::string(":") + ggml_unary_op_name(ggml_get_unary_op(node->src[1])); } + if (node->src[1]->op == GGML_OP_RESHAPE && node->src[1]->src[0]) { + name += std::string("(") + ggml_op_name(node->src[1]->src[0]->op) + ")"; + } + return fusion_str + name; + } if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { const uint64_t m = node->ne[0]; const uint64_t n = node->ne[1]; @@ -4353,6 +4369,23 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { l_warptile = { 256, 128, 128, 16, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 }; l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 }; l_warptile_mmq_int_k = { 256, 128, 128, 32, mm_warp_16, 64, 1, 4, 2, 1, mm_warp_16 }; + + // EXPERIMENT (GGML_VK_MMID_WG256=1): the dense large tile above runs 256 threads on a + // 128x128 tile, but the mul_mat_id variants still run 128. Give MoE the same thread + // count per tile: same BM/BN/BK (so same shared memory), twice the threads sharing each + // A/B tile load, half the accumulators per thread. Warp split stays legal: + // (BM/WM)*(BN/WN) == wg/subgroup == 4, WNITER == (WM*WN)/(WARP*TM*TN*WMITER) == 2. + // A 256-expert MoE at ub=2048 sees only ~64 rows per expert, so the tile that actually + // runs is the medium one, not the large one. Override both. + static const char * mmid_wg256_env = getenv("GGML_VK_MMID_WG256"); + if (mmid_wg256_env && atoi(mmid_wg256_env) != 0) { + l_warptile_mmqid = { 256, 128, 128, 32, mul_mat_subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_8 }; + l_warptile_mmqid_int = { 256, 128, 128, 32, mul_mat_subgroup_size_8, 64, 2, 4, 4, 1, mul_mat_subgroup_size_8 }; + // BM=BN=64 at 4 warps needs WM=WN=32: (BM/WM)*(BN/WN) == 4, cms_per_row/col == 2. + m_warptile_mmqid = { 256, 64, 64, 32, 32, 32, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_8 }; + m_warptile_mmqid_int = { 256, 64, 64, 32, 32, 32, 2, 2, 2, 1, mul_mat_subgroup_size_8 }; + fprintf(stderr, "ggml_vulkan: MUL_MAT_ID medium+large tiles at 256 threads (GGML_VK_MMID_WG256)\n"); + } } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { // Xe2/Xe3 with coopmat enabled - warptile performance tuning l_warptile = { 512, 128, 128, 16, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 }; @@ -4653,7 +4686,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { return spec; }; - const int mul_mat_id_param_count = 5; + const int mul_mat_id_param_count = 6; // a, b, d, ids, expert_counts, fused scale #if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) if (device->coopmat2) { @@ -5669,6 +5702,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_concat_i8, "concat_i8", concat_i8_len, concat_i8_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_concat_i16, "concat_i16", concat_i16_len, concat_i16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_concat_i32, "concat_i32", concat_i32_len, concat_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + // One workgroup per 32x32 tile: elements are passed as (rows, cols, 1). + ggml_vk_create_pipeline(device, device->pipeline_concat_transpose_i32, "concat_transpose_i32", concat_transpose_i32_len, concat_transpose_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {32, 32, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_concat_i64, "concat_i64", concat_i64_len, concat_i64_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_upscale_nearest_f32, "upscale_f32", upscale_f32_len, upscale_f32_data, "main", 2, sizeof(vk_op_upscale_push_constants), {512, 1, 1}, {GGML_SCALE_MODE_NEAREST}, 1); @@ -8982,14 +9017,14 @@ static void ggml_vk_matmul_id( uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11, - bool hoist_row_ids) { + bool hoist_row_ids, const vk_subbuffer & fused_scale, uint32_t fusion_flags) { VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " << "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, - nei0, nei1, nbi1, ne11, n_as, uint32_t(hoist_row_ids) }; - ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as }); + nei0, nei1, nbi1, ne11, n_as, uint32_t(hoist_row_ids), fusion_flags }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf, fused_scale }, pc, { m, nei1, n_as }); } static bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor) { @@ -10143,7 +10178,7 @@ static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, c } } -static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { +static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_tensor * fused_scale = nullptr, ggml_tensor * fused_dst = nullptr) { VK_LOG_DEBUG("ggml_vk_mul_mat_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << 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=" << 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 << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; @@ -10181,7 +10216,9 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& hoisted_row_id_words * sizeof(uint32_t) <= ctx->device->properties.limits.maxStorageBufferRange; - ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + // When the following MUL is fused in, write the scaled result straight to its destination. + const ggml_tensor * out_dst = fused_dst ? fused_dst : dst; + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)out_dst->buffer->context; ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; @@ -10269,6 +10306,18 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type); + // PROBE (GGML_VK_MMID_PROBE=1): which mmid tile actually runs, and with how many threads. + static const char * mmid_probe_env = getenv("GGML_VK_MMID_PROBE"); + if (mmid_probe_env && atoi(mmid_probe_env) != 0) { + static std::set seen; + std::string key = pipeline->name + ":" + std::to_string(n_for_tile); + if (seen.insert(key).second) { + fprintf(stderr, "ggml_vulkan: mmid pipeline=%s n_for_tile=%u m=%u wg=(%u,%u,%u)\n", + pipeline->name.c_str(), n_for_tile, (uint32_t)ne01, + pipeline->wg_denoms[0], pipeline->wg_denoms[1], pipeline->wg_denoms[2]); + } + } + if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) { pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline); } @@ -10359,7 +10408,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } vk_buffer d_D = dst_buf_ctx->dev_buffer; - const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + const uint64_t d_buf_offset = vk_tensor_offset(out_dst) + out_dst->view_offs; GGML_ASSERT(d_D != nullptr); vk_buffer d_X; uint64_t x_buf_offset = 0; @@ -10495,7 +10544,9 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& { d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf, ne01, ne21, ne10, ne10, stride_b_y, ne01, stride_batch_x, stride_batch_y, ne20*ne21, - n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, hoist_row_ids + n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, hoist_row_ids, + fused_scale ? ggml_vk_tensor_subbuffer(ctx, fused_scale) : vk_subbuffer{ d_D, d_buf_offset, d_sz }, + fused_scale ? 1u : 0u ); // NOLINT if (x_non_contig || qx_needs_dequant) { @@ -10759,7 +10810,16 @@ static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx if (ggml_vk_use_mul_mat_vec_id(cgraph, node_idx)) { ggml_vk_mul_mat_vec_id_q_f16(ctx, subctx, cgraph, node_idx); } else { - ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst); + // Fused scale epilogue: the MUL's other operand is applied as the matmul writes out, + // and the result goes straight to the MUL's destination. + const ggml_tensor * fused_scale = nullptr; + ggml_tensor * fused_dst = nullptr; + if (ctx->num_additional_fused_ops == 1) { + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + fused_scale = (mul->src[0] == dst) ? mul->src[1] : mul->src[0]; + fused_dst = mul; + } + ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst, fused_scale, fused_dst); } } @@ -11252,6 +11312,31 @@ static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, u } } +// EXPERIMENT (GGML_VK_CONCAT_TRANSPOSE=1): the delta-net conv-state path does +// ggml_transpose() straight into a dim-0 ggml_concat(), so the generic concat kernel reads +// src1 fully de-coalesced. Measured on Qwen3.6-35B-A3B: CONCAT is ~22% of pp2048 at ub=2048 +// and grows 3.1x for a 2x ubatch. Route that exact shape to a tiled-transpose kernel. +static bool ggml_vk_concat_is_transposed(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { + static const char * env = getenv("GGML_VK_CONCAT_TRANSPOSE"); + if (!(env && atoi(env) != 0)) { + return false; + } + if (ggml_get_op_params_i32(dst, 0) != 0) { // dim 0 only + return false; + } + if (src0->ne[2] != 1 || src0->ne[3] != 1 || src1->ne[2] != 1 || src1->ne[3] != 1) { + return false; + } + const size_t ts = ggml_type_size(src0->type); + if (src0->nb[0] != ts || dst->nb[0] != ts) { // src0 and dst rows must be contiguous + return false; + } + if (src1->nb[0] <= src1->nb[1]) { // src1 must actually be transposed + return false; + } + return src0->ne[1] == src1->ne[1] && dst->ne[1] == src1->ne[1]; +} + static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * dst, ggml_op op) { switch (op) { case GGML_OP_GET_ROWS: @@ -11344,6 +11429,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const if (!ggml_vk_concat_supported(src0, src1, dst)) { return nullptr; } + // Tiled-transpose path handles unquantized 4-byte elements only. + if (!ggml_is_quantized(src0->type) && ggml_vk_concat_unit_size(src0->type) == 4 && + ggml_vk_concat_is_transposed(src0, src1, dst)) { + return ctx->device->pipeline_concat_transpose_i32; + } switch (ggml_vk_concat_unit_size(src0->type)) { case 1: return ctx->device->pipeline_concat_i8; @@ -12368,6 +12458,11 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co case GGML_OP_GLU: case GGML_OP_CONV_2D_DW: { + // The tiled concat kernel is dispatched per 32x32 tile, not per element. + if (op == GGML_OP_CONCAT && pipeline == ctx->device->pipeline_concat_transpose_i32) { + elements = { (uint32_t)src1->ne[1], (uint32_t)src1->ne[0], 1 }; + break; + } uint32_t ne = ggml_nelements(dst); if (op == GGML_OP_CPY && ggml_is_quantized(src0->type) && ggml_is_quantized(dst->type)) { // Convert from number of logical elements to 2- or 4-byte units. @@ -15844,6 +15939,22 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; } + // Fused silu(x)*y: run it as a swiglu split, writing straight to the MUL's destination. + if (ctx->num_additional_fused_ops == 1) { + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + ggml_tensor * other = (mul->src[0] == node) ? mul->src[1] : mul->src[0]; + + ggml_tensor fused = *mul; + fused.op = GGML_OP_GLU; + memset(fused.op_params, 0, sizeof(fused.op_params)); + ggml_set_op_params_i32(&fused, 0, (int32_t) GGML_GLU_OP_SWIGLU); + fused.src[0] = node->src[0]; + fused.src[1] = other; + + ggml_vk_glu(ctx, compute_ctx, fused.src[0], fused.src[1], &fused); + break; + } + switch (ggml_get_unary_op(node)) { case GGML_UNARY_OP_ELU: case GGML_UNARY_OP_EXP: @@ -16860,15 +16971,61 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g } } + // EXPERIMENT (GGML_VK_FUSE_UNARY_MUL=1): silu(x)*y is emitted as two nodes by the delta-net + // path, so the silu result makes a full round trip through memory. That is the same shape + // swiglu-split already computes in one pass, so route the pair to the existing GLU pipeline. + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_MUL) { + static const char * env = getenv("GGML_VK_FUSE_UNARY_MUL"); + if (!(env && atoi(env) != 0)) { + return false; + } + const ggml_tensor * unary = cgraph->nodes[node_idx]; + const ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + + if (ggml_get_unary_op(unary) != GGML_UNARY_OP_SILU) { + return false; + } + if (mul->src[0] != unary && mul->src[1] != unary) { + return false; + } + const ggml_tensor * other = (mul->src[0] == unary) ? mul->src[1] : mul->src[0]; + // The GLU split shader walks both inputs and the output with the same element count. + if (unary->type != GGML_TYPE_F32 || other->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_are_same_shape(unary, other) || !ggml_are_same_shape(unary, mul)) { + return false; + } + if (!ggml_is_contiguous(unary->src[0]) || !ggml_is_contiguous(other) || !ggml_is_contiguous(mul)) { + return false; + } + return true; + } + auto const &mmid_mul_ok = [&](const ggml_tensor *mmid, const ggml_tensor *mul) { const ggml_tensor *scale = mul->src[1]; if (mmid != mul->src[0]) { return false; } - // mat-vec only + // EXPERIMENT (GGML_VK_MMID_SCALE_EPILOGUE=1): the tile shader can apply the scale as it + // writes out, which removes a full write+read of the matmul result at prefill. The + // coopmat2 shader has the binding but not the epilogue, so it stays on the old path. if (!ggml_vk_use_mul_mat_vec_id(cgraph, node_idx)) { - return false; + static const char * env = getenv("GGML_VK_MMID_SCALE_EPILOGUE"); + if (!(env && atoi(env) != 0) || ctx->device->coopmat2) { + return false; + } + // Shader indexes the scale as [token * nei0 + expert_slot]. + if (scale->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale)) { + return false; + } + if (get_misalign_bytes(ctx, scale) != 0) { + return false; + } + return scale->ne[0] == 1 && + scale->ne[1] == mmid->ne[1] && scale->ne[2] == mmid->ne[2] && scale->ne[3] == mmid->ne[3] && + ggml_are_same_shape(mul, mmid); } // shaders assume the types match if (mmid->type != scale->type) { @@ -17524,6 +17681,9 @@ 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(ctx, cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL })) { + ctx->num_additional_fused_ops = 1; + fusion_string = "SILU_MUL"; } 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) && diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/concat_transpose.comp b/ggml/src/ggml-vulkan/vulkan-shaders/concat_transpose.comp new file mode 100644 index 000000000000..653aeaa011b2 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/concat_transpose.comp @@ -0,0 +1,43 @@ +#version 450 + +#include "types.glsl" +#include "generic_binary_head.glsl" + +layout(local_size_x = 32, local_size_y = 8, local_size_z = 1) in; + +// dim-0 concat whose src1 is transposed. The generic kernel reads src1 with the transposed +// stride, so neighbouring lanes touch different cache lines. Stage a 32x32 tile in shared +// memory instead, which keeps both the load and the store coalesced. 33 columns pads away +// the shared-memory bank conflicts. +shared A_TYPE tmp[32][33]; + +void main() { + const uint tx = gl_LocalInvocationID.x; + const uint ty = gl_LocalInvocationID.y; + + const uint row = gl_WorkGroupID.x * 32 + tx; + + // src0 is already contiguous, copy it straight through. + if (gl_WorkGroupID.y == 0 && row < p.ne01) { + for (uint i0 = ty; i0 < p.ne00; i0 += 8) { + data_d[get_doffset() + row*p.nb21 + i0*p.nb20] = D_TYPE(data_a[get_aoffset() + row*p.nb01 + i0*p.nb00]); + } + } + + [[unroll]] for (uint j = 0; j < 32; j += 8) { + const uint c = gl_WorkGroupID.y * 32 + ty + j; + if (c < p.ne10 && row < p.ne11) { + tmp[ty + j][tx] = A_TYPE(data_b[get_boffset() + c*p.nb10 + row*p.nb11]); + } + } + + barrier(); + + const uint col = gl_WorkGroupID.y * 32 + tx; + [[unroll]] for (uint j = 0; j < 32; j += 8) { + const uint r = gl_WorkGroupID.x * 32 + ty + j; + if (col < p.ne10 && r < p.ne11) { + data_d[get_doffset() + r*p.nb21 + (p.ne00 + col)*p.nb20] = D_TYPE(tmp[tx][ty + j]); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 63c4aaebcb1a..7d32124b969b 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -68,6 +68,9 @@ layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; #ifdef MUL_MAT_ID layout (binding = 3) readonly buffer IDS {int data_ids[];}; layout (binding = 4) readonly buffer Counts {int data_expert_count[];}; +// Fused MUL epilogue: one scale per (expert slot, token), broadcast down the M dimension. +// Always bound; p.fusion_flags == 0 means ignore it. +layout (binding = 5) readonly buffer FusedScale {float data_fscale[];}; #endif layout (push_constant) uniform parameter @@ -90,6 +93,7 @@ layout (push_constant) uniform parameter uint ne11; uint n_experts; uint hoist_row_ids; + uint fusion_flags; #else uint base_work_group_z; uint num_batches; @@ -394,9 +398,13 @@ void main() { if (row_i >= _ne1) break; const u16vec2 row_idx = row_ids[row_i - ic * BN]; - if (dr + cm_row * TM + store_r < p.M) { - data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + const uint didx = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r; + if (p.fusion_flags != 0) { + data_d[didx] = D_TYPE(float(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]) * data_fscale[row_idx.y * p.nei0 + row_idx.x]); + } else { + data_d[didx] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + } } } barrier(); @@ -449,15 +457,19 @@ void main() { if (row_i >= _ne1) break; const u16vec2 row_idx = row_ids[row_i - ic * BN]; + const bool do_scale = p.fusion_flags != 0; + const float fscale = do_scale ? data_fscale[row_idx.y * p.nei0 + row_idx.x] : 1.0f; #endif // MUL_MAT_ID [[unroll]] for (uint cr = 0; cr < TM / 2; cr++) { const uint sums_idx = (wsic * TN + cc) * WMITER * (TM / 2) + wsir * (TM / 2) + cr; #ifdef MUL_MAT_ID if (dr_warp + 2 * cr < p.M) { - data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr] = D_TYPE(sums[sums_idx].x); + const uint didx = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr; + data_d[didx] = do_scale ? D_TYPE(float(sums[sums_idx].x) * fscale) : D_TYPE(sums[sums_idx].x); } if (dr_warp + 2 * cr + 1 < p.M) { - data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr + 1] = D_TYPE(sums[sums_idx].y); + const uint didx = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr + 1; + data_d[didx] = do_scale ? D_TYPE(float(sums[sums_idx].y) * fscale) : D_TYPE(sums[sums_idx].y); } #else if (dr_warp + 2 * cr < p.M && dc_warp + cc < p.N) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp index 27f3178e7f26..9bfad031d1b6 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -106,6 +106,9 @@ layout (binding = 1) readonly buffer B4 {B_TYPEV4 data_b_v4[];}; #ifdef MUL_MAT_ID layout (binding = 3) readonly buffer IDS {int data_ids[];}; layout (binding = 4) readonly buffer Counts {int data_expert_count[];}; +// Bound for descriptor-layout parity with the other mul_mat_id shaders. The fused MUL +// epilogue is not implemented here, so the host never enables it on coopmat2 devices. +layout (binding = 5) readonly buffer FusedScale {float data_fscale[];}; shared u16vec4 row_ids[BN]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp index 1fbcbf6c9332..b36d056c1c9c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp @@ -36,6 +36,8 @@ layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; #ifdef MUL_MAT_ID layout (binding = 3) readonly buffer IDS {int data_ids[];}; layout (binding = 4) readonly buffer Counts {int data_expert_count[];}; +// Fused MUL epilogue, see mul_mm.comp. +layout (binding = 5) readonly buffer FusedScale {float data_fscale[];}; #endif layout (push_constant) uniform parameter @@ -58,6 +60,7 @@ layout (push_constant) uniform parameter uint ne11; uint n_experts; uint hoist_row_ids; + uint fusion_flags; #else uint base_work_group_z; uint num_batches; @@ -303,7 +306,10 @@ void main() { const uint sums_idx = (wsic * TN + cc) * WMITER * TM + wsir * TM + cr; #ifdef MUL_MAT_ID if (dr_warp + cr < p.M) { - data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[sums_idx].x); + const uint didx = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr; + data_d[didx] = p.fusion_flags != 0 + ? D_TYPE(float(sums[sums_idx].x) * data_fscale[row_idx.y * p.nei0 + row_idx.x]) + : D_TYPE(sums[sums_idx].x); } #else if (dr_warp + cr < p.M && dc_warp + cc < p.N) { 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 d375c2d12771..5487f76b094d 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -902,6 +902,7 @@ void process_shaders() { string_to_spv("concat_i16", "concat.comp", {{"A_TYPE", "uint16_t"}, {"B_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); string_to_spv("concat_i32", "concat.comp", {{"A_TYPE", "uint"}, {"B_TYPE", "uint"}, {"D_TYPE", "uint"}}); string_to_spv("concat_i64", "concat.comp", {{"A_TYPE", "uvec2"}, {"B_TYPE", "uvec2"}, {"D_TYPE", "uvec2"}}); + string_to_spv("concat_transpose_i32", "concat_transpose.comp", {{"A_TYPE", "uint"}, {"B_TYPE", "uint"}, {"D_TYPE", "uint"}}); string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); From 7a1eb6d85b37fa8342e1b529fd76adc3cf742a40 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 6 Aug 2026 08:48:36 +0000 Subject: [PATCH 069/109] vulkan: reject ne[3] > 1 in the mul_mat_id scale epilogue The fused epilogue derives the scale index from row_ids as [token * nei0 + expert_slot], which carries no 4th dimension, but the gate admitted any ne[3] as long as the scale and the matmul agreed. A tensor with ne[3] > 1 would read the wrong scale for every batch past the first and return quietly wrong results. test-backend-ops never generates such a case, so the suite passed throughout; found by reading the gate against the shader. Co-Authored-By: Claude Opus 5 (cherry picked from commit 016e906788057b9734ab7727de645d58f8080716) --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index bfe1c5a93ed0..23cd87e5693d 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17023,8 +17023,10 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g if (get_misalign_bytes(ctx, scale) != 0) { return false; } - return scale->ne[0] == 1 && - scale->ne[1] == mmid->ne[1] && scale->ne[2] == mmid->ne[2] && scale->ne[3] == mmid->ne[3] && + // The shader indexes the scale as [token * nei0 + expert_slot] from row_ids, which + // carries no 4th dimension, so ne[3] must be 1 or later batches read the wrong scale. + return scale->ne[0] == 1 && mmid->ne[3] == 1 && scale->ne[3] == 1 && + scale->ne[1] == mmid->ne[1] && scale->ne[2] == mmid->ne[2] && ggml_are_same_shape(mul, mmid); } // shaders assume the types match From 0d542b834db3670d4da53a68d5084418c52df961 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Mon, 17 Aug 2026 06:13:27 +0000 Subject: [PATCH 070/109] vulkan : default the transposed-concat path on The delta-net conv-state path transposes straight into a dim-0 concat, so the generic concat kernel walks src1 with a conv_channels * 4 byte stride. On qwen35 that is 40960 B, which is 160 * 256 B with 160 % 16 == 0, so every read lands on the same one of the 16 memory channels: 13.7 GB/s against 138.9 GB/s for the tiled path. The tiled-transpose route has been behind GGML_VK_CONCAT_TRANSPOSE=1 since it landed. Turn it on by default and keep GGML_VK_CONCAT_TRANSPOSE=0 as the opt-out. Qwen3.8-27B pp2048: +0.4% at ub 256, +4.7% at ub 1024, +7.2% at ub 2048. Assisted-by: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 23cd87e5693d..7f5f4bed26a4 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11312,13 +11312,15 @@ static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, u } } -// EXPERIMENT (GGML_VK_CONCAT_TRANSPOSE=1): the delta-net conv-state path does -// ggml_transpose() straight into a dim-0 ggml_concat(), so the generic concat kernel reads -// src1 fully de-coalesced. Measured on Qwen3.6-35B-A3B: CONCAT is ~22% of pp2048 at ub=2048 -// and grows 3.1x for a 2x ubatch. Route that exact shape to a tiled-transpose kernel. +// The delta-net conv-state path does ggml_transpose() straight into a dim-0 ggml_concat(), so +// the generic concat kernel walks src1 with a conv_channels * 4 byte stride. On qwen35 that is +// 40960 B = 160 * 256 B and 160 % 16 == 0, so every read lands on one of the 16 memory channels: +// 13.7 GB/s against 138.9 GB/s for the tiled path. Route that exact shape to a tiled-transpose +// kernel. On by default; GGML_VK_CONCAT_TRANSPOSE=0 opts out. +// Qwen3.8-27B pp2048: +0.4% at ub 256, +4.7% at ub 1024, +7.2% at ub 2048. static bool ggml_vk_concat_is_transposed(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { static const char * env = getenv("GGML_VK_CONCAT_TRANSPOSE"); - if (!(env && atoi(env) != 0)) { + if (env && env[0] == '0') { return false; } if (ggml_get_op_params_i32(dst, 0) != 0) { // dim 0 only From ce93d296ffd884b2d052743a93ddcfc904b9e72f Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 2 Aug 2026 00:13:47 +0000 Subject: [PATCH 071/109] vulkan: flush pending compute ctx before perf logger timestamps The scheduler's async input copies between graph splits land in the compute ctx on devices without a separate transfer queue, so the perf logger's fresh-ctx assert fired under partial offload (--n-cpu-moe). Assisted-by: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8fbb1359f406..8b9408305342 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17390,6 +17390,12 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg std::fill(ctx->query_nodes.begin(), ctx->query_nodes.end(), nullptr); std::fill(ctx->query_node_idx.begin(), ctx->query_node_idx.end(), 0); + // Under partial offload the scheduler's async input copies between graph + // splits can leave commands in a pending compute ctx. Flush it so the + // timestamp stream starts on a fresh command buffer. + if (!ctx->compute_ctx.expired()) { + ggml_vk_synchronize(ctx); + } GGML_ASSERT(ctx->compute_ctx.expired()); compute_ctx = ggml_vk_get_compute_ctx(ctx); ctx->query_idx = 0; From 8a8c2187bd9c8fec6b314d10f9658eccd22695e1 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 2 Aug 2026 11:56:37 +0000 Subject: [PATCH 072/109] vulkan: bound command buffers by memory traffic, not just flops Nodes with no flops estimate (large copies, set_rows, mask fills) can pack a command buffer whose execution time grows with context length until it exceeds the amdgpu ring timeout (10s on the compute ring), causing the ring resets and DeviceLost reported at long context. Add a bytes-per-submit cap (default 8 GiB, GGML_VK_MAX_MB_PER_SUBMIT to override, 0 disables) alongside the existing flops and node-count gates. Assisted-by: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 34 ++++++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8b9408305342..11846bda462b 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -837,6 +837,7 @@ struct vk_device_struct { bool add_rms_fusion; uint32_t partials_binding_alignment; uint32_t max_nodes_per_submit; + uint64_t max_bytes_per_submit; bool shader_64b_indexing; @@ -2166,6 +2167,20 @@ static bool vk_enable_sync_logger = false; static uint32_t vk_perf_logger_frequency = 1; static std::string vk_pipeline_stats_filter; +// Total memory traffic of a node (dst + srcs). Used to bound command buffer +// execution time for bandwidth-bound ops with no flops estimate (large copies, +// set_rows, mask fills at long context) - packing too many of them into one +// submission can exceed the driver timeout. +static uint64_t ggml_vk_get_node_bytes(const ggml_tensor * node) { + uint64_t bytes = ggml_nbytes(node); + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (node->src[i]) { + bytes += ggml_nbytes(node->src[i]); + } + } + return bytes; +} + static uint64_t ggml_vk_get_node_flops(const ggml_tensor * node) { if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { const uint64_t m = node->ne[0]; @@ -6514,6 +6529,15 @@ static vk_device ggml_vk_get_device(size_t idx) { device->max_nodes_per_submit = std::max(max_nodes_per_submit, 1u); } + // Also submit once a batch has accumulated enough memory traffic, so that + // bandwidth-bound nodes with no flops estimate cannot grow a command buffer + // past the driver timeout. 0 disables the limit. + device->max_bytes_per_submit = 8ull * 1024 * 1024 * 1024; + const char* GGML_VK_MAX_MB_PER_SUBMIT = getenv("GGML_VK_MAX_MB_PER_SUBMIT"); + if (GGML_VK_MAX_MB_PER_SUBMIT != nullptr) { + device->max_bytes_per_submit = std::stoull(GGML_VK_MAX_MB_PER_SUBMIT) * 1024 * 1024; + } + const bool force_disable_f16 = getenv("GGML_VK_DISABLE_F16") != nullptr; device->fp16 = !force_disable_f16 && fp16_storage && fp16_compute; @@ -17422,6 +17446,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg uint32_t submitted_nodes = 0; uint32_t submit_count = 0; uint64_t batch_flops = 0; + uint64_t batch_bytes = 0; uint64_t total_flops = 0; uint64_t flops_cap = 200'000'000'000ULL; @@ -17459,6 +17484,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg first_node_in_batch = true; submitted_nodes = 0; batch_flops = 0; + batch_bytes = 0; if (submit_count < 3) { flops_per_submit *= 2; } @@ -17473,9 +17499,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg { auto node_flops = ggml_vk_get_node_flops(cgraph->nodes[i]); total_flops += node_flops; + auto node_bytes = ggml_vk_get_node_bytes(cgraph->nodes[i]); - // Flush the current batch before recording a node that would push it over the flop threshold - if (flops_per_submit != 0 && submitted_nodes > 0 && batch_flops + node_flops >= flops_per_submit) { + // Flush the current batch before recording a node that would push it over the flop or byte threshold + if ((flops_per_submit != 0 && submitted_nodes > 0 && batch_flops + node_flops >= flops_per_submit) || + (ctx->device->max_bytes_per_submit != 0 && submitted_nodes > 0 && batch_bytes + node_bytes >= ctx->device->max_bytes_per_submit)) { vk_context flush_ctx = ggml_vk_get_compute_ctx(ctx); ggml_vk_ctx_end(flush_ctx); flush_ctx->exit_tensor_idx = -1; @@ -17486,6 +17514,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } batch_flops += node_flops; + batch_bytes += node_bytes; } // op_srcs_fused_elementwise indicates whether an op's srcs all contribute to @@ -17708,6 +17737,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg bool almost_ready = (cgraph->n_nodes - i) < cgraph->n_nodes / 5; bool submit = (submitted_nodes >= ctx->device->max_nodes_per_submit) || (flops_per_submit != 0 && batch_flops >= flops_per_submit) || + (ctx->device->max_bytes_per_submit != 0 && batch_bytes >= ctx->device->max_bytes_per_submit) || (i + ctx->num_additional_fused_ops >= last_node) || (almost_ready && !ctx->almost_ready_fence_pending); From 9125bb2531e9fee7e952974ffc6e653f74622015 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 9 Aug 2026 02:48:56 +0000 Subject: [PATCH 073/109] ggml: cut backend splits on the input constant, not the grown capacity #22789 replaced the fixed 30-entry split input array with a growable one and, in the same edit, changed the split-cutting heuristic from the constant to split->inputs_capacity: - if (split->n_inputs == GGML_SCHED_MAX_SPLIT_INPUTS) { + if (split->n_inputs >= split->inputs_capacity) { inputs_capacity starts at GGML_SCHED_MAX_SPLIT_INPUTS but doubles on demand and is never reset for the life of the sched, so once a split slot grows, the scheduler stops cutting there and the cut point ratchets up for every later graph build. Longer splits mean every cross-backend input copy is materialised at the split's start and stays live to its last use inside it, which raises the peak the compute-buffer allocator has to cover - n_copies times over under pipeline parallelism. Only multi-backend configurations can reach this. Keep the growable array, which is what fixes the original >30-input assert, and cut on the constant again as before #22789. >= rather than == so the check keeps firing for splits that did have to grow. DeepSeek-V4-Flash UD-IQ3_XXS, gfx1151, -c 400000 -ub 2048 -fa 1 --fit off: Vulkan0 compute buffer 4714.00 MiB and 9157 graph nodes, byte-identical to the unpatched tree, and neither run grows a split past 30 inputs. Expected - one Vulkan device plus the CPU backend cannot exercise the path on this box. The reported case is 3 devices with pipeline parallelism. test-backend-ops -o FLASH_ATTN_EXT, run alone on gfx1151: 13257/13295 on both this and the unpatched tree, with the same 38 failing cases (identical case list, all type_K=q8_0 prec=def kv_view=1). Pre-existing on the branch, not touched by this change. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-backend.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index fec7d7c92bf3..c3a0a12889b5 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1339,7 +1339,10 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra } // check if the split has too many inputs // FIXME: count the number of inputs instead of only checking when full - if (split->n_inputs >= split->inputs_capacity) { + // cut on the constant, not on inputs_capacity: capacity doubles on demand and + // is never reset, so using it lets the cut point drift up and keeps every input + // copy of an ever-longer split live at once + if (split->n_inputs >= GGML_SCHED_MAX_SPLIT_INPUTS) { const size_t id = hash_id(src); int src_backend_id = sched->hv_tensor_backend_ids[id]; bool supported = ggml_backend_sched_buffer_supported(sched, src, cur_backend_id); From e4615cc84b42f4b167fb15d55b59d1d15379d490 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Mon, 17 Aug 2026 06:13:14 +0000 Subject: [PATCH 074/109] vulkan : take the copy path for bulk UMA reads from uncached mappings ggml_vk_buffer_read_2d took the direct-CPU-read path whenever the buffer was host-visible on a UMA device. That holds only for host-cached mappings. A write-combined mapping (host-visible without HOST_CACHED, which is what amdgpu hands out for GTT) reads back at uncached speed, around 200 MB/s, so every bulk read crawled. Gate the direct path on HOST_CACHED, and keep small reads direct regardless so they do not pay the fence round-trip. Measured on hybrid-attention context checkpoints, where this cost about 600 ms per prompt. Assisted-by: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 11846bda462b..7e374274e07b 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -8699,7 +8699,12 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si // If the device is not an UMA device the memory is host-accessible through rebar. While writing // through PCIe is sufficient fast reading back data from PCIe is slower than going through // the HW device to host copy path. - if(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible && src->device->uma) { + // On UMA, direct CPU reads are only fast from host-cached mappings. Write-combined/uncached + // mappings (e.g. GTT without HOST_CACHED) read at uncached speed (~200 MB/s), so bulk reads + // must go through the device copy path. Small reads stay direct to avoid the fence round-trip. + const bool host_cached = bool(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCached); + if((src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) && src->device->uma && + (host_cached || width * height <= 64 * 1024)) { GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); std::lock_guard guard(src->device->mutex); From 69e30522fa4687ab3338d6200182eb8991c33329 Mon Sep 17 00:00:00 2001 From: Gaetan Puleo <12990773+gaetan-puleo@users.noreply.github.com> Date: Sat, 1 Aug 2026 14:57:00 +0200 Subject: [PATCH 075/109] vulkan: DeepSeek V4 lightning indexer kernels + indexed sparse FA Implements GGML_OP_LIGHTNING_INDEXER on Vulkan (scalar subgroup shader for small batches, coopmat 16x16 tiles for prefill, dedicated decode variant) and an indexed sparse flash-attention path that consumes the indexer's top-k selection directly via a new ggml_flash_attn_ext_add_top_k() API (FA src[5] + op_param[4] = n_kv_raw dense prefix), instead of attending densely over the full compressed KV. The sparse path engages only for V4's CSA shape (hd 512, 64 heads, MQA, f16 K==V latent) when dense_kv >= 3x active_kv; everything else falls through to the dense path, which stays correct because the kq_mask still carries the top-k selection. Dropped from the original: the mul_mat_id tokens-per-expert pipeline selection, which duplicates GGML_VK_MMID_SMALLN already on this branch. Originally by Gaetan Puleo (llama-cpp-nathan-toolbox-deepseek-v4-poc, branch deepseek-v4-flash-strix-halo); cherry-picked with the mmid hunk dropped. --- ggml/include/ggml.h | 5 + ggml/src/ggml-vulkan/ggml-vulkan.cpp | 170 ++++++++++++++++++ .../vulkan-shaders/flash_attn_top_k.comp | 144 +++++++++++++++ .../vulkan-shaders/lightning_indexer_cm.comp | 125 +++++++++++++ .../lightning_indexer_decode_cm.comp | 110 ++++++++++++ .../lightning_indexer_scalar64.comp | 92 ++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 7 + ggml/src/ggml.c | 14 ++ src/llama-graph.cpp | 9 +- src/llama-graph.h | 4 +- src/models/deepseek4.cpp | 2 +- tests/test-backend-ops.cpp | 10 +- 12 files changed, 687 insertions(+), 5 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 5f6774a630c0..d81b8e53050e 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2450,6 +2450,11 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * sinks); + GGML_API void ggml_flash_attn_ext_add_top_k( + struct ggml_tensor * a, + struct ggml_tensor * top_k, + int64_t n_kv_raw); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 528c112645a7..8d921f51c09a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1086,6 +1086,10 @@ struct vk_device_struct { vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT]; // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128 vk_pipeline pipeline_gated_delta_net[4][2]; + vk_pipeline pipeline_lightning_indexer_f16; + vk_pipeline pipeline_lightning_indexer_cm_f16; + vk_pipeline pipeline_lightning_indexer_decode_cm_f16; + vk_pipeline pipeline_flash_attn_top_k_f16; vk_pipeline pipeline_ssm_scan_f32_d128; vk_pipeline pipeline_ssm_scan_f32_d256; vk_pipeline pipeline_ssm_conv_f32; @@ -1906,6 +1910,29 @@ struct vk_op_gated_delta_net_push_constants { uint32_t K; }; +// push constants for the fork's wave64 f16 lightning-indexer kernels (scalar-64 + CM family) +struct vk_op_lightning_indexer_cm_push_constants { + uint32_t n_kv, n_batch, n_stream, nem3; + uint32_t nb1, nb3; + uint32_t nbq1, nbq2, nbq3; + uint32_t nbk2, nbk3; + uint32_t nbw1, nbw3; + uint32_t nbm1, nbm3; +}; +static_assert(sizeof(vk_op_lightning_indexer_cm_push_constants) <= 128); + +struct vk_op_flash_attn_top_k_push_constants { + uint32_t n_batch, n_kv, n_kv_raw, n_top_k, n_head; + uint32_t nbq1, nbq2, nbq3; + uint32_t nbk1, nbk3; + uint32_t nbm1, nbm3; + uint32_t nbt1, nbt3; + uint32_t nb1, nb2, nb3; + float scale; + uint32_t has_sinks; +}; +static_assert(sizeof(vk_op_flash_attn_top_k_push_constants) <= 128); + struct vk_op_ssm_scan_push_constants { uint32_t nb02, nb03, nb12, nb13; uint32_t nb21, nb22, nb31; @@ -5997,6 +6024,29 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } } + if (device->subgroup_arithmetic && device->subgroup_size == 64) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f16, + "lightning_indexer_f16", lightning_indexer_f16_len, lightning_indexer_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {8, 1, 1}, {device->subgroup_size}, 1, true, true, + device->subgroup_size); +#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + if (device->coopmat_support && device->coopmat_support_16x16x16_f32acc && device->subgroup_size_control) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_f16, + "lightning_indexer_cm_f16", lightning_indexer_cm_f16_len, lightning_indexer_cm_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 16, 1}, {device->subgroup_size}, 1, true, true, + device->subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_decode_cm_f16, + "lightning_indexer_decode_cm_f16", lightning_indexer_decode_cm_f16_len, lightning_indexer_decode_cm_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 1, 1}, {device->subgroup_size}, 1, true, true, + device->subgroup_size); + } +#endif + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_top_k_f16, + "flash_attn_top_k_f16", flash_attn_top_k_f16_len, flash_attn_top_k_f16_data, "main", 6, + sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true, + device->subgroup_size); + } + if (device->subgroup_arithmetic && device->subgroup_require_full_support) { ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d128, "ssm_scan_128_f32", ssm_scan_subgroup_f32_len, ssm_scan_subgroup_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {128, device->subgroup_size}, 1, true, true); ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_256_f32", ssm_scan_subgroup_f32_len, ssm_scan_subgroup_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, device->subgroup_size}, 1, true, true); @@ -10985,6 +11035,69 @@ static void ggml_vk_perf_mark_subop(ggml_backend_vk_context * ctx, vk_context& s subctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++); } +static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & subctx, + const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, + const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { + const ggml_tensor * top_k = dst->src[5]; + if (!top_k || !ctx->device->pipeline_flash_attn_top_k_f16 || + q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || + !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || + q->ne[0] != 512 || q->ne[1] < 64 || k->ne[0] != 512 || v->ne[0] != 512 || + q->ne[2] != 64 || k->ne[2] != 1 || v->ne[2] != 1 || + q->ne[1] != top_k->ne[1] || q->ne[3] != top_k->ne[3] || + k->ne[1] != v->ne[1] || k->buffer != v->buffer || k->data != v->data || + !ggml_is_contiguous(mask) || !ggml_is_contiguous(top_k)) { + return false; + } + + float scale = 0.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 (max_bias != 0.0f || logit_softcap != 0.0f) { + return false; + } + + const int32_t n_kv_raw = ggml_get_op_params_i32(dst, 4); + if (n_kv_raw < 0 || n_kv_raw > k->ne[1] || top_k->ne[0] > k->ne[1] - n_kv_raw) { + return false; + } + const int64_t n_kv_active = n_kv_raw + top_k->ne[0]; + if (k->ne[1] < 3 * n_kv_active) { + return false; + } + + const vk_op_flash_attn_top_k_push_constants pc = { + (uint32_t) q->ne[1], (uint32_t) k->ne[1], (uint32_t) n_kv_raw, + (uint32_t) top_k->ne[0], (uint32_t) q->ne[2], + (uint32_t) (q->nb[1] / sizeof(float)), + (uint32_t) (q->nb[2] / sizeof(float)), + (uint32_t) (q->nb[3] / sizeof(float)), + (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (mask->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), + (uint32_t) (top_k->nb[3] / sizeof(int32_t)), + (uint32_t) (dst->nb[1] / sizeof(float)), + (uint32_t) (dst->nb[2] / sizeof(float)), + (uint32_t) (dst->nb[3] / sizeof(float)), + scale, sinks != nullptr, + }; + + const vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q); + const vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; + vk_pipeline pipeline = ctx->device->pipeline_flash_attn_top_k_f16; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf, + ggml_vk_tensor_subbuffer(ctx, top_k), ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {(uint32_t) q->ne[1], (uint32_t) CEIL_DIV(q->ne[2], 8), (uint32_t) q->ne[3]}); + return true; +} + static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; @@ -11036,6 +11149,9 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx assert(dst->type == GGML_TYPE_F32); assert(q->type == GGML_TYPE_F32); + if (ggml_vk_flash_attn_top_k(ctx, subctx, q, k, v, mask, sinks, dst)) { + return; + } uint32_t gqa_ratio = 1; uint32_t qk_ratio = neq2 / nek2; uint32_t workgroups_x = (uint32_t)neq1; @@ -11941,6 +12057,17 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const } return nullptr; case GGML_OP_LIGHTNING_INDEXER: + // fork fast path: f16 K on wave64 subgroup-arithmetic devices routes to the tuned + // scalar-64/CM kernels; anything else falls through to the generic pipeline table + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32 && + src0->ne[0] == 128 && src0->ne[1] == 64 && src1->ne[1] == 1 && + ctx->device->pipeline_lightning_indexer_f16) { + if (ctx->device->pipeline_lightning_indexer_decode_cm_f16 && src0->ne[2] == 1) { + return ctx->device->pipeline_lightning_indexer_decode_cm_f16; + } + return ctx->device->pipeline_lightning_indexer_cm_f16 && src0->ne[2] >= 16 ? + ctx->device->pipeline_lightning_indexer_cm_f16 : ctx->device->pipeline_lightning_indexer_f16; + } // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer() if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) { return ctx->device->pipeline_lightning_indexer_f32[src1->type]; @@ -13021,6 +13148,8 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context& pc, { (uint32_t)(n_seqs * n_heads), 1, 1 }); } +static void ggml_vk_lightning_indexer_cm(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, vk_pipeline pipeline); + static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const ggml_tensor * q = dst->src[0]; const ggml_tensor * k = dst->src[1]; @@ -13030,6 +13159,14 @@ static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, q, k, w, dst, dst->op); GGML_ASSERT(pipeline != nullptr); + // the fork's wave64 f16 kernels take their own push-constant layout + if (pipeline == ctx->device->pipeline_lightning_indexer_f16 || + pipeline == ctx->device->pipeline_lightning_indexer_cm_f16 || + pipeline == ctx->device->pipeline_lightning_indexer_decode_cm_f16) { + ggml_vk_lightning_indexer_cm(ctx, subctx, dst, pipeline); + return; + } + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); const uint32_t n_kv = k->ne[2]; @@ -13127,6 +13264,39 @@ static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& s pc, { H, n_seqs, S_v }); } +static void ggml_vk_lightning_indexer_cm(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, vk_pipeline pipeline) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * w = dst->src[2]; + const ggml_tensor * m = dst->src[3]; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_lightning_indexer_cm_push_constants pc = { + (uint32_t) k->ne[2], + (uint32_t) q->ne[2], + (uint32_t) q->ne[3], + (uint32_t) m->ne[3], + (uint32_t) (dst->nb[1] / sizeof(float)), + (uint32_t) (dst->nb[3] / sizeof(float)), + (uint32_t) (q->nb[1] / sizeof(float)), + (uint32_t) (q->nb[2] / sizeof(float)), + (uint32_t) (q->nb[3] / sizeof(float)), + (uint32_t) (k->nb[2] / sizeof(ggml_fp16_t)), + (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (w->nb[1] / sizeof(float)), + (uint32_t) (w->nb[3] / sizeof(float)), + (uint32_t) (m->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (m->nb[3] / sizeof(ggml_fp16_t)), + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, q), ggml_vk_tensor_subbuffer(ctx, k), + ggml_vk_tensor_subbuffer(ctx, w), ggml_vk_tensor_subbuffer(ctx, m), + ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {(uint32_t) k->ne[2], (uint32_t) q->ne[2], (uint32_t) q->ne[3]}); +} + static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp new file mode 100644 index 000000000000..7f82266781f1 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp @@ -0,0 +1,144 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_basic : require + +layout(constant_id = 0) const uint WORKGROUP_SIZE = 512; +layout(constant_id = 1) const uint SUBGROUP_SIZE = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer MaskBuf { float16_t data_m[]; }; +layout(binding = 3) readonly buffer SinkBuf { float data_s[]; }; +layout(binding = 4) readonly buffer TopBuf { int data_top[]; }; +layout(binding = 5) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_batch; + uint n_kv; + uint n_kv_raw; + uint n_top_k; + uint n_head; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk1; + uint nbk3; + uint nbm1; + uint nbm3; + uint nbt1; + uint nbt3; + uint nb1; + uint nb2; + uint nb3; + float scale; + uint has_sinks; +} p; + +const uint HEAD_SIZE = 512; +const uint HEADS_PER_GROUP = 8; +const uint KEYS_PER_BLOCK = 16; + +shared float16_t key_sh[KEYS_PER_BLOCK * HEAD_SIZE]; +shared uint key_idx[KEYS_PER_BLOCK]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint head = gl_WorkGroupID.y * HEADS_PER_GROUP + gl_SubgroupID; + const uint token = gl_WorkGroupID.x; + const uint stream = gl_WorkGroupID.z; + + if (token >= p.n_batch || head >= p.n_head) { + return; + } + + float accum[HEAD_SIZE / SUBGROUP_SIZE]; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + accum[i] = 0.0; + } + + float row_max = uintBitsToFloat(0xff800000); + float row_sum = 0.0; + const uint q_base = stream * p.nbq3 + head * p.nbq2 + token * p.nbq1; + const uint mask_base = stream * p.nbm3 + token * p.nbm1; + const uint top_base = stream * p.nbt3 + token * p.nbt1; + const uint total_keys = p.n_kv_raw + p.n_top_k; + + for (uint kb = 0; kb < total_keys; kb += KEYS_PER_BLOCK) { + if (tid < KEYS_PER_BLOCK) { + const uint selected = kb + tid; + uint key = p.n_kv; + if (selected < p.n_kv_raw) { + key = selected; + } else if (selected < total_keys) { + const int compressed = data_top[top_base + selected - p.n_kv_raw]; + if (compressed >= 0 && uint(compressed) < p.n_kv - p.n_kv_raw) { + key = p.n_kv_raw + uint(compressed); + } + } + key_idx[tid] = key; + } + barrier(); + + for (uint idx = tid; idx < KEYS_PER_BLOCK * HEAD_SIZE; idx += WORKGROUP_SIZE) { + const uint col = idx / HEAD_SIZE; + const uint dim = idx % HEAD_SIZE; + const uint key = key_idx[col]; + key_sh[idx] = key < p.n_kv ? data_k[stream * p.nbk3 + key * p.nbk1 + dim] : float16_t(0.0); + } + barrier(); + + [[unroll]] for (uint col = 0; col < KEYS_PER_BLOCK; ++col) { + const uint selected = kb + col; + const uint key = key_idx[col]; + if (selected >= total_keys || key >= p.n_kv) { + continue; + } + + float partial = 0.0; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + const uint dim = lane + i * SUBGROUP_SIZE; + partial += data_q[q_base + dim] * float(key_sh[col * HEAD_SIZE + dim]); + } + const float mask = float(data_m[mask_base + key]); + const float score = subgroupAdd(partial) * p.scale + mask; + if (mask < -65500.0) { + continue; + } + + const float new_max = max(row_max, score); + const float old_scale = row_sum == 0.0 ? 0.0 : exp(row_max - new_max); + const float value_scale = exp(score - new_max); + row_sum = row_sum * old_scale + value_scale; + row_max = new_max; + + [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + const uint dim = lane + i * SUBGROUP_SIZE; + accum[i] = accum[i] * old_scale + value_scale * float(key_sh[col * HEAD_SIZE + dim]); + } + } + barrier(); + } + + if (p.has_sinks != 0) { + const float sink = data_s[head]; + const float new_max = max(row_max, sink); + const float old_scale = row_sum == 0.0 ? 0.0 : exp(row_max - new_max); + const float sink_scale = exp(sink - new_max); + row_sum = row_sum * old_scale + sink_scale; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + accum[i] *= old_scale; + } + } + + const uint dst_base = stream * p.nb3 + token * p.nb2 + head * p.nb1; + const float inv_sum = row_sum == 0.0 ? 0.0 : 1.0 / row_sum; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + data_dst[dst_base + lane + i * SUBGROUP_SIZE] = accum[i] * inv_sum; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp new file mode 100644 index 000000000000..a0a3639d2546 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp @@ -0,0 +1,125 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer WBuf { float data_w[]; }; +layout(binding = 3) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_batch; + uint n_stream; + uint nem3; + uint nb1; + uint nb3; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk2; + uint nbk3; + uint nbw1; + uint nbw3; + uint nbm1; + uint nbm3; +} p; + +const uint TILE = 16; +const uint HEAD_SIZE = 128; +const uint N_HEAD = 64; +const uint VEC_PER_HEAD = HEAD_SIZE / 4; +const uint TILE_STRIDE = VEC_PER_HEAD + 2; +const uint SCORE_STRIDE = TILE / 4 + 1; + +shared f16vec4 q_sh[TILE * TILE_STRIDE]; +shared f16vec4 k_sh[TILE * TILE_STRIDE]; +shared vec4 score_sh[TILE * SCORE_STRIDE]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint kv_base = gl_WorkGroupID.x * TILE; + const uint token_base = gl_WorkGroupID.y * TILE; + const uint stream = gl_WorkGroupID.z; + + float totals[4]; + [[unroll]] for (uint i = 0; i < 4; ++i) { + totals[i] = 0.0; + } + + for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { + const uint key = idx / VEC_PER_HEAD; + const uint d4 = idx % VEC_PER_HEAD; + const uint kv = kv_base + key; + f16vec4 value = f16vec4(0.0); + if (kv < p.n_kv) { + const uint offset = stream * p.nbk3 + kv * p.nbk2 + d4 * 4; + value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); + } + k_sh[key * TILE_STRIDE + d4] = value; + } + barrier(); + + for (uint head = 0; head < N_HEAD; ++head) { + for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { + const uint token_local = idx / VEC_PER_HEAD; + const uint d4 = idx % VEC_PER_HEAD; + const uint token = token_base + token_local; + f16vec4 value = f16vec4(0.0); + if (token < p.n_batch) { + const uint offset = stream * p.nbq3 + token * p.nbq2 + head * p.nbq1 + d4 * 4; + value = f16vec4(data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); + } + q_sh[token_local * TILE_STRIDE + d4] = value; + } + barrier(); + + coopmat scores = + coopmat(0.0); + coopmat kmat; + coopmat qmat; + + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(kmat, k_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(qmat, q_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + scores = coopMatMulAdd(kmat, qmat, scores); + } + + coopMatStore(scores, score_sh, 0, SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + barrier(); + + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint idx = tid + i * SUBGROUP_SIZE; + const uint key = idx / TILE; + const uint token_local = idx % TILE; + const uint token = token_base + token_local; + if (token < p.n_batch && kv_base + key < p.n_kv) { + const float score = score_sh[key * SCORE_STRIDE + token_local / 4][token_local % 4]; + const float weight = data_w[stream * p.nbw3 + token * p.nbw1 + head]; + totals[i] += max(score, 0.0) * weight; + } + } + barrier(); + } + + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint idx = tid + i * SUBGROUP_SIZE; + const uint key = idx / TILE; + const uint token_local = idx % TILE; + const uint kv = kv_base + key; + const uint token = token_base + token_local; + if (kv < p.n_kv && token < p.n_batch) { + const uint mask_base = (stream % p.nem3) * p.nbm3 + token * p.nbm1; + const uint dst_base = stream * p.nb3 + token * p.nb1; + data_dst[dst_base + kv] = totals[i] + float(data_m[mask_base + kv]); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp new file mode 100644 index 000000000000..fd555c76806d --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp @@ -0,0 +1,110 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer WBuf { float data_w[]; }; +layout(binding = 3) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_batch; + uint n_stream; + uint nem3; + uint nb1; + uint nb3; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk2; + uint nbk3; + uint nbw1; + uint nbw3; + uint nbm1; + uint nbm3; +} p; + +const uint TILE = 16; +const uint HEAD_SIZE = 128; +const uint N_HEAD = 64; +const uint VEC_PER_HEAD = HEAD_SIZE / 4; +const uint TILE_STRIDE = VEC_PER_HEAD + 2; +const uint SCORE_STRIDE = TILE / 4 + 1; + +shared f16vec4 q_sh[TILE * TILE_STRIDE]; +shared f16vec4 k_sh[TILE * TILE_STRIDE]; +shared vec4 score_sh[TILE * SCORE_STRIDE]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint kv_base = gl_WorkGroupID.x * TILE; + const uint token = gl_WorkGroupID.y; + const uint stream = gl_WorkGroupID.z; + + for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { + const uint key = idx / VEC_PER_HEAD; + const uint d4 = idx % VEC_PER_HEAD; + const uint kv = kv_base + key; + f16vec4 value = f16vec4(0.0); + if (kv < p.n_kv) { + const uint offset = stream * p.nbk3 + kv * p.nbk2 + d4 * 4; + value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); + } + k_sh[key * TILE_STRIDE + d4] = value; + } + barrier(); + + float total = 0.0; + for (uint head_base = 0; head_base < N_HEAD; head_base += TILE) { + for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { + const uint head_local = idx / VEC_PER_HEAD; + const uint d4 = idx % VEC_PER_HEAD; + const uint head = head_base + head_local; + const uint offset = stream * p.nbq3 + token * p.nbq2 + head * p.nbq1 + d4 * 4; + q_sh[head_local * TILE_STRIDE + d4] = + f16vec4(data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); + } + barrier(); + + coopmat scores = + coopmat(0.0); + coopmat kmat; + coopmat qmat; + + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(kmat, k_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(qmat, q_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + scores = coopMatMulAdd(kmat, qmat, scores); + } + + coopMatStore(scores, score_sh, 0, SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + barrier(); + + if (tid < TILE && kv_base + tid < p.n_kv) { + [[unroll]] for (uint head_local = 0; head_local < TILE; ++head_local) { + const float score = score_sh[tid * SCORE_STRIDE + head_local / 4][head_local % 4]; + const float weight = data_w[stream * p.nbw3 + token * p.nbw1 + head_base + head_local]; + total += max(score, 0.0) * weight; + } + } + barrier(); + } + + if (tid < TILE) { + const uint kv = kv_base + tid; + if (kv < p.n_kv) { + const uint mask_base = (stream % p.nem3) * p.nbm3 + token * p.nbm1; + const uint dst_base = stream * p.nb3 + token * p.nb1; + data_dst[dst_base + kv] = total + float(data_m[mask_base + kv]); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp new file mode 100644 index 000000000000..693bb3ece8e8 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp @@ -0,0 +1,92 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_basic : require + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer WBuf { float data_w[]; }; +layout(binding = 3) readonly buffer MaskBuf { float16_t data_m[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_batch; + uint n_stream; + uint nem3; + uint nb1; + uint nb3; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk2; + uint nbk3; + uint nbw1; + uint nbw3; + uint nbm1; + uint nbm3; +} p; + +const uint K_PER_GROUP = 8; +const uint N_HEAD = 64; + +void main() { + const uint lane = gl_SubgroupInvocationID; + const uint token = gl_WorkGroupID.y; + const uint stream = gl_WorkGroupID.z; + const uint kv_base = gl_WorkGroupID.x * K_PER_GROUP; + + if (token >= p.n_batch || stream >= p.n_stream) { + return; + } + + float k0[K_PER_GROUP]; + float k1[K_PER_GROUP]; + [[unroll]] for (uint j = 0; j < K_PER_GROUP; ++j) { + const uint kv = kv_base + j; + if (kv < p.n_kv) { + const uint k_base = stream * p.nbk3 + kv * p.nbk2; + k0[j] = float(data_k[k_base + lane]); + k1[j] = float(data_k[k_base + lane + SUBGROUP_SIZE]); + } else { + k0[j] = 0.0; + k1[j] = 0.0; + } + } + + float score[K_PER_GROUP]; + [[unroll]] for (uint j = 0; j < K_PER_GROUP; ++j) { + score[j] = 0.0; + } + + for (uint head = 0; head < N_HEAD; ++head) { + const uint q_base = stream * p.nbq3 + token * p.nbq2 + head * p.nbq1; + const float q0 = data_q[q_base + lane]; + const float q1 = data_q[q_base + lane + SUBGROUP_SIZE]; + const float weight = data_w[stream * p.nbw3 + token * p.nbw1 + head]; + + [[unroll]] for (uint j = 0; j < K_PER_GROUP; ++j) { + const float qk = subgroupAdd(q0 * k0[j] + q1 * k1[j]); + if (lane == 0) { + score[j] += max(qk, 0.0) * weight; + } + } + } + + if (lane == 0) { + const uint mask_base = (stream % p.nem3) * p.nbm3 + token * p.nbm1; + const uint dst_base = stream * p.nb3 + token * p.nb1; + [[unroll]] for (uint j = 0; j < K_PER_GROUP; ++j) { + const uint kv = kv_base + j; + if (kv < p.n_kv) { + data_dst[dst_base + kv] = score[j] + float(data_m[mask_base + kv]); + } + } + } +} 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 752380f0b7af..d3acc412e8cd 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -801,6 +801,13 @@ void process_shaders() { string_to_spv("get_rows_i32", "get_rows.comp", {{"TEMP_TYPE", "uint"}, {"A_TYPE", "uint"}, {"B_TYPE", "int"}, {"D_TYPE", "uint"}}); + string_to_spv("lightning_indexer_f16", "lightning_indexer_scalar64.comp", {}); +#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + string_to_spv("lightning_indexer_cm_f16", "lightning_indexer_cm.comp", {}); + string_to_spv("lightning_indexer_decode_cm_f16", "lightning_indexer_decode_cm.comp", {}); +#endif + string_to_spv("flash_attn_top_k_f16", "flash_attn_top_k.comp", {}); + string_to_spv("mul_mat_vec_p021_f16_f32_subgroup_add", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}); string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}); string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index e0b615c07edf..5df67fa6e11c 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5511,6 +5511,20 @@ void ggml_flash_attn_ext_add_sinks( a->src[4] = sinks; } +void ggml_flash_attn_ext_add_top_k( + struct ggml_tensor * a, + struct ggml_tensor * top_k, + int64_t n_kv_raw) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->src[5] == NULL); + GGML_ASSERT(top_k->type == GGML_TYPE_I32); + GGML_ASSERT(top_k->ne[1] == a->src[0]->ne[1]); + GGML_ASSERT(n_kv_raw >= 0 && n_kv_raw <= a->src[1]->ne[1]); + + a->src[5] = top_k; + ggml_set_op_params_i32(a, 4, (int32_t) n_kv_raw); +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back( diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0ef..04bc11f7edbd 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2546,8 +2546,10 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * kq_mask, ggml_tensor * sinks, ggml_tensor * v_mla, - float kq_scale, - int il) const { + float kq_scale, + int il, + ggml_tensor * top_k, + int64_t n_kv_raw) const { const bool v_trans = v->nb[1] > v->nb[2]; // split the batch into streams if needed @@ -2583,6 +2585,9 @@ ggml_tensor * llm_graph_context::build_attn_mha( res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); ggml_flash_attn_ext_add_sinks(cur, sinks); + if (top_k) { + ggml_flash_attn_ext_add_top_k(cur, top_k, n_kv_raw); + } ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); if (v_mla) { diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb53..bbfcdd9ea27d 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1172,7 +1172,9 @@ struct llm_graph_context { ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] float kq_scale, - int il) const; + int il, + ggml_tensor * top_k = nullptr, + int64_t n_kv_raw = 0) const; llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index fc816e2aeb43..08415fe97a12 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -754,7 +754,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention( ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0); cb(kq_mask, "csa_lid_kq_mask", il); - ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il); + ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, kq_scale, il, top_k, raw_k->ne[2]); if (k_rot) { out = llama_mul_mat_hadamard(ctx0, out, k_rot); } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 437073855e4e..2cdcb1d5ec55 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10153,6 +10153,9 @@ static std::vector> make_test_cases_eval() { } } } + test_cases.emplace_back(new test_lightning_indexer(128, 64, 257, 1, 1, 1, GGML_TYPE_F16)); + test_cases.emplace_back(new test_lightning_indexer(128, 64, 257, 17, 1, 1, GGML_TYPE_F16)); + test_cases.emplace_back(new test_lightning_indexer(128, 64, 512, 512, 1, 1, GGML_TYPE_F16)); for (int kv : { 1, 7, 8, 63, 64, 65 }) { for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0}) { @@ -10190,6 +10193,11 @@ static std::vector> make_test_cases_perf() { GGML_TYPE_F32, {n_kv, 512, 64, 1}, false, {2, 1, 0, 3})); } + for (ggml_type type_a : { GGML_TYPE_IQ2_XS, GGML_TYPE_IQ3_XXS }) { + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 256, 6, false, 2048, 512, 4096)); + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 256, 6, false, 4096, 512, 2048)); + } + // Conv2d: K=CRS=NPQ=4096 matmul performance uint32_t iwh_idx = 0; uint32_t kwh_idx = 1; @@ -10558,7 +10566,7 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 64, 1, 1, false, true)); // KDA PP-64 // lightning_indexer - for (int kv : { 256, 4096, 65536 }) { + for (int kv : { 256, 512, 4096, 65536 }) { for (int bs : { 1, 512, 2048 }) { for (int nh : { 32, 64 }) { for (int ns : { 1, 4 }) { From aa04dcc08e1cda3f8a397fd84f671d5712959672 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sat, 1 Aug 2026 13:30:31 +0000 Subject: [PATCH 076/109] vulkan: harden the sparse-FA shader and document the top-k API - flash_attn_top_k.comp: remove the dead bounds check that would skip barriers for part of the workgroup if it ever fired (barrier divergence is UB; the dispatch gate sizes the grid exactly), pin per-lane sizing to LANES=64 instead of the SUBGROUP_SIZE spec constant (array bounds fold from the spec default at compile time), name the f16 mask threshold constant - ggml.h: document ggml_flash_attn_ext_add_top_k semantics (index base, dense prefix, backends-may-ignore contract) - tests: cover the scalar indexer variant (batch 4 and boundary 15), previously only the cm and decode-cm variants had eval parity cases Co-Authored-By: Claude Fable 5 --- ggml/include/ggml.h | 5 +++ .../vulkan-shaders/flash_attn_top_k.comp | 36 ++++++++++++------- tests/test-backend-ops.cpp | 4 +++ 3 files changed, 32 insertions(+), 13 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d81b8e53050e..9071e3c7056a 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2450,6 +2450,11 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * sinks); + // sparse attention hint: attend only to the first n_kv_raw keys (dense prefix) plus the + // keys selected by top_k. top_k is I32 [n_top_k, n_tokens, 1, n_streams]; each index i + // selects absolute key n_kv_raw + i. Negative or out-of-range indices are ignored. + // Backends may ignore the hint: the kq_mask must still encode the same selection, so a + // dense fallback computes the identical result. GGML_API void ggml_flash_attn_ext_add_top_k( struct ggml_tensor * a, struct ggml_tensor * top_k, diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp index 7f82266781f1..b8b49c677bd5 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp @@ -39,9 +39,18 @@ layout(push_constant) uniform Parameters { uint has_sinks; } p; +// Shape constants pinned by the dispatch gate in ggml_vk_flash_attn_top_k: DeepSeek V4 +// CSA attention only (hd 512, 64 heads MQA over an f16 K==V latent). The gate also pins +// the pipeline to subgroup size 64, so per-lane sizing uses LANES rather than the +// SUBGROUP_SIZE specialization constant (a spec-constant array bound would be folded at +// compile time from the default and silently break under a different runtime subgroup). const uint HEAD_SIZE = 512; const uint HEADS_PER_GROUP = 8; const uint KEYS_PER_BLOCK = 16; +const uint LANES = 64; +// f16 -inf (or the lowest f16 normal some mask writers use in its place) marks a +// masked-out key. +const float MASK_NEG_INF = -65500.0; shared float16_t key_sh[KEYS_PER_BLOCK * HEAD_SIZE]; shared uint key_idx[KEYS_PER_BLOCK]; @@ -53,12 +62,13 @@ void main() { const uint token = gl_WorkGroupID.x; const uint stream = gl_WorkGroupID.z; - if (token >= p.n_batch || head >= p.n_head) { - return; - } + // No bounds check: the grid is sized exactly (x = n_batch, y * HEADS_PER_GROUP covers + // n_head == 64). An early return here would skip the barriers below for part of the + // workgroup, which is undefined behavior — do not reintroduce one without restructuring + // the barrier flow. - float accum[HEAD_SIZE / SUBGROUP_SIZE]; - [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + float accum[HEAD_SIZE / LANES]; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { accum[i] = 0.0; } @@ -101,13 +111,13 @@ void main() { } float partial = 0.0; - [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { - const uint dim = lane + i * SUBGROUP_SIZE; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + const uint dim = lane + i * LANES; partial += data_q[q_base + dim] * float(key_sh[col * HEAD_SIZE + dim]); } const float mask = float(data_m[mask_base + key]); const float score = subgroupAdd(partial) * p.scale + mask; - if (mask < -65500.0) { + if (mask < MASK_NEG_INF) { continue; } @@ -117,8 +127,8 @@ void main() { row_sum = row_sum * old_scale + value_scale; row_max = new_max; - [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { - const uint dim = lane + i * SUBGROUP_SIZE; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + const uint dim = lane + i * LANES; accum[i] = accum[i] * old_scale + value_scale * float(key_sh[col * HEAD_SIZE + dim]); } } @@ -131,14 +141,14 @@ void main() { const float old_scale = row_sum == 0.0 ? 0.0 : exp(row_max - new_max); const float sink_scale = exp(sink - new_max); row_sum = row_sum * old_scale + sink_scale; - [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { accum[i] *= old_scale; } } const uint dst_base = stream * p.nb3 + token * p.nb2 + head * p.nb1; const float inv_sum = row_sum == 0.0 ? 0.0 : 1.0 / row_sum; - [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { - data_dst[dst_base + lane + i * SUBGROUP_SIZE] = accum[i] * inv_sum; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + data_dst[dst_base + lane + i * LANES] = accum[i] * inv_sum; } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 2cdcb1d5ec55..e18981e4cdad 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10153,7 +10153,11 @@ static std::vector> make_test_cases_eval() { } } } + // batch 1 = Vulkan decode-cm variant, 4/15 = scalar subgroup variant (below the cm + // threshold of 16, 15 is the boundary), 17/512 = cm prefill variant test_cases.emplace_back(new test_lightning_indexer(128, 64, 257, 1, 1, 1, GGML_TYPE_F16)); + test_cases.emplace_back(new test_lightning_indexer(128, 64, 257, 4, 1, 1, GGML_TYPE_F16)); + test_cases.emplace_back(new test_lightning_indexer(128, 64, 257, 15, 1, 1, GGML_TYPE_F16)); test_cases.emplace_back(new test_lightning_indexer(128, 64, 257, 17, 1, 1, GGML_TYPE_F16)); test_cases.emplace_back(new test_lightning_indexer(128, 64, 512, 512, 1, 1, GGML_TYPE_F16)); From 5046bf67df3c6b8c578bc0f74f195b034bbe2f9b Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sat, 1 Aug 2026 14:55:05 +0000 Subject: [PATCH 077/109] tests: sparse top-k FA parity + perf coverage (V4 CSA shape) Adds test_flash_attn_ext_top_k: builds the DeepSeek V4 CSA attention shape (hd 512, 64-head MQA, V as a view of K) with a consistent per-token top-k/mask pair, one deliberately invalid index, and cases on both sides of the Vulkan engagement gates. nb >= 64 cases are the first numerical parity coverage the sparse prefill shader has had; nb < 64 and sub-3x-kv cases pin the dense-fallback contract. Perf cases sweep kv 8k/32k/64k at nb 1/8/64/512 with a fixed active set. Measured on gfx1151: the sparse shader is flat vs kv at prefill (~2.2 TFLOPS active-only) while nb < 64 falls back to dense and scales with kv (1326 us at 64k, nb=1) - the gap a sparse decode path needs to close. Co-Authored-By: Claude Fable 5 --- tests/test-backend-ops.cpp | 134 +++++++++++++++++++++++++++++++++++++ 1 file changed, 134 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e18981e4cdad..a8d57ae43261 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7198,6 +7198,117 @@ struct test_flash_attn_ext : public test_case { } }; +// GGML_OP_FLASH_ATTN_EXT with a top-k sparse selection hint (DeepSeek V4 CSA shape). +// The kq_mask encodes the same selection as the top_k indices, so a backend that ignores +// the hint (CPU) computes the identical result densely — this is exactly the contract +// that keeps sparse and dense paths interchangeable, and what this test verifies. +struct test_flash_attn_ext_top_k : public test_case { + const int64_t kv; // total KV size (compressed region + dense prefix) + const int64_t nb; // batch size (query tokens) + const int64_t n_kv_raw; // dense prefix always attended + const int64_t n_top_k; // selected keys per query token + const bool sinks; + + static constexpr int64_t hs = 512; // V4 CSA head size, K == V latent + static constexpr int64_t nh = 64; // V4 CSA query heads (MQA) + + std::string vars() override { + return VARS_TO_STR5(kv, nb, n_kv_raw, n_top_k, sinks); + } + + double max_nmse_err() override { + return 5e-4; + } + + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + // only the active keys contribute compute on a sparse backend; count those so + // perf mode reports the useful-work rate + return 2 * nh * nb * (hs + hs) * (n_kv_raw + n_top_k); + } + + test_flash_attn_ext_top_k(int64_t kv = 768, int64_t nb = 8, int64_t n_kv_raw = 64, int64_t n_top_k = 128, bool sinks = false) + : kv(kv), nb(nb), n_kv_raw(n_kv_raw), n_top_k(n_top_k), sinks(sinks) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hs, nb, nh, 1); + ggml_set_name(q, "q"); + + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, hs, kv, 1, 1); + ggml_set_name(k, "k"); + + // V4 CSA attends over the K latent itself: V is the same cache tensor + ggml_tensor * v = ggml_view_4d(ctx, k, hs, kv, 1, 1, k->nb[1], k->nb[2], k->nb[3], 0); + ggml_set_name(v, "v"); + + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, 1); + ggml_set_name(m, "m"); + + ggml_tensor * t = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, nb, 1, 1); + ggml_set_name(t, "top_k"); + + ggml_tensor * s = nullptr; + if (sinks) { + s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nh); + ggml_set_name(s, "s"); + } + + ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hs), 0.0f, 0.0f); + ggml_flash_attn_ext_add_sinks(out, s); + ggml_flash_attn_ext_add_top_k(out, t, n_kv_raw); + ggml_flash_attn_ext_set_prec (out, GGML_PREC_F32); + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + const int64_t range = kv - n_kv_raw; // size of the selectable compressed region + + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "top_k") == 0 || strcmp(t->name, "m") == 0) { + continue; // filled together below + } + if (strcmp(t->name, "s") == 0) { + init_tensor_uniform(t, -10.0f, 10.0f); + } else { + init_tensor_uniform(t); + } + } + + // build a consistent (top_k, mask) pair: a deterministic per-token selection, + // strided so adjacent tokens select overlapping-but-different keys, with one + // deliberately invalid index (-1) whose mask slot stays -inf + std::vector top(n_top_k * nb); + std::vector mask(kv * nb); + const ggml_fp16_t minus_inf = ggml_fp32_to_fp16(-INFINITY); + const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f); + + for (int64_t b = 0; b < nb; ++b) { + for (int64_t i = 0; i < kv; ++i) { + mask[b * kv + i] = i < n_kv_raw ? zero : minus_inf; + } + for (int64_t j = 0; j < n_top_k; ++j) { + int32_t idx = (int32_t) ((j * range) / n_top_k + b) % (int32_t) range; + if (j == n_top_k - 1 && b == 0) { + idx = -1; // exercise the ignore-invalid-index path + } else { + mask[b * kv + n_kv_raw + idx] = zero; + } + top[b * n_top_k + j] = idx; + } + } + + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "top_k") == 0) { + ggml_backend_tensor_set(t, top.data(), 0, top.size() * sizeof(int32_t)); + } else if (strcmp(t->name, "m") == 0) { + ggml_backend_tensor_set(t, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); + } + } + } +}; + // GGML_OP_CROSS_ENTROPY_LOSS struct test_cross_entropy_loss : public test_case { const ggml_type type; @@ -10167,6 +10278,18 @@ static std::vector> make_test_cases_eval() { } } + // sparse top-k FA: (kv, nb, n_kv_raw, n_top_k, sinks). The Vulkan sparse path engages + // when kv >= 3*(n_kv_raw + n_top_k) AND nb >= 64 (prefill-only); the nb < 64 cases + // and the kv=512 case verify dense-fallback parity with the hint attached, the + // nb=64/128 cases exercise the sparse shader itself. + test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 1, 256, 512, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 8, 64, 128, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 17, 64, 128, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k( 512, 4, 64, 128, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, true)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 128, 256, 512, false)); + return test_cases; } #ifdef _MSC_VER @@ -10582,6 +10705,17 @@ static std::vector> make_test_cases_perf() { } } + // sparse top-k FA at V4 decode/prefill shapes — the A/B instrument for the + // gather-to-compact work (n_active = n_kv_raw + n_top_k stays fixed as kv grows). + // nb 1/8 currently takes the DENSE path (the sparse shader gates on nb >= 64): + // those rows measure the decode cost gather-to-compact must beat. nb 64/512 + // measures the existing sparse prefill shader. + for (int kv : { 8192, 32768, 65536 }) { + for (int nb : { 1, 8, 64, 512 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, nb, 1024, 512, false)); + } + } + return test_cases; } From 8a15795eb4306da43affadcbab98a9e38fc82fbe Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sat, 1 Aug 2026 15:08:58 +0000 Subject: [PATCH 078/109] vulkan: gather-to-compact sparse decode FA for DeepSeek V4 top-k selection The sparse prefill shader gates on q->ne[1] >= 64, so single-token decode attends densely over the whole compressed KV and its cost grows with context. This adds a gather pass (flash_attn_gather.comp): copy the active rows (dense prefix + top-k selection; MQA, so all 64 query heads share one set) plus their mask values into a compact contiguous scratch in prealloc_y, then run the ordinary dense FA over the compacted K/V/mask. V is the K latent, so one gather serves both. Invalid indices and padding get zeroed K and -inf mask. The FA function itself only has its inputs swapped: KV, mask geometry, strides and the K/V/mask bindings are overridden up front and every downstream decision (pipeline choice, split-k, workgroup sizing, use_mask_opt) sizes itself to the compact KV unchanged. Engages for the V4 CSA decode shape when kv >= 2x the padded active set; GGML_VK_FA_TOPK_GATHER=0 disables. Measured (gfx1151, test-backend-ops perf, active set 1536): kv=8192 nb=1: 255.6 us -> 55.8 us (4.6x) kv=32768 nb=1: 986.4 us -> 56.1 us (17.6x) kv=65536 nb=1: 1333.8 us -> 58.7 us (22.7x) Time is flat vs context. nb>1 still falls back to dense pending a union gather. FLASH_ATTN_EXT eval suite green incl. the top-k parity cases (the kv=4096 nb=1 case exercises this path end-to-end vs CPU). Co-Authored-By: Claude Fable 5 Fold-in note for the toolbox branch: use_dequant_kv is additionally gated on !fa_compact.active (the compact scratch is already contiguous f16, and the two scratch layers must not stack), and the compact stride overrides chain through the toolbox's nb*_eff values so the contiguize/dequant path keeps its strides when the gather is inactive. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 143 +++++++++++++++++- .../vulkan-shaders/flash_attn_gather.comp | 70 +++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + 3 files changed, 207 insertions(+), 7 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8d921f51c09a..9eb7f819b910 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1090,6 +1090,7 @@ struct vk_device_struct { vk_pipeline pipeline_lightning_indexer_cm_f16; vk_pipeline pipeline_lightning_indexer_decode_cm_f16; vk_pipeline pipeline_flash_attn_top_k_f16; + vk_pipeline pipeline_flash_attn_gather_f16; vk_pipeline pipeline_ssm_scan_f32_d128; vk_pipeline pipeline_ssm_scan_f32_d256; vk_pipeline pipeline_ssm_conv_f32; @@ -1921,6 +1922,12 @@ struct vk_op_lightning_indexer_cm_push_constants { }; static_assert(sizeof(vk_op_lightning_indexer_cm_push_constants) <= 128); +struct vk_op_flash_attn_gather_push_constants { + uint32_t n_kv, n_kv_raw, n_top_k, kv_c; + uint32_t nbk1, nbk3, nbt3, nbm3, nem3; +}; +static_assert(sizeof(vk_op_flash_attn_gather_push_constants) <= 128); + struct vk_op_flash_attn_top_k_push_constants { uint32_t n_batch, n_kv, n_kv_raw, n_top_k, n_head; uint32_t nbq1, nbq2, nbq3; @@ -6045,6 +6052,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { "flash_attn_top_k_f16", flash_attn_top_k_f16_len, flash_attn_top_k_f16_data, "main", 6, sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true, device->subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_f16, + "flash_attn_gather_f16", flash_attn_gather_f16_len, flash_attn_gather_f16_data, "main", 5, + sizeof(vk_op_flash_attn_gather_push_constants), {1, 1, 1}, {}, 1, true, true, + device->subgroup_size); } if (device->subgroup_arithmetic && device->subgroup_require_full_support) { @@ -11098,6 +11109,95 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & return true; } +struct vk_fa_compact_state { + bool active = false; + uint32_t kv_c = 0; + vk_subbuffer kc_buf, mc_buf; +}; + +// V4 sparse decode (gather-to-compact): the sparse prefill shader above gates on +// q->ne[1] >= 64, so single-token decode otherwise attends densely over the whole +// compressed KV, at a cost that grows with context. Instead, gather the active rows +// (dense prefix + top-k selection; MQA, so all query heads share one set) into a +// compact contiguous scratch in prealloc_y, and let the ordinary dense FA below run +// over the compacted K/V/mask. Correct by the same contract as the sparse shader: +// the source mask carries the selection, and the gathered mask preserves it. +static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_context & subctx, + const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, + const ggml_tensor * mask, ggml_tensor * dst, vk_fa_compact_state & st) { + const ggml_tensor * top_k = dst->src[5]; + static const char * gather_env = getenv("GGML_VK_FA_TOPK_GATHER"); + if ((gather_env && gather_env[0] == '0') || + !top_k || !ctx->device->pipeline_flash_attn_gather_f16 || + q->ne[1] != 1 || // single-token decode only; batched queries need a union gather + q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || + !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || + q->ne[0] != 512 || k->ne[0] != 512 || v->ne[0] != 512 || q->ne[2] != 64 || + k->ne[2] != 1 || v->ne[2] != 1 || + q->ne[1] != top_k->ne[1] || q->ne[3] != top_k->ne[3] || + k->ne[1] != v->ne[1] || k->buffer != v->buffer || k->data != v->data || + !ggml_is_contiguous(mask) || !ggml_is_contiguous(top_k)) { + return false; + } + + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || logit_softcap != 0.0f) { + return false; + } + + const int32_t n_kv_raw = ggml_get_op_params_i32(dst, 4); + if (n_kv_raw < 0 || n_kv_raw > k->ne[1] || top_k->ne[0] > k->ne[1] - n_kv_raw) { + return false; + } + + const uint32_t kv_c = GGML_PAD((uint32_t)(n_kv_raw + top_k->ne[0]), 256u); + // the gather writes then re-reads ~the active bytes; dense reads the source KV once, + // so compaction only pays when the source is comfortably larger than the active set + if ((uint64_t) k->ne[1] < 2ull * kv_c) { + return false; + } + + const uint32_t ns = (uint32_t) q->ne[3]; + const size_t kc_sz = (size_t) ns * kv_c * 512 * sizeof(ggml_fp16_t); + const size_t mc_sz = (size_t) ns * kv_c * sizeof(ggml_fp16_t); + + if (ctx->prealloc_size_y < kc_sz + mc_sz) { + ctx->prealloc_size_y = kc_sz + mc_sz; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_y_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + + vk_pipeline pipeline = ctx->device->pipeline_flash_attn_gather_f16; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_flash_attn_gather_push_constants pc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, (uint32_t) top_k->ne[0], kv_c, + (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (top_k->nb[3] / sizeof(int32_t)), + (uint32_t) (mask->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) mask->ne[3], + }; + + st.kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); + st.mc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, kc_sz); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, top_k), + ggml_vk_tensor_subbuffer(ctx, mask), st.kc_buf, st.mc_buf }, + pc, { kv_c, 1, ns }); + ggml_vk_sync_buffers(ctx, subctx); + ctx->prealloc_y_need_sync = true; + + st.active = true; + st.kv_c = kv_c; + return true; +} + static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { VK_LOG_DEBUG("ggml_vk_flash_attn((" << q << ", name=" << q->name << ", type=" << q->type << ", ne0=" << q->ne[0] << ", ne1=" << q->ne[1] << ", ne2=" << q->ne[2] << ", ne3=" << q->ne[3] << ", nb0=" << q->nb[0] << ", nb1=" << q->nb[1] << ", nb2=" << q->nb[2] << ", nb3=" << q->nb[3]; std::cerr << "), (" << k << ", name=" << k->name << ", type=" << k->type << ", ne0=" << k->ne[0] << ", ne1=" << k->ne[1] << ", ne2=" << k->ne[2] << ", ne3=" << k->ne[3] << ", nb0=" << k->nb[0] << ", nb1=" << k->nb[1] << ", nb2=" << k->nb[2] << ", nb3=" << k->nb[3]; @@ -11117,15 +11217,15 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - const uint32_t nem0 = mask ? mask->ne[0] : 0; - const uint32_t nem1 = mask ? mask->ne[1] : 0; - const uint32_t nem2 = mask ? mask->ne[2] : 0; - const uint32_t nem3 = mask ? mask->ne[3] : 0; + uint32_t nem0 = mask ? mask->ne[0] : 0; + uint32_t nem1 = mask ? mask->ne[1] : 0; + uint32_t nem2 = mask ? mask->ne[2] : 0; + uint32_t nem3 = mask ? mask->ne[3] : 0; const uint32_t HSK = nek0; const uint32_t HSV = nev0; uint32_t N = neq1; - const uint32_t KV = nek1; + uint32_t KV = nek1; GGML_ASSERT(ne0 == HSV); GGML_ASSERT(ne2 == N); @@ -11152,6 +11252,17 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx if (ggml_vk_flash_attn_top_k(ctx, subctx, q, k, v, mask, sinks, dst)) { return; } + // V4 sparse decode: gather the active set into a compact scratch and run the dense + // FA below on it. Overrides KV, the mask geometry, and (further down) the K/V/mask + // bindings and strides; every other decision then sizes itself to the compact KV. + vk_fa_compact_state fa_compact; + if (ggml_vk_flash_attn_gather_compact(ctx, subctx, q, k, v, mask, dst, fa_compact)) { + KV = fa_compact.kv_c; + nem0 = fa_compact.kv_c; + nem1 = N; + nem2 = 1; + nem3 = (uint32_t) q->ne[3]; + } uint32_t gqa_ratio = 1; uint32_t qk_ratio = neq2 / nek2; uint32_t workgroups_x = (uint32_t)neq1; @@ -11198,6 +11309,9 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const bool kv_needs_dequant = !ggml_vk_fa_kv_native(k->type, ctx->device->coopmat2) || !ggml_vk_fa_kv_native(v->type, ctx->device->coopmat2); const bool use_dequant_kv = !fa_dequant_off && + // the gather-to-compact scratch is already contiguous f16; the + // dequant/contiguize pass must not run on top of it + !fa_compact.active && ((k_quant && v_quant) || kv_needs_dequant || (fa_kv_contig && kv_f16_strided)) && neq1 >= 64 && is_dense_kv_cache(k) && is_dense_kv_cache(v) && kv_f16_sz <= ctx->device->properties.limits.maxStorageBufferRange && @@ -11236,6 +11350,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx 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)); uint32_t v_stride = (uint32_t)(nbv1 / ggml_type_size(v->type)); + if (fa_compact.active) { + k_stride = 512; + v_stride = 512; + } // For F32, the shader treats it as a block of size 4 (for vec4 loads) if (k->type == GGML_TYPE_F32) { @@ -11383,6 +11501,11 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx vk_subbuffer v_buf = ggml_vk_tensor_subbuffer(ctx, v); vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); vk_subbuffer mask_buf = mask ? ggml_vk_tensor_subbuffer(ctx, mask) : q_buf; + if (fa_compact.active) { + k_buf = fa_compact.kc_buf; + v_buf = fa_compact.kc_buf; // V is the K latent; one gather serves both + mask_buf = fa_compact.mc_buf; + } vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; vk_subbuffer mask_opt_buf = use_mask_opt ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf; @@ -11441,6 +11564,12 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx ggml_vk_sync_buffers(ctx, subctx); } + // compact scratch layout: [512, kv_c, 1, ns] f16, tightly packed + const uint32_t eff_nbk2 = fa_compact.active ? fa_compact.kv_c * 512 * (uint32_t)sizeof(ggml_fp16_t) : nbk2_eff; + const uint32_t eff_nbk3 = fa_compact.active ? fa_compact.kv_c * 512 * (uint32_t)sizeof(ggml_fp16_t) : nbk3_eff; + const uint32_t eff_nbv2 = fa_compact.active ? eff_nbk2 : nbv2_eff; + const uint32_t eff_nbv3 = fa_compact.active ? eff_nbk3 : nbv3_eff; + const vk_flash_attn_push_constants pc = { N, KV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, (uint32_t)neq2, (uint32_t)neq3, @@ -11448,8 +11577,8 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx (uint32_t)nev2, (uint32_t)nev3, nem1, nem2, nem3, q_stride, (uint32_t)nbq2, (uint32_t)nbq3, - k_stride, nbk2_eff, nbk3_eff, - v_stride, nbv2_eff, nbv3_eff, + k_stride, eff_nbk2, eff_nbk3, + v_stride, eff_nbv2, eff_nbv3, scale, max_bias, logit_softcap, mask_n_head_log2, m0, m1, gqa_ratio, split_kv, split_k }; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp new file mode 100644 index 000000000000..4e3dc1c42623 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp @@ -0,0 +1,70 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +// Gathers the active KV rows of a top-k sparse attention (DeepSeek V4 CSA decode) into a +// compact contiguous scratch: rows [0, n_kv_raw) of the source (the dense prefix), then the +// n_top_k selected rows, then zero padding up to kv_c. The gathered mask row keeps the +// per-key mask values so causality/validity survive compaction; invalid top-k indices and +// padding get -inf mask and zeroed K (softmax-neutral either way, zeroed so no NaN*0). +// One workgroup per compact row; V is the K latent (V==K), so a single gather serves both. + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 1) readonly buffer TopBuf { int data_top[]; }; +layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; +layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; // total source KV rows + uint n_kv_raw; // dense prefix length + uint n_top_k; // selected rows for the (single) query token + uint kv_c; // padded compact row count == dispatch row range + uint nbk1; // K source row stride, elements + uint nbk3; // K source stream stride, elements + uint nbt3; // top_k stream stride, elements + uint nbm3; // mask source stream stride, elements + uint nem3; // mask ne[3], for stream broadcast +} p; + +const uint HEAD_SIZE = 512; +const uint LANES = 64; + +void main() { + const uint row = gl_WorkGroupID.x; + const uint stream = gl_WorkGroupID.z; + const uint tid = gl_LocalInvocationIndex; + + // map compact row -> source row; p.n_kv is the invalid sentinel + uint src = p.n_kv; + if (row < p.n_kv_raw) { + src = row; + } else if (row < p.n_kv_raw + p.n_top_k) { + const int idx = data_top[stream * p.nbt3 + (row - p.n_kv_raw)]; + if (idx >= 0 && uint(idx) < p.n_kv - p.n_kv_raw) { + src = p.n_kv_raw + uint(idx); + } + } + + const uint dst_base = (stream * p.kv_c + row) * HEAD_SIZE; + if (src < p.n_kv) { + const uint src_base = stream * p.nbk3 + src * p.nbk1; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + data_kc[dst_base + tid + i * LANES] = data_k[src_base + tid + i * LANES]; + } + if (tid == 0) { + data_mc[stream * p.kv_c + row] = data_m[(stream % p.nem3) * p.nbm3 + src]; + } + } else { + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + data_kc[dst_base + tid + i * LANES] = float16_t(0.0); + } + if (tid == 0) { + data_mc[stream * p.kv_c + row] = float16_t(uintBitsToFloat(0xff800000)); + } + } +} 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 d3acc412e8cd..4a75993680fc 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -807,6 +807,7 @@ void process_shaders() { string_to_spv("lightning_indexer_decode_cm_f16", "lightning_indexer_decode_cm.comp", {}); #endif string_to_spv("flash_attn_top_k_f16", "flash_attn_top_k.comp", {}); + string_to_spv("flash_attn_gather_f16", "flash_attn_gather.comp", {}); string_to_spv("mul_mat_vec_p021_f16_f32_subgroup_add", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}); string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}); From a27636c9256d62a66a4ea325aadbfdfa4ee449d7 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 2 Aug 2026 00:04:55 +0000 Subject: [PATCH 079/109] vulkan: fused DeepSeek V4 hyper-connection ops (HC pre / comb / post) Ports ggml-cuda/dsv4-hc.cu to Vulkan: hc_pre (mix the HC input streams down to one embedding), hc_comb (per-token 4x4 stream-mixing matrix - per-source softmax then eps-stabilized alternating column/row sinkhorn normalization, whole matrix in registers, one thread per token), and hc_post (redistribute the layer output back into the streams with the mixed residual). Same launch geometry as the CUDA kernels; plain f32 compute, no subgroup or coopmat requirements, so the pipelines are created unconditionally. The value at decode is dispatch-count collapse: the unfused fallback runs the decomposed graph (measured on gfx1151 config-a partial offload: SUM_ROWS alone 79 dispatches x 39.7us = 3.1ms per graph, plus DIV/MUL/ ADD shares at 4x4 shapes) where the fused form is 3 dispatches per layer. resolve_fused_ops now keeps all three fusions enabled on Vulkan instead of printing 'not supported, set to disabled'. Parity: upstream test-backend-ops DSV4_HC_PRE/COMB/POST cases green vs CPU on first build. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 153 ++++++++++++++++++ .../vulkan-shaders/dsv4_hc_comb.comp | 101 ++++++++++++ .../vulkan-shaders/dsv4_hc_post.comp | 44 +++++ .../vulkan-shaders/dsv4_hc_pre.comp | 37 +++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 3 + 5 files changed, 338 insertions(+) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 9eb7f819b910..a712d2bdd165 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1091,6 +1091,9 @@ struct vk_device_struct { vk_pipeline pipeline_lightning_indexer_decode_cm_f16; vk_pipeline pipeline_flash_attn_top_k_f16; vk_pipeline pipeline_flash_attn_gather_f16; + vk_pipeline pipeline_dsv4_hc_pre_f32; + vk_pipeline pipeline_dsv4_hc_comb_f32; + vk_pipeline pipeline_dsv4_hc_post_f32; vk_pipeline pipeline_ssm_scan_f32_d128; vk_pipeline pipeline_ssm_scan_f32_d256; vk_pipeline pipeline_ssm_conv_f32; @@ -1922,6 +1925,35 @@ struct vk_op_lightning_indexer_cm_push_constants { }; static_assert(sizeof(vk_op_lightning_indexer_cm_push_constants) <= 128); +struct vk_op_dsv4_hc_pre_push_constants { + uint32_t n_embd, hc, nr; + uint32_t sx0, sx1, sx2; + uint32_t sw0, sw1; + uint32_t sd0, sd1; +}; +static_assert(sizeof(vk_op_dsv4_hc_pre_push_constants) <= 128); + +struct vk_op_dsv4_hc_comb_push_constants { + uint32_t n_tokens; + uint32_t sm0, sm1; + uint32_t ss0; + uint32_t sb0; + uint32_t sd0, sd1, sd2; + float eps; + int32_t n_iter; +}; +static_assert(sizeof(vk_op_dsv4_hc_comb_push_constants) <= 128); + +struct vk_op_dsv4_hc_post_push_constants { + uint32_t n_embd, hc, nr; + uint32_t sx0, sx1; + uint32_t sr0, sr1, sr2; + uint32_t sp0, sp1; + uint32_t sc0, sc1, sc2; + uint32_t sd0, sd1, sd2; +}; +static_assert(sizeof(vk_op_dsv4_hc_post_push_constants) <= 128); + struct vk_op_flash_attn_gather_push_constants { uint32_t n_kv, n_kv_raw, n_top_k, kv_c; uint32_t nbk1, nbk3, nbt3, nbm3, nem3; @@ -6058,6 +6090,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { device->subgroup_size); } + // DSv4 fused hyper-connection ops: plain f32 compute, no subgroup/coopmat requirements + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_f32, + "dsv4_hc_pre_f32", dsv4_hc_pre_f32_len, dsv4_hc_pre_f32_data, "main", 3, + sizeof(vk_op_dsv4_hc_pre_push_constants), {256, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_comb_f32, + "dsv4_hc_comb_f32", dsv4_hc_comb_f32_len, dsv4_hc_comb_f32_data, "main", 4, + sizeof(vk_op_dsv4_hc_comb_push_constants), {256, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32, + "dsv4_hc_post_f32", dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, + sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, {}, 1); + if (device->subgroup_arithmetic && device->subgroup_require_full_support) { ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d128, "ssm_scan_128_f32", ssm_scan_subgroup_f32_len, ssm_scan_subgroup_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {128, device->subgroup_size}, 1, true, true); ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_256_f32", ssm_scan_subgroup_f32_len, ssm_scan_subgroup_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, device->subgroup_size}, 1, true, true); @@ -13426,6 +13469,91 @@ static void ggml_vk_lightning_indexer_cm(ggml_backend_vk_context * ctx, vk_conte pc, {(uint32_t) k->ne[2], (uint32_t) q->ne[2], (uint32_t) q->ne[3]}); } +// DSv4 fused hyper-connection ops — ports of ggml-cuda/dsv4-hc.cu. Strides are passed in +// f32 elements; grids mirror the CUDA launch geometry (flat 1D, 256 threads per workgroup). + +static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * w = dst->src[1]; + + const uint32_t n_embd = (uint32_t) x->ne[0]; + const uint32_t hc = (uint32_t) x->ne[1]; + const uint32_t n_tokens = (uint32_t) x->ne[2]; + const uint32_t nr = n_embd * n_tokens; + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_pre_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_dsv4_hc_pre_push_constants pc = { + n_embd, hc, nr, + (uint32_t)(x->nb[0] / sizeof(float)), (uint32_t)(x->nb[1] / sizeof(float)), (uint32_t)(x->nb[2] / sizeof(float)), + (uint32_t)(w->nb[0] / sizeof(float)), (uint32_t)(w->nb[1] / sizeof(float)), + (uint32_t)(dst->nb[0] / sizeof(float)), (uint32_t)(dst->nb[1] / sizeof(float)), + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, x), ggml_vk_tensor_subbuffer(ctx, w), + ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {nr, 1, 1}); +} + +static void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * mixes = dst->src[0]; + const ggml_tensor * scale = dst->src[1]; + const ggml_tensor * base = dst->src[2]; + + const uint32_t n_tokens = (uint32_t) mixes->ne[1]; + const float eps = ggml_get_op_params_f32(dst, 0); + const int32_t n_iter = ggml_get_op_params_i32(dst, 1); + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_comb_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_dsv4_hc_comb_push_constants pc = { + n_tokens, + (uint32_t)(mixes->nb[0] / sizeof(float)), (uint32_t)(mixes->nb[1] / sizeof(float)), + (uint32_t)(scale->nb[0] / sizeof(float)), + (uint32_t)(base->nb[0] / sizeof(float)), + (uint32_t)(dst->nb[0] / sizeof(float)), (uint32_t)(dst->nb[1] / sizeof(float)), (uint32_t)(dst->nb[2] / sizeof(float)), + eps, n_iter, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, mixes), ggml_vk_tensor_subbuffer(ctx, scale), + ggml_vk_tensor_subbuffer(ctx, base), ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {n_tokens, 1, 1}); +} + +static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * residual = dst->src[1]; + const ggml_tensor * post = dst->src[2]; + const ggml_tensor * comb = dst->src[3]; + + const uint32_t n_embd = (uint32_t) x->ne[0]; + const uint32_t n_tokens = (uint32_t) x->ne[1]; + const uint32_t hc = (uint32_t) residual->ne[1]; + const uint32_t nr = n_embd * hc * n_tokens; + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_post_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_dsv4_hc_post_push_constants pc = { + n_embd, hc, nr, + (uint32_t)(x->nb[0] / sizeof(float)), (uint32_t)(x->nb[1] / sizeof(float)), + (uint32_t)(residual->nb[0] / sizeof(float)), (uint32_t)(residual->nb[1] / sizeof(float)), (uint32_t)(residual->nb[2] / sizeof(float)), + (uint32_t)(post->nb[0] / sizeof(float)), (uint32_t)(post->nb[1] / sizeof(float)), + (uint32_t)(comb->nb[0] / sizeof(float)), (uint32_t)(comb->nb[1] / sizeof(float)), (uint32_t)(comb->nb[2] / sizeof(float)), + (uint32_t)(dst->nb[0] / sizeof(float)), (uint32_t)(dst->nb[1] / sizeof(float)), (uint32_t)(dst->nb[2] / sizeof(float)), + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, x), ggml_vk_tensor_subbuffer(ctx, residual), + ggml_vk_tensor_subbuffer(ctx, post), ggml_vk_tensor_subbuffer(ctx, comb), + ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {nr, 1, 1}); +} + static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -16505,6 +16633,21 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; + case GGML_OP_DSV4_HC_PRE: + ggml_vk_dsv4_hc_pre(ctx, compute_ctx, node); + + break; + + case GGML_OP_DSV4_HC_COMB: + ggml_vk_dsv4_hc_comb(ctx, compute_ctx, node); + + break; + + case GGML_OP_DSV4_HC_POST: + ggml_vk_dsv4_hc_post(ctx, compute_ctx, node); + + break; + case GGML_OP_SSM_SCAN: ggml_vk_ssm_scan(ctx, compute_ctx, node); @@ -19311,6 +19454,16 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_GATED_LINEAR_ATTN: // the shader block size is hardcoded to head_size 64 return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64; + case GGML_OP_DSV4_HC_PRE: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32; + case GGML_OP_DSV4_HC_COMB: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; + case GGML_OP_DSV4_HC_POST: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32; case GGML_OP_LIGHTNING_INDEXER: { const ggml_tensor * q = op->src[0]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp new file mode 100644 index 000000000000..0449715bf458 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp @@ -0,0 +1,101 @@ +#version 450 + +// DeepSeek V4 fused hyper-connection "comb": build the 4x4 stream-mixing matrix per token — +// per-source softmax over scaled+biased logits, then eps-stabilized alternating column/row +// (sinkhorn) normalization. Port of ggml-cuda/dsv4-hc.cu (hc_comb): one THREAD per token, +// the whole 4x4 lives in registers. At decode this is a single active thread by design — +// the fusion's value is collapsing the ~dozen decomposed graph ops (and their intermediate +// tensors) into one dispatch, not throughput. + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer MBuf { float data_mixes[]; }; +layout(binding = 1) readonly buffer SBuf { float data_scale[]; }; +layout(binding = 2) readonly buffer BBuf { float data_base[]; }; +layout(binding = 3) writeonly buffer DBuf { float data_d[]; }; + +layout(push_constant) uniform Parameters { + uint n_tokens; + uint sm0, sm1; + uint ss0; + uint sb0; + uint sd0, sd1, sd2; + float eps; + int n_iter; +} p; + +const uint HC = 4; +const uint COMB_OFFSET = 2 * HC; // comb logits start after the pre/post blocks of the mix vector + +void norm_cols(inout float comb[HC * HC]) { + for (uint idst = 0; idst < HC; ++idst) { + float sum = p.eps; + for (uint isrc = 0; isrc < HC; ++isrc) { + sum += comb[idst + HC * isrc]; + } + const float inv_sum = 1.0 / sum; + for (uint isrc = 0; isrc < HC; ++isrc) { + comb[idst + HC * isrc] *= inv_sum; + } + } +} + +void norm_rows(inout float comb[HC * HC]) { + for (uint isrc = 0; isrc < HC; ++isrc) { + float sum = p.eps; + for (uint idst = 0; idst < HC; ++idst) { + sum += comb[idst + HC * isrc]; + } + const float inv_sum = 1.0 / sum; + for (uint idst = 0; idst < HC; ++idst) { + comb[idst + HC * isrc] *= inv_sum; + } + } +} + +void main() { + const uint it = gl_GlobalInvocationID.x; + if (it >= p.n_tokens) { + return; + } + + const float scale_comb = data_scale[2 * p.ss0]; + float comb[HC * HC]; + + for (uint isrc = 0; isrc < HC; ++isrc) { + float vmax = uintBitsToFloat(0xff800000); // -inf + for (uint idst = 0; idst < HC; ++idst) { + const uint idx = idst + HC * isrc; + const float v = data_mixes[(COMB_OFFSET + idx) * p.sm0 + it * p.sm1] * scale_comb + + data_base[(COMB_OFFSET + idx) * p.sb0]; + comb[idx] = v; + vmax = max(vmax, v); + } + + float sum = 0.0; + for (uint idst = 0; idst < HC; ++idst) { + const uint idx = idst + HC * isrc; + const float v = exp(comb[idx] - vmax); + comb[idx] = v; + sum += v; + } + + const float inv_sum = 1.0 / sum; + for (uint idst = 0; idst < HC; ++idst) { + const uint idx = idst + HC * isrc; + comb[idx] = comb[idx] * inv_sum + p.eps; + } + } + + norm_cols(comb); + for (int i = 1; i < p.n_iter; ++i) { + norm_rows(comb); + norm_cols(comb); + } + + for (uint isrc = 0; isrc < HC; ++isrc) { + for (uint idst = 0; idst < HC; ++idst) { + data_d[idst * p.sd0 + isrc * p.sd1 + it * p.sd2] = comb[idst + HC * isrc]; + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp new file mode 100644 index 000000000000..212109dfd434 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp @@ -0,0 +1,44 @@ +#version 450 + +// DeepSeek V4 fused hyper-connection "post": redistribute the layer output back into the HC +// streams with the sinkhorn-mixed residual, +// dst[i0, idst, it] = x[i0, it] * post[idst, it] + sum_isrc residual[i0, isrc, it] * comb[idst, isrc, it]. +// Port of ggml-cuda/dsv4-hc.cu (hc_post). Flat elementwise over n_embd * hc * n_tokens. + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer XBuf { float data_x[]; }; +layout(binding = 1) readonly buffer RBuf { float data_r[]; }; +layout(binding = 2) readonly buffer PBuf { float data_p[]; }; +layout(binding = 3) readonly buffer CBuf { float data_c[]; }; +layout(binding = 4) writeonly buffer DBuf { float data_d[]; }; + +layout(push_constant) uniform Parameters { + uint n_embd; + uint hc; + uint nr; // n_embd * hc * n_tokens + uint sx0, sx1; + uint sr0, sr1, sr2; + uint sp0, sp1; + uint sc0, sc1, sc2; + uint sd0, sd1, sd2; +} p; + +void main() { + const uint ir = gl_GlobalInvocationID.x; + if (ir >= p.nr) { + return; + } + + const uint i0 = ir % p.n_embd; + const uint idst = (ir / p.n_embd) % p.hc; + const uint it = ir / (p.n_embd * p.hc); + + float sum = data_x[i0 * p.sx0 + it * p.sx1] * data_p[idst * p.sp0 + it * p.sp1]; + for (uint isrc = 0; isrc < p.hc; ++isrc) { + sum += data_r[i0 * p.sr0 + isrc * p.sr1 + it * p.sr2] + * data_c[idst * p.sc0 + isrc * p.sc1 + it * p.sc2]; + } + + data_d[i0 * p.sd0 + idst * p.sd1 + it * p.sd2] = sum; +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp new file mode 100644 index 000000000000..b6cc09fe2cfa --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp @@ -0,0 +1,37 @@ +#version 450 + +// DeepSeek V4 fused hyper-connection "pre": mix the HC input streams down to one embedding, +// dst[i0, it] = sum_ih x[i0, ih, it] * w[ih, it]. Port of ggml-cuda/dsv4-hc.cu (hc_pre). +// Flat elementwise kernel over n_embd * n_tokens; strides are in f32 elements. + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer XBuf { float data_x[]; }; +layout(binding = 1) readonly buffer WBuf { float data_w[]; }; +layout(binding = 2) writeonly buffer DBuf { float data_d[]; }; + +layout(push_constant) uniform Parameters { + uint n_embd; + uint hc; + uint nr; // n_embd * n_tokens + uint sx0, sx1, sx2; + uint sw0, sw1; + uint sd0, sd1; +} p; + +void main() { + const uint ir = gl_GlobalInvocationID.x; + if (ir >= p.nr) { + return; + } + + const uint i0 = ir % p.n_embd; + const uint it = ir / p.n_embd; + + float sum = data_x[i0 * p.sx0 + it * p.sx2] * data_w[it * p.sw1]; + for (uint ih = 1; ih < p.hc; ++ih) { + sum += data_x[i0 * p.sx0 + ih * p.sx1 + it * p.sx2] * data_w[ih * p.sw0 + it * p.sw1]; + } + + data_d[i0 * p.sd0 + it * p.sd1] = sum; +} 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 4a75993680fc..57ea2eb804b2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -808,6 +808,9 @@ void process_shaders() { #endif string_to_spv("flash_attn_top_k_f16", "flash_attn_top_k.comp", {}); string_to_spv("flash_attn_gather_f16", "flash_attn_gather.comp", {}); + string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); + string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); + string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {}); string_to_spv("mul_mat_vec_p021_f16_f32_subgroup_add", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}); string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}); From 631fd4858e70d8063361a815fa1711eaa4128d5f Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 2 Aug 2026 16:07:19 +0000 Subject: [PATCH 080/109] llama: keep DeepSeek lightning-indexer key cache f16 under quantized -ctk The fused indexer kernels read f16 keys only, so quantizing the small (128-dim) indexer cache silently disables them and falls back to the decomposed full-KV indexer path, whose contiguize cost grows superlinearly with depth (measured 0.8ms -> 97ms per dispatch by 12k context on Vulkan). Pin the indexer key cache to f16 in both DSA cache variants; the memory cost vs q8_0 is ~120 bytes per token per layer. Assisted-by: Claude Fable 5 --- src/llama-kv-cache-dsa.cpp | 5 ++++- src/llama-kv-cache-dsv4.cpp | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp index 96cb045d2e5d..e926e34314b4 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -47,8 +47,11 @@ llama_kv_cache_dsa::llama_kv_cache_dsa( LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); + // keep indexer keys f16 regardless of type_k: the fused indexer kernels read + // f16 only, and quantizing this small cache (128 dims) saves little while + // forcing the much slower decomposed indexer path kv_lid = std::make_unique( - model, hparams_lid, type_k, type_v, + model, hparams_lid, GGML_TYPE_F16, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_lid, reuse, nullptr); } diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 948d08146fbf..ca46cf1b06e9 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -1305,8 +1305,11 @@ llama_kv_cache_dsv4::llama_kv_cache_dsv4( LLAMA_LOG_INFO("%s: creating DSV4 lightning-indexer KV cache, size = %u cells\n", __func__, dsv4_comp_size(kv_size, DSV4_CSA_RATIO)); + // keep indexer keys f16 regardless of type_k: the fused indexer kernels read + // f16 only, and quantizing this small cache (128 dims) saves little while + // forcing the much slower decomposed indexer path kv_lid = std::make_unique( - model, hparams_lid, type_k, type_v, + model, hparams_lid, GGML_TYPE_F16, type_v, v_trans, offload, unified_compressed, GGML_PAD(dsv4_comp_size(kv_size, DSV4_CSA_RATIO), 256u), n_seq_max, n_pad, 0, LLAMA_SWA_TYPE_NONE, nullptr, filter_csa, nullptr, nullptr); From 0e71642c57a5492375e2f4649dc7a9ed30092c9b Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Mon, 3 Aug 2026 01:48:02 +0000 Subject: [PATCH 081/109] llama: contiguize grouped o-proj input for small multi-token batches (DSv4) The permuted view feeding the grouped wo_a matmul gives B a 128 KiB power-of-2 token stride. Single-token decode never touches it, but small multi-token batches (speculative verify, n=2-4) hit a strided-B matmul path that runs ~11x slower than contiguous (40 vs 460 GFLOPS measured on gfx1151, ~54% of GPU time in a draft-verify window). A contiguous copy for 2..8 tokens is far cheaper than the stride tax; n=1 and large prefill batches are unaffected. Assisted-by: Claude Fable 5 --- src/models/deepseek4.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 08415fe97a12..5501c2cba177 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -1205,6 +1205,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl( out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, nt); out = ggml_permute(ctx0, out, 0, 2, 1, 3); + // small multi-token batches (speculative verify) hit a pathological strided-B + // path in the grouped matmul below; a contiguous copy is much cheaper + if (nt > 1 && nt <= 8) { + out = ggml_cont(ctx0, out); + } ggml_tensor * oa = ggml_mul_mat(ctx0, layer.wo_a, out); cb(oa, "attn_wo_a", il); oa = ggml_permute(ctx0, oa, 0, 2, 1, 3); From aca6680eb64df06f01fb413e8e14cdbea9c6be0f Mon Sep 17 00:00:00 2001 From: Jaap Buurman Date: Wed, 12 Aug 2026 22:25:40 +0200 Subject: [PATCH 082/109] vulkan: accelerate DeepSeek V4 sparse prefill FA Assisted-by: OpenAI Codex --- .../DSV4-vulkan-sparse-prefill-progress.md | 259 +++++++++++++++++ ggml/src/ggml-vulkan/ggml-vulkan.cpp | 16 +- .../vulkan-shaders/flash_attn_top_k_cm.comp | 269 ++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 3 + tests/test-backend-ops.cpp | 1 + 5 files changed, 545 insertions(+), 3 deletions(-) create mode 100644 docs/development/DSV4-vulkan-sparse-prefill-progress.md create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp diff --git a/docs/development/DSV4-vulkan-sparse-prefill-progress.md b/docs/development/DSV4-vulkan-sparse-prefill-progress.md new file mode 100644 index 000000000000..80bf562a0623 --- /dev/null +++ b/docs/development/DSV4-vulkan-sparse-prefill-progress.md @@ -0,0 +1,259 @@ +# DeepSeek V4 Vulkan sparse prefill progress + +This file is a self-contained handoff for the DeepSeek V4 sparse-attention prompt-processing optimization on AMD Strix Halo. Read the repository `AGENTS.md` and `CONTRIBUTING.md` before continuing. + +## Repository state + +- Repository: `https://github.com/Nathanw1014/llama.cpp` +- Branch: `strix-halo-vulkan` +- Starting commit: `baf0025de861c6f6ea3720fa81c52ae1b2e6c078` +- Target GPU: AMD Radeon 8060S / gfx1151, RADV, Vulkan, wave64 +- The device reports `GL_KHR_cooperative_matrix`, f16 inputs with f32 accumulation, 64 KiB shared memory, and a maximum 512-thread workgroup used by this path. + +The implementation is not upstream `ggml-org/llama.cpp`. It builds on this branch's DeepSeek V4 graph, Lightning Indexer, sparse top-K hint, decode gather path, fused HC kernels, and Vulkan profiler changes. + +## Important execution constraints + +The system is an APU. CPU compilation and GPU benchmarking share power and memory bandwidth. Never build and benchmark at the same time. Serialize all builds, correctness tests, and performance tests. + +GPU commands must run with host GPU access. In an agent sandbox, request elevated/out-of-sandbox execution. A sandboxed benchmark showed only CPU activity and is invalid. + +Redirect the full model benchmark to a log. Inspect only the last profiler block with `tail`; do not load the full log into agent context. + +## Build and ccache + +The build directory is `build`, configured as Release with Vulkan enabled. The default ccache directory was read-only in the agent environment, so use a writable directory: + +```bash +cmake -S . -B build \ + -DGGML_VULKAN=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + +CCACHE_DIR=/tmp/llama-cpp-ccache cmake --build build --config Release \ + --target llama-bench test-backend-ops -j "$(nproc)" + +CCACHE_DIR=/tmp/llama-cpp-ccache ccache -s +``` + +ccache was verified active. The final build reported direct hits. Note that an incremental change to `ggml-vulkan.cpp` is one large C++ translation unit and therefore uses one compiler core even with `-j`. Shader object regeneration can run in parallel. + +## Canonical benchmark command + +Do not change or omit switches for the 32k acceptance run: + +```bash +GGML_VK_PERF_LOGGER=1 ./build/bin/llama-bench \ + -m ~/Projects/docker/localLLaMA/models/models--unsloth--DeepSeek-V4-Flash-0731-GGUF/snapshots/109848da2469efe1f1aab9e11acea08a065ccd4f/UD-IQ3_XXS/DeepSeek-V4-Flash-0731-UD-IQ3_XXS-00001-of-00004.gguf \ + -r 1 -d 32768 -p 2048 -ub 2048 -fa 1 -n 0 \ + > /tmp/dsv4-vulkan-32k.log 2>&1 + +tail -n 180 /tmp/dsv4-vulkan-32k.log +``` + +The known model path exists on the target system. + +## Graph and dispatch findings + +DeepSeek V4 builds the Lightning Indexer and sparse attention in `src/models/deepseek4.cpp`: + +1. `build_lid_top_k()` creates indexer Q/K/weights and calls `ggml_lightning_indexer()`. +2. `ggml_top_k()` selects up to `hparams.indexer_top_k` compressed-cache indices for every query token. +3. `build_csa_lid_attention()` concatenates the raw SWA K prefix with compressed CSA K, builds a dense mask carrying the same sparse selection, and calls `build_attn_mha(..., top_k, raw_k->ne[2])`. +4. `build_attn_mha()` attaches `top_k` and `n_kv_raw` to `GGML_OP_FLASH_ATTN_EXT` through `ggml_flash_attn_ext_add_top_k()`. + +At the final 32k PP2048 batch: + +- total K rows: 11,008 +- `n_kv_raw`: 2,304 raw SWA rows, always attended subject to the mask +- `n_top_k`: 512 selected compressed rows per query token +- active rows: 2,816 +- selectable compressed region: 8,704 + +The custom Vulkan path is selected by `ggml_vk_flash_attn_top_k()` before ordinary FA. Its gate requires the DeepSeek V4 shape and `total_k >= 3 * (n_kv_raw + n_top_k)`. The final shape satisfies `11008 >= 3 * 2816`. + +The old `flash_attn_top_k.comp` shader is scalar/subgroup code. One 512-thread workgroup covers eight heads for one query token. It stages 16 selected 512-wide K/V rows, computes QK with scalar FMAs and `subgroupAdd`, updates online softmax one key at a time, and accumulates PV manually. It does not use cooperative matrices. + +The top-K set differs by query token but is shared by all 64 query heads for that token. This makes the attention for one token a regular matrix problem across heads and selected keys despite sparse per-token indexing. + +## Root cause evidence + +The existing Vulkan timestamp infrastructure was extended with `ggml_vk_perf_mark_subop()` after the sparse dispatch. This reports the sparse kernel separately as `FA_TOP_K_SPARSE (sub-op)` or `FA_TOP_K_CM (sub-op)`. + +Focused test shape: + +```bash +GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=32768,nb=512,n_kv_raw=1024,n_top_k=512,sinks=0' +``` + +Results: + +- old scalar sparse kernel: 61.42 ms, 1.68 TFLOPS of useful active-set work +- ordinary dense FA diagnostic (`GGML_VK_FA_TOPK=0`): 255.97 ms, about 8.7 TFLOPS over the full dense work +- final cooperative sparse kernel: 32.55 ms, 3.17 TFLOPS of useful active-set work + +The residual `FLASH_ATTN_EXT` interval after the sparse timestamp is only about 4-7 us. The cost is inside the shader, not dispatch or surrounding synchronization. + +A temporary uniform stage-profiling mode was used and removed. For the final 32-head tile: + +- selected K gather + cooperative QK: 14.30 ms +- gather + QK + serial softmax: 26.34 ms +- gather + QK + parallel softmax: 15.53 ms +- full kernel: 32.55 ms +- the remaining cooperative PV/output portion is about 17.0 ms + +The old scalar shader was compute/issue inefficient. Dense FA proved matrix hardware is much faster but was still too expensive because it processes all K rows. The final implementation preserves sparsity and uses the matrix hardware for both QK and PV. + +## Implementation + +Files changed: + +- `ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp` + - New cooperative-matrix sparse prefill shader. + - One 512-thread workgroup covers 32 query heads for one token. + - Eight wave64 subgroups cover two 16-head tiles by four 16-key or 16-output-dimension tiles. + - Processes 64 selected keys per online-softmax block. + - Stages only indexed selected K/V tiles, never the full K range. + - Uses f16 cooperative-matrix inputs and f32 accumulation for QK and PV. + - Uses a 16-lane segmented softmax per head. XOR subgroup shuffles reduce max and sum for four independent heads per wave without workgroup barriers. + - Keeps f32 output accumulators and normalizes after all active blocks. + - Preserves the raw prefix, top-K index validation, mask, sinks, stream strides, and K == V latent behavior. +- `ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp` + - Embeds the new shader when cooperative-matrix shader support is available. +- `ggml/src/ggml-vulkan/ggml-vulkan.cpp` + - Adds the cooperative sparse pipeline when the device supports the required 16x16x16 f16/f32 cooperative matrix shape. + - Selects it by capability and keeps the scalar shader as fallback. + - Adds `GGML_VK_FA_TOPK=0` to force ordinary dense FA for diagnostics. + - Adds `GGML_VK_FA_TOPK_CM=0` to force the old scalar sparse shader for A/B tests. + - Adds sparse sub-operation timestamps through the existing profiler. + +The path is capability-based, not hardcoded to Strix Halo. The current sparse shape gate remains DeepSeek V4-specific. Devices without the required cooperative matrix support keep the correct scalar sparse or dense fallback. + +Two discarded prototypes are useful context: + +- A 16-head, 64-key streaming cooperative tile was correct and reduced the focused test from 61.4 to 44.6 ms. +- Keeping 32 complete 512-wide K/V rows in LDS grew shared memory to about 41 KiB, reduced residency, doubled block/barrier count, and regressed to 73.8 ms. Do not retry full-row LDS staging without solving occupancy. +- A 32-head, 64-key streaming tile halved irregular row loads but initially stayed near 44.7 ms because serial softmax cost about 11.5 ms. Parallel segmented softmax produced the final 32.55 ms result. + +## Correctness validation + +Run: + +```bash +./build/bin/test-backend-ops test -b Vulkan0 -o FLASH_ATTN_EXT -p 'n_top_k=' +``` + +Final result: 8/8 sparse top-K FA cases passed against the CPU reference. Cases include: + +- decode and short batches that use dense/gather fallback +- sparse prefill batch sizes 64 and 128 +- `n_kv_raw` plus top-K selection +- invalid top-K index handling from the test fixture +- sinks enabled and disabled +- sparse threshold transitions +- an active-key count of 193, which exercises a partial final 64-key block + +The test uses the existing FA tolerance of NMSE <= `5e-4`. No NaN or Inf failure occurred. The implementation changes Q and probability inputs to f16 cooperative-matrix operands with f32 accumulation, matching the precision strategy of ordinary Vulkan cooperative FA. + +Still desirable before broader submission: + +- compare model logits on controlled prompts between `GGML_VK_FA_TOPK_CM=0` and the default cooperative path + +## Canonical 32k results + +Exact clean runs, same command and machine, no concurrent build: + +```text +32k context, PP 2048, ub 2048 + +Before (commit baf0025de): +112.29 tok/s +Total Vulkan: 18.1957 s +Sparse FA: 8.84496 s, 421.189 ms/layer +Lightning Indexer: 1.19438 s +TOP_K: 0.076948 s + +After: +152.32 tok/s +Total Vulkan: 13.4032 s +Sparse FA: 4.44701 s, 211.762 ms/layer +Lightning Indexer: 1.13156 s +TOP_K: 0.073945 s + +Change: +Throughput: +35.65% +Total Vulkan time: -26.34% +Sparse FA time: -49.72% +Sparse FA saved: 4.398 s +Total GPU time saved: 4.793 s +``` + +The profiler now lists the optimized dispatch as `FA_TOP_K_CM (sub-op)`. The following residual `FLASH_ATTN_EXT` line is only the post-mark interval and must not be interpreted as the kernel time. + +## Context-depth measurements + +All points use PP2048, ub2048, FA enabled, one repetition, and no token generation. They were run sequentially with no compiler active: + +```text +Existing depth tok/s Total Vulkan Final large FA Lightning Indexer TOP_K +0 253.44 8.041 s 0.294 s 0.095 s 0.001 s +8192 211.01 9.666 s 1.625 s 0.379 s 0.022 s +16384 177.19 11.518 s 3.031 s 0.678 s 0.044 s +32768 152.32 13.403 s 4.447 s 1.132 s 0.074 s +``` + +At 0, 8k, and 16k, total K is below the existing sparse-path gate `total_k >= 3 * (n_kv_raw + n_top_k)`. These points use the unchanged ordinary dense FA implementation, so the cooperative sparse change does not affect or regress them. At 32k, total K is 11,008 and the cooperative sparse path engages. The 32k `Final large FA` value is the `FA_TOP_K_CM (sub-op)` total; the lower-depth values are the large ordinary `FLASH_ATTN_EXT` totals. + +Logs: + +- `/tmp/dsv4-vulkan-cm-0k.log` +- `/tmp/dsv4-vulkan-cm-8k.log` +- `/tmp/dsv4-vulkan-cm-16k.log` +- `/tmp/dsv4-vulkan-cm-32k.log` +- `/tmp/dsv4-vulkan-baseline-clean.log` +- `/tmp/dsv4-fa-cm-correctness-final.log` + +## Next optimization target + +The cooperative sparse FA remains the largest context-dependent cost at about 4.45 s total. Stage profiling indicates approximately 14.3 ms of focused-test time in gather/QK, about 1.2 ms in parallel softmax, and about 17 ms in PV/output. + +The next useful work is PV and output accumulation, not TOP_K. Investigate: + +- reducing repeated selected V staging across the two 32-head workgroups per token without increasing LDS enough to lose occupancy +- reducing the eight output-dimension passes or retaining more PV state in cooperative fragments/registers +- checking register count and spills for the 32 f32 output accumulators per invocation using RADV shader statistics +- alternate 32-head layouts that keep the same eight-wave occupancy but improve PV scheduling +- query-tile overlap/union gathering only if measured top-K overlap is high enough; a whole-2048-query union is unlikely to help + +Do not optimize TOP_K first. At 32k it is only about 74 ms total. Lightning Indexer is about 1.13 s and is the next context-dependent target only after sparse FA improves further. + +## Useful diagnostics + +Force old scalar sparse path: + +```bash +GGML_VK_FA_TOPK_CM=0 GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=32768,nb=512,n_kv_raw=1024,n_top_k=512,sinks=0' +``` + +Force ordinary dense FA: + +```bash +GGML_VK_FA_TOPK=0 GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=32768,nb=512,n_kv_raw=1024,n_top_k=512,sinks=0' +``` + +Default cooperative sparse path: + +```bash +GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=32768,nb=512,n_kv_raw=1024,n_top_k=512,sinks=0' +``` + +Always run these sequentially. Do not run a compiler concurrently on this APU. diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index a712d2bdd165..25d78b77b59a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1090,6 +1090,7 @@ struct vk_device_struct { vk_pipeline pipeline_lightning_indexer_cm_f16; vk_pipeline pipeline_lightning_indexer_decode_cm_f16; vk_pipeline pipeline_flash_attn_top_k_f16; + vk_pipeline pipeline_flash_attn_top_k_cm_f16; vk_pipeline pipeline_flash_attn_gather_f16; vk_pipeline pipeline_dsv4_hc_pre_f32; vk_pipeline pipeline_dsv4_hc_comb_f32; @@ -6078,6 +6079,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { "lightning_indexer_decode_cm_f16", lightning_indexer_decode_cm_f16_len, lightning_indexer_decode_cm_f16_data, "main", 5, sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 1, 1}, {device->subgroup_size}, 1, true, true, device->subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_top_k_cm_f16, + "flash_attn_top_k_cm_f16", flash_attn_top_k_cm_f16_len, flash_attn_top_k_cm_f16_data, "main", 6, + sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true, + device->subgroup_size); } #endif ggml_vk_create_pipeline(device, device->pipeline_flash_attn_top_k_f16, @@ -11093,7 +11098,9 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { const ggml_tensor * top_k = dst->src[5]; - if (!top_k || !ctx->device->pipeline_flash_attn_top_k_f16 || + static const char * top_k_env = getenv("GGML_VK_FA_TOPK"); + if ((top_k_env && top_k_env[0] == '0') || + !top_k || (!ctx->device->pipeline_flash_attn_top_k_f16 && !ctx->device->pipeline_flash_attn_top_k_cm_f16) || q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || q->ne[0] != 512 || q->ne[1] < 64 || k->ne[0] != 512 || v->ne[0] != 512 || @@ -11143,12 +11150,15 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & const vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q); const vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; - vk_pipeline pipeline = ctx->device->pipeline_flash_attn_top_k_f16; + static const char * top_k_cm_env = getenv("GGML_VK_FA_TOPK_CM"); + const bool use_cm = (!top_k_cm_env || top_k_cm_env[0] != '0') && ctx->device->pipeline_flash_attn_top_k_cm_f16; + vk_pipeline pipeline = use_cm ? ctx->device->pipeline_flash_attn_top_k_cm_f16 : ctx->device->pipeline_flash_attn_top_k_f16; ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, {q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf, ggml_vk_tensor_subbuffer(ctx, top_k), ggml_vk_tensor_subbuffer(ctx, dst)}, - pc, {(uint32_t) q->ne[1], (uint32_t) CEIL_DIV(q->ne[2], 8), (uint32_t) q->ne[3]}); + pc, {(uint32_t) q->ne[1], (uint32_t) CEIL_DIV(q->ne[2], use_cm ? 32 : 8), (uint32_t) q->ne[3]}); + ggml_vk_perf_mark_subop(ctx, subctx, use_cm ? "FA_TOP_K_CM (sub-op)" : "FA_TOP_K_SPARSE (sub-op)"); return true; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp new file mode 100644 index 000000000000..420a05375a7b --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp @@ -0,0 +1,269 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require + +layout(constant_id = 0) const uint WORKGROUP_SIZE = 512; +layout(constant_id = 1) const uint SUBGROUP_SIZE = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer MaskBuf { float16_t data_m[]; }; +layout(binding = 3) readonly buffer SinkBuf { float data_s[]; }; +layout(binding = 4) readonly buffer TopBuf { int data_top[]; }; +layout(binding = 5) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_batch; + uint n_kv; + uint n_kv_raw; + uint n_top_k; + uint n_head; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk1; + uint nbk3; + uint nbm1; + uint nbm3; + uint nbt1; + uint nbt3; + uint nb1; + uint nb2; + uint nb3; + float scale; + uint has_sinks; +} p; + +const uint TILE = 16; +const uint HEAD_SIZE = 512; +const uint HEADS_PER_GROUP = 32; +const uint KEYS_PER_BLOCK = 64; +const uint DIMS_PER_BLOCK = 64; +const uint QK_STRIDE = TILE / 4 + 2; +const uint SCORE_STRIDE = HEADS_PER_GROUP / 4 + 1; +const uint P_STRIDE = KEYS_PER_BLOCK / 4 + 2; +const uint V_STRIDE = DIMS_PER_BLOCK / 4 + 2; +const uint PV_STRIDE = DIMS_PER_BLOCK / 4; +const float MASK_NEG_INF = -65500.0; + +shared uint key_idx[KEYS_PER_BLOCK]; +shared f16vec4 q_sh[HEADS_PER_GROUP * QK_STRIDE]; +shared f16vec4 k_sh[KEYS_PER_BLOCK * QK_STRIDE]; +shared vec4 score_sh[KEYS_PER_BLOCK * SCORE_STRIDE]; +shared f16vec4 p_sh[HEADS_PER_GROUP * P_STRIDE]; +shared f16vec4 v_sh[KEYS_PER_BLOCK * V_STRIDE]; +shared vec4 pv_sh[HEADS_PER_GROUP * PV_STRIDE]; +shared float old_scale_sh[HEADS_PER_GROUP]; +shared float row_max_sh[HEADS_PER_GROUP]; +shared float row_sum_sh[HEADS_PER_GROUP]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint token = gl_WorkGroupID.x; + const uint head_base = gl_WorkGroupID.y * HEADS_PER_GROUP; + const uint stream = gl_WorkGroupID.z; + const uint mask_base = stream * p.nbm3 + token * p.nbm1; + const uint top_base = stream * p.nbt3 + token * p.nbt1; + const uint total_keys = p.n_kv_raw + p.n_top_k; + + float accum[HEADS_PER_GROUP * HEAD_SIZE / WORKGROUP_SIZE]; + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + accum[i] = 0.0; + } + + if (tid < HEADS_PER_GROUP) { + row_max_sh[tid] = uintBitsToFloat(0xff800000); + row_sum_sh[tid] = 0.0; + } + barrier(); + + for (uint kb = 0; kb < total_keys; kb += KEYS_PER_BLOCK) { + if (tid < KEYS_PER_BLOCK) { + const uint selected = kb + tid; + uint key = p.n_kv; + if (selected < p.n_kv_raw) { + key = selected; + } else if (selected < total_keys) { + const int compressed = data_top[top_base + selected - p.n_kv_raw]; + if (compressed >= 0 && uint(compressed) < p.n_kv - p.n_kv_raw) { + key = p.n_kv_raw + uint(compressed); + } + } + key_idx[tid] = key; + } + barrier(); + + coopmat scores = + coopmat(0.0); + coopmat kmat; + coopmat qmat; + + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + if (tid < KEYS_PER_BLOCK * (TILE / 4)) { + const uint key_local = tid / (TILE / 4); + const uint d4 = tid % (TILE / 4); + const uint key = key_idx[key_local]; + f16vec4 value = f16vec4(0.0); + if (key < p.n_kv) { + const uint offset = stream * p.nbk3 + key * p.nbk1 + d + d4 * 4; + value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); + } + k_sh[key_local * QK_STRIDE + d4] = value; + } + if (tid < HEADS_PER_GROUP * (TILE / 4)) { + const uint head_local = tid / (TILE / 4); + const uint d4 = tid % (TILE / 4); + const uint head = head_base + head_local; + const uint offset = stream * p.nbq3 + head * p.nbq2 + token * p.nbq1 + d + d4 * 4; + q_sh[head_local * QK_STRIDE + d4] = f16vec4( + data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); + } + barrier(); + + const uint key_chunk = gl_SubgroupID % (KEYS_PER_BLOCK / TILE); + const uint head_tile = gl_SubgroupID / (KEYS_PER_BLOCK / TILE); + coopMatLoad(kmat, k_sh, key_chunk * TILE * QK_STRIDE, + QK_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(qmat, q_sh, head_tile * TILE * QK_STRIDE, + QK_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + scores = coopMatMulAdd(kmat, qmat, scores); + barrier(); + } + + const uint score_key_chunk = gl_SubgroupID % (KEYS_PER_BLOCK / TILE); + const uint score_head_tile = gl_SubgroupID / (KEYS_PER_BLOCK / TILE); + coopMatStore(scores, score_sh, + score_key_chunk * TILE * SCORE_STRIDE + score_head_tile * (TILE / 4), + SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + barrier(); + + { + const uint head_local = tid / (SUBGROUP_SIZE / 4); + const uint softmax_lane = tid % (SUBGROUP_SIZE / 4); + float block_max = uintBitsToFloat(0xff800000); + [[unroll]] for (uint i = 0; i < KEYS_PER_BLOCK / (SUBGROUP_SIZE / 4); ++i) { + const uint key_local = softmax_lane + i * (SUBGROUP_SIZE / 4); + const uint key = key_idx[key_local]; + const float mask = key < p.n_kv ? float(data_m[mask_base + key]) : uintBitsToFloat(0xff800000); + const float score = float(score_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; + block_max = mask < MASK_NEG_INF ? block_max : max(block_max, score); + } + [[unroll]] for (uint delta = 1; delta < SUBGROUP_SIZE / 4; delta *= 2) { + block_max = max(block_max, subgroupShuffleXor(block_max, delta)); + } + + const float old_row_max = row_max_sh[head_local]; + const float old_row_sum = row_sum_sh[head_local]; + const float new_max = max(old_row_max, block_max); + const float old_scale = old_row_sum == 0.0 ? 0.0 : exp(old_row_max - new_max); + float block_sum = 0.0; + [[unroll]] for (uint i = 0; i < KEYS_PER_BLOCK / (SUBGROUP_SIZE / 4); ++i) { + const uint key_local = softmax_lane + i * (SUBGROUP_SIZE / 4); + const uint key = key_idx[key_local]; + const float mask = key < p.n_kv ? float(data_m[mask_base + key]) : uintBitsToFloat(0xff800000); + float weight = 0.0; + if (mask >= MASK_NEG_INF) { + const float score = float(score_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; + weight = exp(score - new_max); + block_sum += weight; + } + p_sh[head_local * P_STRIDE + key_local / 4][key_local % 4] = float16_t(weight); + } + [[unroll]] for (uint delta = 1; delta < SUBGROUP_SIZE / 4; delta *= 2) { + block_sum += subgroupShuffleXor(block_sum, delta); + } + + if (softmax_lane == 0) { + row_sum_sh[head_local] = old_row_sum * old_scale + block_sum; + row_max_sh[head_local] = new_max; + old_scale_sh[head_local] = old_scale; + } + } + barrier(); + + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + const uint out_idx = tid + i * WORKGROUP_SIZE; + const uint head_local = out_idx / HEAD_SIZE; + accum[i] *= old_scale_sh[head_local]; + } + + [[unroll]] for (uint dim_base = 0; dim_base < HEAD_SIZE; dim_base += DIMS_PER_BLOCK) { + [[unroll]] for (uint idx = tid; idx < KEYS_PER_BLOCK * (DIMS_PER_BLOCK / 4); idx += WORKGROUP_SIZE) { + const uint key_local = idx / (DIMS_PER_BLOCK / 4); + const uint d4 = idx % (DIMS_PER_BLOCK / 4); + const uint key = key_idx[key_local]; + f16vec4 value = f16vec4(0.0); + if (key < p.n_kv) { + const uint offset = stream * p.nbk3 + key * p.nbk1 + dim_base + d4 * 4; + value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); + } + v_sh[key_local * V_STRIDE + d4] = value; + } + barrier(); + + coopmat pv = + coopmat(0.0); + coopmat pmat; + coopmat vmat; + + const uint pv_head_tile = gl_SubgroupID / (DIMS_PER_BLOCK / TILE); + const uint pv_dim_tile = gl_SubgroupID % (DIMS_PER_BLOCK / TILE); + [[unroll]] for (uint key_chunk = 0; key_chunk < KEYS_PER_BLOCK / TILE; ++key_chunk) { + coopMatLoad(pmat, p_sh, + pv_head_tile * TILE * P_STRIDE + key_chunk * (TILE / 4), + P_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(vmat, v_sh, + key_chunk * TILE * V_STRIDE + pv_dim_tile * (TILE / 4), + V_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + pv = coopMatMulAdd(pmat, vmat, pv); + } + + coopMatStore(pv, pv_sh, + pv_head_tile * TILE * PV_STRIDE + pv_dim_tile * (TILE / 4), PV_STRIDE, + gl_CooperativeMatrixLayoutRowMajor); + barrier(); + + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + const uint out_idx = tid + i * WORKGROUP_SIZE; + const uint head_local = out_idx / HEAD_SIZE; + const uint dim = out_idx % HEAD_SIZE; + if (dim >= dim_base && dim < dim_base + DIMS_PER_BLOCK) { + accum[i] += pv_sh[head_local * PV_STRIDE + (dim - dim_base) / 4][dim % 4]; + } + } + barrier(); + } + } + + if (p.has_sinks != 0 && tid < HEADS_PER_GROUP) { + const float sink = data_s[head_base + tid]; + const float new_max = max(row_max_sh[tid], sink); + const float old_scale = row_sum_sh[tid] == 0.0 ? 0.0 : exp(row_max_sh[tid] - new_max); + row_sum_sh[tid] = row_sum_sh[tid] * old_scale + exp(sink - new_max); + row_max_sh[tid] = new_max; + old_scale_sh[tid] = old_scale; + } + barrier(); + + if (p.has_sinks != 0) { + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + const uint out_idx = tid + i * WORKGROUP_SIZE; + accum[i] *= old_scale_sh[out_idx / HEAD_SIZE]; + } + } + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + const uint out_idx = tid + i * WORKGROUP_SIZE; + const uint head_local = out_idx / HEAD_SIZE; + const uint dim = out_idx % HEAD_SIZE; + const uint dst_base = stream * p.nb3 + token * p.nb2 + (head_base + head_local) * p.nb1; + const float inv_sum = row_sum_sh[head_local] == 0.0 ? 0.0 : 1.0 / row_sum_sh[head_local]; + data_dst[dst_base + dim] = accum[i] * inv_sum; + } +} 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 57ea2eb804b2..845720ad8b1e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -807,6 +807,9 @@ void process_shaders() { string_to_spv("lightning_indexer_decode_cm_f16", "lightning_indexer_decode_cm.comp", {}); #endif string_to_spv("flash_attn_top_k_f16", "flash_attn_top_k.comp", {}); +#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + string_to_spv("flash_attn_top_k_cm_f16", "flash_attn_top_k_cm.comp", {}); +#endif string_to_spv("flash_attn_gather_f16", "flash_attn_gather.comp", {}); string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a8d57ae43261..008760f47754 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10288,6 +10288,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k( 512, 4, 64, 128, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, true)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(1024, 64, 65, 128, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 128, 256, 512, false)); return test_cases; From 10c54572d363f37111eb227c6650ac73b85c43af Mon Sep 17 00:00:00 2001 From: Jaap Buurman Date: Thu, 13 Aug 2026 11:12:21 +0200 Subject: [PATCH 083/109] vulkan: split sparse prefill attention Assisted-by: Codex --- .../DSV4-vulkan-sparse-prefill-progress.md | 92 +++++++++++++++++ ggml/src/ggml-vulkan/ggml-vulkan.cpp | 99 ++++++++++++++++++- .../vulkan-shaders/flash_attn_base.glsl | 39 ++++---- .../vulkan-shaders/flash_attn_cm1.comp | 16 +-- .../vulkan-shaders/flash_attn_top_k.comp | 2 + .../vulkan-shaders/flash_attn_top_k_cm.comp | 38 +++++-- tests/test-backend-ops.cpp | 1 + 7 files changed, 251 insertions(+), 36 deletions(-) diff --git a/docs/development/DSV4-vulkan-sparse-prefill-progress.md b/docs/development/DSV4-vulkan-sparse-prefill-progress.md index 80bf562a0623..93c42476cfc1 100644 --- a/docs/development/DSV4-vulkan-sparse-prefill-progress.md +++ b/docs/development/DSV4-vulkan-sparse-prefill-progress.md @@ -257,3 +257,95 @@ GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ ``` Always run these sequentially. Do not run a compiler concurrently on this APU. + +## Experimental raw-prefix split prototype + +An uncommitted follow-up prototype was tested after commit `24c7ead76cc1a9631fa1b42b0bfa53a15169e1ba`. Check `git status` before continuing. It was initially tested with `GGML_VK_FA_TOPK_SPLIT=1`. After the successful 32k run, the split path was changed to default-on for a llama-server coherence test. Set `GGML_VK_FA_TOPK_SPLIT=0` to restore the single cooperative sparse kernel. + +Stage profiling on the exact production sparse shape (`kv=19200,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0`) showed: + +```text +cooperative QK and selected-K gather: 88.885 ms +softmax increment: 4.909 ms +cooperative PV and output increment: 128.399 ms +full cooperative sparse kernel: 222.193 ms +``` + +PV and output were 57.8% of the kernel. More importantly, most active keys are not sparse: all 2,304 raw-prefix rows are contiguous, while only 512 compressed rows use top-K indices. The prototype therefore makes two attention partitions: + +1. Ordinary optimized Vulkan cooperative FA processes the contiguous raw prefix. +2. The cooperative top-K shader processes only the 512 selected compressed rows. +3. The existing split-K reduction combines both online-softmax partitions and applies sinks. + +This preserves the exact raw-prefix, sparse top-K, causal mask, and softmax semantics. It reuses the existing ordinary FA and split-K reduction rather than adding a new subsystem. + +The exact-shape microbenchmark improved from 222.19 ms to 111.33 ms. Its steady split stages were about 62-64 ms raw-prefix FA, 43-46 ms selected sparse FA, and 3.6-3.9 ms reduction. + +Correctness validation: + +```bash +GGML_VK_FA_TOPK_SPLIT=1 ./build/bin/test-backend-ops test \ + -b Vulkan0 -o FLASH_ATTN_EXT -p 'n_top_k=' + +./build/bin/test-backend-ops test -b Vulkan0 -o FLASH_ATTN_EXT +``` + +Results were 8/8 sparse top-K cases and 13,296/13,296 complete Vulkan FA cases against the CPU reference. No NaN or Inf failure occurred. + +Canonical 32k prototype command: + +```bash +GGML_VK_FA_TOPK_SPLIT=1 GGML_VK_PERF_LOGGER=1 ./build/bin/llama-bench \ + -m ~/Projects/docker/localLLaMA/models/models--unsloth--DeepSeek-V4-Flash-0731-GGUF/snapshots/109848da2469efe1f1aab9e11acea08a065ccd4f/UD-IQ3_XXS/DeepSeek-V4-Flash-0731-UD-IQ3_XXS-00001-of-00004.gguf \ + -r 1 -d 32768 -p 2048 -ub 2048 -fa 1 -n 0 \ + > /tmp/dsv4-vulkan-split-32k.log 2>&1 +``` + +Final 32k result: + +```text +32k context, PP 2048, ub 2048 + +Old scalar: 112.29 tok/s, 18.1957 s total, 8.84496 s sparse FA +Committed coopmat: 152.32 tok/s, 13.4032 s total, 4.44701 s sparse FA +Experimental split: 208.70 tok/s, 9.7704 s total, 1.20170 s split sparse FA + +Experimental split stages: +raw-prefix FA: 0.210775 s total, 10.037 ms/layer +selected sparse FA: 0.908190 s total, 43.247 ms/layer +split reduction: 0.082732 s total, 3.940 ms/layer +Lightning Indexer: 1.110410 s total, 52.877 ms/layer +TOP_K: 0.077394 s total, 3.685 ms/layer +``` + +Relative to the committed cooperative path, throughput improved 37.0%, total Vulkan time fell 27.1%, and sparse FA time fell 73.0%. Relative to the old scalar path, throughput improved 85.9%, total Vulkan time fell 46.3%, and sparse FA time fell 86.4%. + +The main unresolved tradeoff is scratch memory. At PP2048, the two output partitions use about 539 MiB because each stores an f32 partial output for 512 dimensions x 64 heads x 2048 queries, plus L/M data. The device maximum storage-buffer range is checked and the code falls back when the allocation is unavailable. The path is temporarily default-on for a llama-server coherence test but should not be finalized without discussing this footprint. A likely next step is to avoid materializing both full output partitions, for example by directly merging the selected partition into the raw result or processing query tiles, while retaining the measured split-path speed. + +The old sparse selection heuristic required `total_k >= 3 * active_k`. For the coherence test it now selects sparse attention whenever `total_k > active_k`; equality still uses dense FA because no keys are pruned. The shape, capability, and allocation gates remain unchanged. + +Prototype logs: + +- `/tmp/dsv4-fa-exact-qk.log` +- `/tmp/dsv4-fa-exact-softmax.log` +- `/tmp/dsv4-fa-exact-full.log` +- `/tmp/dsv4-fa-exact-split.log` +- `/tmp/dsv4-fa-split-correctness.log` +- `/tmp/dsv4-fa-all-correctness.log` +- `/tmp/dsv4-vulkan-split-32k.log` + +## Llama-server coherence check + +After making the split path default-on and changing the sparse crossover to `total_k > active_k`, `llama-server` produced a coherent response from a 2,044-token prompt at significant context depth. The final profiler block confirmed that all 21 sparse-attention layers used the split path: + +```text +FA_TOP_K_RAW: 0.179223 s total, 8.534 ms/layer +FA_TOP_K_SELECTED: 0.886629 s total, 42.220 ms/layer +FA_TOP_K_REDUCE: 0.083240 s total, 3.964 ms/layer +Split sparse FA: 1.149092 s total, 54.719 ms/layer +Lightning Indexer: 1.054610 s total, 50.220 ms/layer +TOP_K: 0.044540 s total, 2.121 ms/layer +Total Vulkan: 9.797560 s +``` + +This closely matches the canonical PP2048 llama-bench result of 9.770 s total and 1.202 s split sparse FA. The focused sparse CPU-reference test was rerun after the crossover change and passed 8/8 cases. diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 25d78b77b59a..141e7098496a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1970,6 +1970,8 @@ struct vk_op_flash_attn_top_k_push_constants { uint32_t nb1, nb2, nb3; float scale; uint32_t has_sinks; + uint32_t profile_stage; + uint32_t split_mode; }; static_assert(sizeof(vk_op_flash_attn_top_k_push_constants) <= 128); @@ -11126,11 +11128,11 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & return false; } const int64_t n_kv_active = n_kv_raw + top_k->ne[0]; - if (k->ne[1] < 3 * n_kv_active) { + if (k->ne[1] <= n_kv_active) { return false; } - const vk_op_flash_attn_top_k_push_constants pc = { + vk_op_flash_attn_top_k_push_constants pc = { (uint32_t) q->ne[1], (uint32_t) k->ne[1], (uint32_t) n_kv_raw, (uint32_t) top_k->ne[0], (uint32_t) q->ne[2], (uint32_t) (q->nb[1] / sizeof(float)), @@ -11145,14 +11147,105 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & (uint32_t) (dst->nb[1] / sizeof(float)), (uint32_t) (dst->nb[2] / sizeof(float)), (uint32_t) (dst->nb[3] / sizeof(float)), - scale, sinks != nullptr, + scale, sinks != nullptr, 0, 0, }; const vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q); const vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; static const char * top_k_cm_env = getenv("GGML_VK_FA_TOPK_CM"); const bool use_cm = (!top_k_cm_env || top_k_cm_env[0] != '0') && ctx->device->pipeline_flash_attn_top_k_cm_f16; + static const char * top_k_profile_env = getenv("GGML_VK_FA_TOPK_PROFILE"); + pc.profile_stage = use_cm && top_k_profile_env ? atoi(top_k_profile_env) : 0; vk_pipeline pipeline = use_cm ? ctx->device->pipeline_flash_attn_top_k_cm_f16 : ctx->device->pipeline_flash_attn_top_k_f16; + + static const char * top_k_split_env = getenv("GGML_VK_FA_TOPK_SPLIT"); + const uint32_t mask_stride = (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)); + const bool try_split = use_cm && (!top_k_split_env || top_k_split_env[0] != '0') && n_kv_raw > 0 && top_k->ne[0] > 0 && mask_stride <= 0xffff; + if (try_split) { + const uint32_t N = (uint32_t) q->ne[1]; + const uint32_t D = 512; + const uint32_t NH = 64; + const uint32_t NS = (uint32_t) q->ne[3]; + const uint32_t raw_kv = (uint32_t) n_kv_raw; + const uint32_t partitions = 2; + const bool f32acc = true; + vk_fa_tuning_params tuning = get_fa_tuning_params(ctx->device, D, D, N, raw_kv, GGML_TYPE_F16, GGML_TYPE_F16, f32acc); + + const uint32_t q_stride = (uint32_t) (q->nb[1] / sizeof(float)); + const uint32_t k_stride = (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)); + const bool aligned = raw_kv % tuning.block_cols == 0 && (q_stride & 7) == 0 && (k_stride & 7) == 0; + const vk_fa_pipeline_state raw_state = get_fa_pipeline_state(ctx->device, tuning, D, D, aligned, f32acc, + true, false, false, GGML_TYPE_F16, GGML_TYPE_F16); + if (raw_state.path == FA_COOPMAT1 && ctx->device->pipeline_flash_attn_split_k_reduce) { + vk_pipeline raw_pipeline; + { + std::lock_guard guard(ctx->device->compile_mutex); + auto & pipelines = ctx->device->pipeline_flash_attn_f32_f16; + auto it = pipelines.find(raw_state); + if (it != pipelines.end()) { + raw_pipeline = it->second; + } else { + pipelines[raw_state] = raw_pipeline = std::make_shared(); + } + } + ggml_pipeline_request_descriptor_sets(ctx, raw_pipeline, 1); + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1); + + const uint64_t split_size = ((uint64_t) D * NH * sizeof(float) + NH * 2 * sizeof(float)) * partitions * N * NS; + if (split_size <= ctx->device->properties.limits.maxStorageBufferRange) { + if (ctx->prealloc_size_split_k < split_size) { + ctx->prealloc_size_split_k = split_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_split_k_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + + const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + const uint32_t n_head_log2 = 64; + const uint32_t packed_gqa = (mask_stride << 16) | 1; + const uint32_t packed_partitions = (partitions << 16) | 1; + const vk_flash_attn_push_constants raw_pc = { + N, raw_kv, + NH, N, NS, + NH, NS, + 1, NS, + 1, NS, + (uint32_t) mask->ne[1], (uint32_t) mask->ne[2], (uint32_t) mask->ne[3], + q_stride, (uint32_t) q->nb[2], (uint32_t) q->nb[3], + k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], + k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], + scale, 0.0f, 0.0f, + n_head_log2, 1.0f, 1.0f, + packed_gqa, raw_kv, packed_partitions, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, raw_pipeline, + {q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, k), + ggml_vk_tensor_subbuffer(ctx, mask), q_buf, split_buf, q_buf}, + raw_pc, {N, NH, NS}); + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_RAW (sub-op)"); + + pc.split_mode = 1; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf, + ggml_vk_tensor_subbuffer(ctx, top_k), split_buf}, + pc, {N, (uint32_t) CEIL_DIV(q->ne[2], 32), NS}); + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_SELECTED (sub-op)"); + + ggml_vk_sync_buffers(ctx, subctx); + const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = {D, NH, N, NS, partitions, sinks != nullptr}; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, + {split_buf, sinks_buf, ggml_vk_tensor_subbuffer(ctx, dst)}, + reduce_pc, {NH, D, N * NS}); + ctx->prealloc_split_k_need_sync = true; + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_REDUCE (sub-op)"); + return true; + } + } + } + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, {q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf, diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 0ce4503a8847..800a79a97d17 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -124,7 +124,7 @@ ACC_TYPE perElemOpStoreCol0(const in uint32_t r, const in uint32_t c, const in A // Load the slope matrix, indexed by Q's dimension 2. ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2) { - const uint32_t h = iq2 + (r % p.gqa_ratio); + const uint32_t h = iq2 + (r % (p.gqa_ratio & 0xffff)); uint32_t n_head_log2 = p.mask_n_head_log2 & N_LOG2_MASK; @@ -137,32 +137,40 @@ ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const i // Load the sink value, indexed by Q's dimension 2. ACC_TYPE perElemOpGetSink(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2) { - const uint32_t h = iq2 + (r % p.gqa_ratio); + const uint32_t h = iq2 + (r % (p.gqa_ratio & 0xffff)); return ACC_TYPE(data_s[h]); } uint32_t i, N, KV, split_k_index, Tr, start_j, end_j, gqa_iq1, iq2, iq3, rk2, rk3, rv2, rv3, ik2, ik3, iv2, iv3, - q_stride, k_stride, v_stride, m_stride; + q_stride, k_stride, v_stride, m_stride, gqa_ratio, split_k_num, output_k_num; +bool partial_output; void init_indices() { N = p.N; KV = p.KV; + gqa_ratio = p.gqa_ratio & 0xffff; + split_k_num = p.k_num & 0xffff; + output_k_num = p.k_num >> 16; + partial_output = output_k_num != 0; + if (!partial_output) { + output_k_num = split_k_num; + } - if (p.k_num > 1) { - if (p.gqa_ratio > 1) { + if (split_k_num > 1) { + if (gqa_ratio > 1) { i = 0; // batch and split_k share gl_WorkGroupID.x - gqa_iq1 = gl_WorkGroupID.x / p.k_num; - split_k_index = gl_WorkGroupID.x % p.k_num; + gqa_iq1 = gl_WorkGroupID.x / split_k_num; + split_k_index = gl_WorkGroupID.x % split_k_num; } else { gqa_iq1 = 0; - split_k_index = gl_WorkGroupID.x % p.k_num; - i = gl_WorkGroupID.x / p.k_num; + split_k_index = gl_WorkGroupID.x % split_k_num; + i = gl_WorkGroupID.x / split_k_num; } - } else if (p.gqa_ratio > 1) { + } else if (gqa_ratio > 1) { i = 0; gqa_iq1 = gl_WorkGroupID.x; split_k_index = 0; @@ -179,7 +187,7 @@ void init_indices() // When not using grouped query attention, all rows share the same iq2, equal to gl_WorkGroupID.y. // When using grouped query attention, each workgroup does gqa_ratio consecutive values of iq2. - iq2 = gl_WorkGroupID.y * p.gqa_ratio; + iq2 = gl_WorkGroupID.y * gqa_ratio; iq3 = gl_WorkGroupID.z; // broadcast factors @@ -200,14 +208,11 @@ void init_indices() // nb?1 are already divided by the type size and are in units of elements. // When using grouped query attention, Q is indexed by iq2, so the stride // should be nb02 (which is in bytes). - q_stride = p.gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01; + q_stride = gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01; k_stride = p.nb11; v_stride = p.nb21; - // When using grouped query attention, all rows use the same mask (stride 0). - // "p.gqa_ratio >> 16" is just a roundabout way of writing zero - // that prevents the compiler from folding the "&" through the select - // and breaking the alignment detection. - m_stride = (p.gqa_ratio > 1) ? (p.gqa_ratio >> 16) : KV; + const uint32_t mask_stride_override = p.gqa_ratio >> 16; + m_stride = mask_stride_override != 0 ? mask_stride_override : (gqa_ratio > 1 ? 0 : KV); } // Bias applied to softmax to stay in fp16 range. 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 057ed739aa8d..e3ea909a59ce 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -169,7 +169,7 @@ void main() { } // Only load if the block is not all zeros if (mask_opt_bits != MASK_OPT_ALL_ZERO) { - bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; + bool nem1_bounds_check = !(gqa_ratio > 1) && (p.nem1 % Br) != 0; float max_mask = NEG_FLT_MAX_OVER_2; [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { @@ -533,10 +533,10 @@ void main() { // If there is split_k, then the split_k resolve shader does the final // division by L. Store the intermediate O value and per-row m and L values. - if (p.k_num > 1) { - if (p.gqa_ratio > 1) { + if (partial_output || split_k_num > 1) { + if (gqa_ratio > 1) { // note: O and Q have swapped coord 1,2. - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { @@ -549,7 +549,7 @@ void main() { } } - o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); + o_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)); [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { perElemOpStoreCol0(tile_row(r), 0u, ACC_TYPE(Lf[r]), o_offset, iq2, N); @@ -562,7 +562,7 @@ void main() { const uint global_row = i * Br + row; if (global_row < N) { - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)) / 4; + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)) / 4; [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { const uint d = d0 + col_tid; @@ -572,7 +572,7 @@ void main() { } if (global_row < N && col_tid == 0) { - uint32_t lm_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); + uint32_t lm_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)); data_o[lm_offset + iq2] = D_TYPE(Lf[r]); data_o[lm_offset + p.ne1 + iq2] = D_TYPE(Mf[r]); } @@ -621,7 +621,7 @@ void main() { uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4; - if (p.gqa_ratio > 1) { + if (gqa_ratio > 1) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp index b8b49c677bd5..62807471bdd1 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp @@ -37,6 +37,8 @@ layout(push_constant) uniform Parameters { uint nb3; float scale; uint has_sinks; + uint profile_stage; + uint split_mode; } p; // Shape constants pinned by the dispatch gate in ggml_vk_flash_attn_top_k: DeepSeek V4 diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp index 420a05375a7b..f3dfbe59ac00 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp @@ -39,6 +39,8 @@ layout(push_constant) uniform Parameters { uint nb3; float scale; uint has_sinks; + uint profile_stage; + uint split_mode; } p; const uint TILE = 16; @@ -71,7 +73,7 @@ void main() { const uint stream = gl_WorkGroupID.z; const uint mask_base = stream * p.nbm3 + token * p.nbm1; const uint top_base = stream * p.nbt3 + token * p.nbt1; - const uint total_keys = p.n_kv_raw + p.n_top_k; + const uint total_keys = p.split_mode != 0 ? p.n_top_k : p.n_kv_raw + p.n_top_k; float accum[HEADS_PER_GROUP * HEAD_SIZE / WORKGROUP_SIZE]; [[unroll]] for (uint i = 0; i < accum.length(); ++i) { @@ -88,10 +90,11 @@ void main() { if (tid < KEYS_PER_BLOCK) { const uint selected = kb + tid; uint key = p.n_kv; - if (selected < p.n_kv_raw) { + if (p.split_mode == 0 && selected < p.n_kv_raw) { key = selected; } else if (selected < total_keys) { - const int compressed = data_top[top_base + selected - p.n_kv_raw]; + const uint top_pos = p.split_mode != 0 ? selected : selected - p.n_kv_raw; + const int compressed = data_top[top_base + top_pos]; if (compressed >= 0 && uint(compressed) < p.n_kv - p.n_kv_raw) { key = p.n_kv_raw + uint(compressed); } @@ -144,6 +147,10 @@ void main() { SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); barrier(); + if (p.profile_stage == 1) { + continue; + } + { const uint head_local = tid / (SUBGROUP_SIZE / 4); const uint softmax_lane = tid % (SUBGROUP_SIZE / 4); @@ -194,6 +201,10 @@ void main() { accum[i] *= old_scale_sh[head_local]; } + if (p.profile_stage == 2) { + continue; + } + [[unroll]] for (uint dim_base = 0; dim_base < HEAD_SIZE; dim_base += DIMS_PER_BLOCK) { [[unroll]] for (uint idx = tid; idx < KEYS_PER_BLOCK * (DIMS_PER_BLOCK / 4); idx += WORKGROUP_SIZE) { const uint key_local = idx / (DIMS_PER_BLOCK / 4); @@ -242,7 +253,7 @@ void main() { } } - if (p.has_sinks != 0 && tid < HEADS_PER_GROUP) { + if (p.split_mode == 0 && p.has_sinks != 0 && tid < HEADS_PER_GROUP) { const float sink = data_s[head_base + tid]; const float new_max = max(row_max_sh[tid], sink); const float old_scale = row_sum_sh[tid] == 0.0 ? 0.0 : exp(row_max_sh[tid] - new_max); @@ -252,7 +263,7 @@ void main() { } barrier(); - if (p.has_sinks != 0) { + if (p.split_mode == 0 && p.has_sinks != 0) { [[unroll]] for (uint i = 0; i < accum.length(); ++i) { const uint out_idx = tid + i * WORKGROUP_SIZE; accum[i] *= old_scale_sh[out_idx / HEAD_SIZE]; @@ -262,8 +273,19 @@ void main() { const uint out_idx = tid + i * WORKGROUP_SIZE; const uint head_local = out_idx / HEAD_SIZE; const uint dim = out_idx % HEAD_SIZE; - const uint dst_base = stream * p.nb3 + token * p.nb2 + (head_base + head_local) * p.nb1; - const float inv_sum = row_sum_sh[head_local] == 0.0 ? 0.0 : 1.0 / row_sum_sh[head_local]; - data_dst[dst_base + dim] = accum[i] * inv_sum; + if (p.split_mode != 0) { + const uint part_idx = 1; + const uint matrix_base = HEAD_SIZE * p.n_head * (part_idx + 2 * (token + p.n_batch * stream)); + const uint lm_base = HEAD_SIZE * p.n_head * p.n_batch * 2 + p.n_head * 2 * (part_idx + 2 * (token + p.n_batch * stream)); + data_dst[matrix_base + (head_base + head_local) * HEAD_SIZE + dim] = accum[i]; + if (dim == 0) { + data_dst[lm_base + head_base + head_local] = row_sum_sh[head_local]; + data_dst[lm_base + p.n_head + head_base + head_local] = row_max_sh[head_local]; + } + } else { + const uint dst_base = stream * p.nb3 + token * p.nb2 + (head_base + head_local) * p.nb1; + const float inv_sum = row_sum_sh[head_local] == 0.0 ? 0.0 : 1.0 / row_sum_sh[head_local]; + data_dst[dst_base + dim] = accum[i] * inv_sum; + } } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 008760f47754..2b6ae01e6eed 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10716,6 +10716,7 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, nb, 1024, 512, false)); } } + test_cases.emplace_back(new test_flash_attn_ext_top_k(19200, 2048, 2304, 512, false)); return test_cases; } From aa5085f9746e56fb2356253018f5b8c942cc97e7 Mon Sep 17 00:00:00 2001 From: Jaap Buurman Date: Thu, 13 Aug 2026 11:50:12 +0200 Subject: [PATCH 084/109] vulkan: tile sparse prefill scratch Assisted-by: Codex --- .../DSV4-vulkan-sparse-prefill-progress.md | 58 ++++++++++- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 97 +++++++++++-------- tests/test-backend-ops.cpp | 1 + 3 files changed, 117 insertions(+), 39 deletions(-) diff --git a/docs/development/DSV4-vulkan-sparse-prefill-progress.md b/docs/development/DSV4-vulkan-sparse-prefill-progress.md index 93c42476cfc1..5ace979a08df 100644 --- a/docs/development/DSV4-vulkan-sparse-prefill-progress.md +++ b/docs/development/DSV4-vulkan-sparse-prefill-progress.md @@ -320,7 +320,7 @@ TOP_K: 0.077394 s total, 3.685 ms/layer Relative to the committed cooperative path, throughput improved 37.0%, total Vulkan time fell 27.1%, and sparse FA time fell 73.0%. Relative to the old scalar path, throughput improved 85.9%, total Vulkan time fell 46.3%, and sparse FA time fell 86.4%. -The main unresolved tradeoff is scratch memory. At PP2048, the two output partitions use about 539 MiB because each stores an f32 partial output for 512 dimensions x 64 heads x 2048 queries, plus L/M data. The device maximum storage-buffer range is checked and the code falls back when the allocation is unavailable. The path is temporarily default-on for a llama-server coherence test but should not be finalized without discussing this footprint. A likely next step is to avoid materializing both full output partitions, for example by directly merging the selected partition into the raw result or processing query tiles, while retaining the measured split-path speed. +The initial split implementation used 538,968,064 bytes (514 MiB) of scratch at PP2048 because each of two partitions stored an f32 partial output for 512 dimensions x 64 heads x 2048 queries, plus L/M data. This was subsequently reduced by query tiling as described below. The old sparse selection heuristic required `total_k >= 3 * active_k`. For the coherence test it now selects sparse attention whenever `total_k > active_k`; equality still uses dense FA because no keys are pruned. The shape, capability, and allocation gates remain unchanged. @@ -349,3 +349,59 @@ Total Vulkan: 9.797560 s ``` This closely matches the canonical PP2048 llama-bench result of 9.770 s total and 1.202 s split sparse FA. The focused sparse CPU-reference test was rerun after the crossover change and passed 8/8 cases. + +## Tiled split scratch optimization + +Commit `4bbe53e4775f0707de8158e4977a16ab770829da` used two full PP2048 output partitions. The same two-partition algorithm now processes at most 256 query tokens per tile and reuses the split scratch between tiles. Q, mask, top-K, and destination descriptors are offset to the tile while K/V remain shared. A Vulkan pipeline barrier separates reuse of each scratch tile. + +Scratch at PP2048 changed from: + +```text +Before: 538,968,064 bytes (514 MiB) +After: 67,371,008 bytes (64.25 MiB) +Change: 8x reduction +``` + +The exact production-shape microbenchmark (`kv=19200,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0`) measured: + +```text +Full-batch two partitions: 114.65 ms +256-query tiled path: 113.70 ms +``` + +Two one-partition alternatives were tested and discarded. Merging the raw partial directly inside the cooperative selected shader measured 119.17 ms. Writing selected output separately and using a lightweight merge kernel measured 120.45 ms. Both cut scratch in half but regressed because writing selected output outside the contiguous split layout increased the selected stage from about 44 ms to about 51 ms. Query tiling preserves the faster memory layout. + +Canonical 32k result after tiling: + +```text +32k context, PP 2048, ub 2048 + +Full-batch split: 208.70 tok/s, 9.77039 s total, 1.20170 s sparse FA +Tiled split: 209.62 tok/s, 9.72897 s total, 1.18721 s sparse FA + +Tiled split stages: +raw-prefix FA: 0.192879 s total, 168 tile dispatches +selected sparse FA: 0.910153 s total, 168 tile dispatches +split reduction: 0.084179 s total, 168 tile dispatches +Lightning Indexer: 1.100230 s total +TOP_K: 0.075158 s total +``` + +The 168 dispatch count is eight tiles x 21 sparse-attention layers. Relative to the full-batch split path, throughput improved 0.44%, total Vulkan time fell 0.42%, and sparse FA time fell 1.21%. The main result is the 8x scratch reduction without a performance regression. + +Final correctness after removing the discarded merge prototypes: + +- focused sparse top-K suite: 9/9 passed against the CPU reference, including a 257-query tile-boundary case +- complete Vulkan Flash Attention suite: 13,296/13,296 passed +- no NaN or Inf failure + +Logs: + +- `/tmp/dsv4-fa-exact-fused.log` +- `/tmp/dsv4-fa-exact-merge.log` +- `/tmp/dsv4-fa-exact-two-part-current.log` +- `/tmp/dsv4-fa-exact-tiled.log` +- `/tmp/dsv4-fa-tiled-final-correctness.log` +- `/tmp/dsv4-fa-tiled-boundary-correctness.log` +- `/tmp/dsv4-fa-tiled-all-correctness.log` +- `/tmp/dsv4-vulkan-tiled-32k.log` diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 141e7098496a..845180c706bb 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11168,6 +11168,8 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & const uint32_t NS = (uint32_t) q->ne[3]; const uint32_t raw_kv = (uint32_t) n_kv_raw; const uint32_t partitions = 2; + const uint32_t tile_size = std::min(N, 256u); + const uint32_t n_tiles = CEIL_DIV(N, tile_size); const bool f32acc = true; vk_fa_tuning_params tuning = get_fa_tuning_params(ctx->device, D, D, N, raw_kv, GGML_TYPE_F16, GGML_TYPE_F16, f32acc); @@ -11188,11 +11190,12 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & pipelines[raw_state] = raw_pipeline = std::make_shared(); } } - ggml_pipeline_request_descriptor_sets(ctx, raw_pipeline, 1); - ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); - ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1); + ggml_pipeline_request_descriptor_sets(ctx, raw_pipeline, n_tiles); + ggml_pipeline_request_descriptor_sets(ctx, pipeline, n_tiles); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, n_tiles); - const uint64_t split_size = ((uint64_t) D * NH * sizeof(float) + NH * 2 * sizeof(float)) * partitions * N * NS; + const uint64_t partition_size = ((uint64_t) D * NH + NH * 2) * sizeof(float) * tile_size * NS; + const uint64_t split_size = partition_size * partitions; if (split_size <= ctx->device->properties.limits.maxStorageBufferRange) { if (ctx->prealloc_size_split_k < split_size) { ctx->prealloc_size_split_k = split_size; @@ -11202,45 +11205,63 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & ggml_vk_sync_buffers(ctx, subctx); } - const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); const uint32_t n_head_log2 = 64; const uint32_t packed_gqa = (mask_stride << 16) | 1; const uint32_t packed_partitions = (partitions << 16) | 1; - const vk_flash_attn_push_constants raw_pc = { - N, raw_kv, - NH, N, NS, - NH, NS, - 1, NS, - 1, NS, - (uint32_t) mask->ne[1], (uint32_t) mask->ne[2], (uint32_t) mask->ne[3], - q_stride, (uint32_t) q->nb[2], (uint32_t) q->nb[3], - k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], - k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], - scale, 0.0f, 0.0f, - n_head_log2, 1.0f, 1.0f, - packed_gqa, raw_kv, packed_partitions, + const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + const vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, k); + const vk_subbuffer mask_buf = ggml_vk_tensor_subbuffer(ctx, mask); + const vk_subbuffer top_buf = ggml_vk_tensor_subbuffer(ctx, top_k); + const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + const auto sliced = [](const vk_subbuffer & buf, uint64_t offset) { + return vk_subbuffer{buf.buffer, buf.offset + offset, buf.size - offset}; }; - ggml_vk_dispatch_pipeline(ctx, subctx, raw_pipeline, - {q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, k), - ggml_vk_tensor_subbuffer(ctx, mask), q_buf, split_buf, q_buf}, - raw_pc, {N, NH, NS}); - ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_RAW (sub-op)"); - - pc.split_mode = 1; - ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf, - ggml_vk_tensor_subbuffer(ctx, top_k), split_buf}, - pc, {N, (uint32_t) CEIL_DIV(q->ne[2], 32), NS}); - ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_SELECTED (sub-op)"); - - ggml_vk_sync_buffers(ctx, subctx); - const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = {D, NH, N, NS, partitions, sinks != nullptr}; - ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, - {split_buf, sinks_buf, ggml_vk_tensor_subbuffer(ctx, dst)}, - reduce_pc, {NH, D, N * NS}); - ctx->prealloc_split_k_need_sync = true; - ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_REDUCE (sub-op)"); + for (uint32_t tile = 0; tile < n_tiles; ++tile) { + if (tile != 0) { + ggml_vk_sync_buffers(ctx, subctx); + } + const uint32_t token_offset = tile * tile_size; + const uint32_t tile_n = std::min(tile_size, N - token_offset); + const vk_subbuffer tile_q = sliced(q_buf, (uint64_t) token_offset * q->nb[1]); + const vk_subbuffer tile_mask = sliced(mask_buf, (uint64_t) token_offset * mask->nb[1]); + const vk_subbuffer tile_top = sliced(top_buf, (uint64_t) token_offset * top_k->nb[1]); + const vk_subbuffer tile_dst = sliced(dst_buf, (uint64_t) token_offset * dst->nb[2]); + const vk_flash_attn_push_constants raw_pc = { + tile_n, raw_kv, + NH, tile_n, NS, + NH, NS, + 1, NS, + 1, NS, + (uint32_t) mask->ne[1], (uint32_t) mask->ne[2], (uint32_t) mask->ne[3], + q_stride, (uint32_t) q->nb[2], (uint32_t) q->nb[3], + k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], + k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], + scale, 0.0f, 0.0f, + n_head_log2, 1.0f, 1.0f, + packed_gqa, raw_kv, packed_partitions, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, raw_pipeline, + {tile_q, k_buf, k_buf, tile_mask, tile_q, split_buf, tile_q}, + raw_pc, {tile_n, NH, NS}); + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_RAW (sub-op)"); + + pc.n_batch = tile_n; + pc.split_mode = 1; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {tile_q, k_buf, tile_mask, sinks_buf, tile_top, split_buf}, + pc, {tile_n, (uint32_t) CEIL_DIV(q->ne[2], 32), NS}); + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_SELECTED (sub-op)"); + + ctx->prealloc_split_k_need_sync = true; + ggml_vk_sync_buffers(ctx, subctx); + const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = {D, NH, tile_n, NS, partitions, sinks != nullptr}; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, + {split_buf, sinks_buf, tile_dst}, reduce_pc, {NH, D, tile_n * NS}); + ctx->prealloc_split_k_need_sync = true; + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_REDUCE (sub-op)"); + } return true; } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 2b6ae01e6eed..18ba88a3c335 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10290,6 +10290,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, true)); test_cases.emplace_back(new test_flash_attn_ext_top_k(1024, 64, 65, 128, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 128, 256, 512, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 257, 256, 512, false)); return test_cases; } From 8a4a20b3f2ed3450d9d1fa0a6488a1bd959ad85c Mon Sep 17 00:00:00 2001 From: Jaap Buurman Date: Thu, 13 Aug 2026 12:19:10 +0200 Subject: [PATCH 085/109] vulkan: reuse sparse FA probability fragments Assisted-by: Codex --- .../DSV4-vulkan-sparse-prefill-progress.md | 85 +++++++++++++++++++ .../vulkan-shaders/flash_attn_top_k_cm.comp | 27 +++--- 2 files changed, 100 insertions(+), 12 deletions(-) diff --git a/docs/development/DSV4-vulkan-sparse-prefill-progress.md b/docs/development/DSV4-vulkan-sparse-prefill-progress.md index 5ace979a08df..82843486bc0f 100644 --- a/docs/development/DSV4-vulkan-sparse-prefill-progress.md +++ b/docs/development/DSV4-vulkan-sparse-prefill-progress.md @@ -405,3 +405,88 @@ Logs: - `/tmp/dsv4-fa-tiled-boundary-correctness.log` - `/tmp/dsv4-fa-tiled-all-correctness.log` - `/tmp/dsv4-vulkan-tiled-32k.log` + +## Selected PV probability reuse + +The selected cooperative-matrix stage remained the largest split sparse-FA component. Profiling the exact tiled production shape showed: + +```text +selected K gather and cooperative QK: about 20.0 ms +QK plus softmax: about 22.3 ms +full selected stage: about 45.5 ms +cooperative PV and output increment: about 23.2 ms +``` + +PV and output were about 51% of the selected stage. The shader previously loaded each of four 16x16 probability cooperative-matrix fragments again for every one of the eight 64-dimension PV passes. The optimized shader loads these four fragments once per 64-key block and retains them across all PV dimension passes. This follows the probability-fragment lifetime used by ordinary cooperative Vulkan FA. + +The shader also aliases the shared score and PV-output matrices because their lifetimes do not overlap. Shader resource statistics changed as follows: + +```text + Before After +VGPRs 192 192 +VGPR spills 0 0 +LDS 36,864 28,672 bytes +static instructions 11,524 11,358 +``` + +The exact production-shape microbenchmark changed from 113.70 ms for the committed tiled path to 110.11 ms. The selected stage fell from about 45.5 ms to about 43.7 ms. The raw-prefix and reduction implementations are unchanged. + +Two canonical 32k runs after this change measured: + +```text +32k context, PP 2048, ub 2048 + +Committed tiled reference: +209.62 tok/s +Total Vulkan: 9.72897 s +Split sparse FA: 1.18721 s + raw-prefix FA: 0.192879 s + selected sparse FA: 0.910153 s + split reduction: 0.084179 s +Lightning Indexer: 1.100230 s +TOP_K: 0.075158 s + +Probability reuse, run 1: +211.03 tok/s +Total Vulkan: 9.66363 s +Split sparse FA: 1.12776 s + raw-prefix FA: 0.192078 s + selected sparse FA: 0.849862 s + split reduction: 0.085821 s +Lightning Indexer: 1.098980 s +TOP_K: 0.071036 s + +Probability reuse, run 2: +209.57 tok/s +Total Vulkan: 9.72968 s +Split sparse FA: 1.13553 s + raw-prefix FA: 0.192987 s + selected sparse FA: 0.859890 s + split reduction: 0.082653 s +Lightning Indexer: 1.102550 s +TOP_K: 0.071474 s +``` + +The two-run selected-stage improvement is 5.5-6.6%, and the complete split sparse-FA improvement is 4.4-5.0%. The two-run throughput mean is 210.30 tok/s, 0.32% above the 209.62 tok/s reference. End-to-end noise in unrelated model kernels is larger than this small total-throughput change, but both profiler runs isolate a consistent gain in the modified selected stage. + +Correctness after the change: + +- focused sparse top-K suite: 9/9 passed against the CPU reference +- complete Vulkan Flash Attention suite: 13,297/13,297 passed +- pipeline statistics probe: 1/1 passed, with no SGPR or VGPR spills +- no NaN or Inf failure + +Logs: + +- `/tmp/dsv4-fa-selected-qk-tiled.log` +- `/tmp/dsv4-fa-selected-softmax-tiled.log` +- `/tmp/dsv4-fa-lds-alias-stats.log` +- `/tmp/dsv4-fa-pmat-stats.log` +- `/tmp/dsv4-fa-exact-lds-alias.log` +- `/tmp/dsv4-fa-exact-pmat.log` +- `/tmp/dsv4-fa-pmat-correctness.log` +- `/tmp/dsv4-fa-pmat-all-correctness.log` +- `/tmp/dsv4-vulkan-32k-pmat.log` +- `/tmp/dsv4-vulkan-32k-pmat-repeat.log` + +The next sparse-FA optimization should continue to target the selected PV/output half. The retained probability fragments remove redundant cooperative loads without increasing reported VGPR allocation. More invasive changes such as doubling the PV dimension tile can reduce barriers but must be designed around the eight available wave64 subgroups, 64 KiB LDS limit, and already high 192-VGPR allocation. Do not build a crossover matrix until the next kernel layout is settled because an optimization can shift the crossover. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp index f3dfbe59ac00..d4bc4011d076 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp @@ -58,10 +58,9 @@ const float MASK_NEG_INF = -65500.0; shared uint key_idx[KEYS_PER_BLOCK]; shared f16vec4 q_sh[HEADS_PER_GROUP * QK_STRIDE]; shared f16vec4 k_sh[KEYS_PER_BLOCK * QK_STRIDE]; -shared vec4 score_sh[KEYS_PER_BLOCK * SCORE_STRIDE]; +shared vec4 matrix_sh[KEYS_PER_BLOCK * SCORE_STRIDE]; shared f16vec4 p_sh[HEADS_PER_GROUP * P_STRIDE]; shared f16vec4 v_sh[KEYS_PER_BLOCK * V_STRIDE]; -shared vec4 pv_sh[HEADS_PER_GROUP * PV_STRIDE]; shared float old_scale_sh[HEADS_PER_GROUP]; shared float row_max_sh[HEADS_PER_GROUP]; shared float row_sum_sh[HEADS_PER_GROUP]; @@ -142,7 +141,7 @@ void main() { const uint score_key_chunk = gl_SubgroupID % (KEYS_PER_BLOCK / TILE); const uint score_head_tile = gl_SubgroupID / (KEYS_PER_BLOCK / TILE); - coopMatStore(scores, score_sh, + coopMatStore(scores, matrix_sh, score_key_chunk * TILE * SCORE_STRIDE + score_head_tile * (TILE / 4), SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); barrier(); @@ -159,7 +158,7 @@ void main() { const uint key_local = softmax_lane + i * (SUBGROUP_SIZE / 4); const uint key = key_idx[key_local]; const float mask = key < p.n_kv ? float(data_m[mask_base + key]) : uintBitsToFloat(0xff800000); - const float score = float(score_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; + const float score = float(matrix_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; block_max = mask < MASK_NEG_INF ? block_max : max(block_max, score); } [[unroll]] for (uint delta = 1; delta < SUBGROUP_SIZE / 4; delta *= 2) { @@ -177,7 +176,7 @@ void main() { const float mask = key < p.n_kv ? float(data_m[mask_base + key]) : uintBitsToFloat(0xff800000); float weight = 0.0; if (mask >= MASK_NEG_INF) { - const float score = float(score_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; + const float score = float(matrix_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; weight = exp(score - new_max); block_sum += weight; } @@ -205,6 +204,14 @@ void main() { continue; } + coopmat pmats[KEYS_PER_BLOCK / TILE]; + [[unroll]] for (uint key_chunk = 0; key_chunk < KEYS_PER_BLOCK / TILE; ++key_chunk) { + const uint pv_head_tile = gl_SubgroupID / (DIMS_PER_BLOCK / TILE); + coopMatLoad(pmats[key_chunk], p_sh, + pv_head_tile * TILE * P_STRIDE + key_chunk * (TILE / 4), + P_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + } + [[unroll]] for (uint dim_base = 0; dim_base < HEAD_SIZE; dim_base += DIMS_PER_BLOCK) { [[unroll]] for (uint idx = tid; idx < KEYS_PER_BLOCK * (DIMS_PER_BLOCK / 4); idx += WORKGROUP_SIZE) { const uint key_local = idx / (DIMS_PER_BLOCK / 4); @@ -221,22 +228,18 @@ void main() { coopmat pv = coopmat(0.0); - coopmat pmat; coopmat vmat; const uint pv_head_tile = gl_SubgroupID / (DIMS_PER_BLOCK / TILE); const uint pv_dim_tile = gl_SubgroupID % (DIMS_PER_BLOCK / TILE); [[unroll]] for (uint key_chunk = 0; key_chunk < KEYS_PER_BLOCK / TILE; ++key_chunk) { - coopMatLoad(pmat, p_sh, - pv_head_tile * TILE * P_STRIDE + key_chunk * (TILE / 4), - P_STRIDE, gl_CooperativeMatrixLayoutRowMajor); coopMatLoad(vmat, v_sh, key_chunk * TILE * V_STRIDE + pv_dim_tile * (TILE / 4), V_STRIDE, gl_CooperativeMatrixLayoutRowMajor); - pv = coopMatMulAdd(pmat, vmat, pv); + pv = coopMatMulAdd(pmats[key_chunk], vmat, pv); } - coopMatStore(pv, pv_sh, + coopMatStore(pv, matrix_sh, pv_head_tile * TILE * PV_STRIDE + pv_dim_tile * (TILE / 4), PV_STRIDE, gl_CooperativeMatrixLayoutRowMajor); barrier(); @@ -246,7 +249,7 @@ void main() { const uint head_local = out_idx / HEAD_SIZE; const uint dim = out_idx % HEAD_SIZE; if (dim >= dim_base && dim < dim_base + DIMS_PER_BLOCK) { - accum[i] += pv_sh[head_local * PV_STRIDE + (dim - dim_base) / 4][dim % 4]; + accum[i] += matrix_sh[head_local * PV_STRIDE + (dim - dim_base) / 4][dim % 4]; } } barrier(); From 13bedfef6b969441e1e64633026c7f22b9ed0002 Mon Sep 17 00:00:00 2001 From: Jaap Buurman Date: Thu, 13 Aug 2026 12:59:50 +0200 Subject: [PATCH 086/109] vulkan: cache sparse FA masks per key block Assisted-by: Codex --- .../DSV4-vulkan-sparse-prefill-progress.md | 103 +++++++++++++++++- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 7 +- .../vulkan-shaders/flash_attn_base.glsl | 3 +- .../vulkan-shaders/flash_attn_top_k_cm.comp | 30 +++-- tests/test-backend-ops.cpp | 5 +- 5 files changed, 123 insertions(+), 25 deletions(-) diff --git a/docs/development/DSV4-vulkan-sparse-prefill-progress.md b/docs/development/DSV4-vulkan-sparse-prefill-progress.md index 82843486bc0f..e880ae63e99e 100644 --- a/docs/development/DSV4-vulkan-sparse-prefill-progress.md +++ b/docs/development/DSV4-vulkan-sparse-prefill-progress.md @@ -262,7 +262,7 @@ Always run these sequentially. Do not run a compiler concurrently on this APU. An uncommitted follow-up prototype was tested after commit `24c7ead76cc1a9631fa1b42b0bfa53a15169e1ba`. Check `git status` before continuing. It was initially tested with `GGML_VK_FA_TOPK_SPLIT=1`. After the successful 32k run, the split path was changed to default-on for a llama-server coherence test. Set `GGML_VK_FA_TOPK_SPLIT=0` to restore the single cooperative sparse kernel. -Stage profiling on the exact production sparse shape (`kv=19200,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0`) showed: +Stage profiling on the previously used 19,200-row sparse shape (`kv=19200,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0`) showed: ```text cooperative QK and selected-K gather: 88.885 ms @@ -362,7 +362,7 @@ After: 67,371,008 bytes (64.25 MiB) Change: 8x reduction ``` -The exact production-shape microbenchmark (`kv=19200,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0`) measured: +The previously used 19,200-row microbenchmark (`kv=19200,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0`) measured: ```text Full-batch two partitions: 114.65 ms @@ -408,7 +408,7 @@ Logs: ## Selected PV probability reuse -The selected cooperative-matrix stage remained the largest split sparse-FA component. Profiling the exact tiled production shape showed: +The selected cooperative-matrix stage remained the largest split sparse-FA component. Profiling the tiled 19,200-row shape showed: ```text selected K gather and cooperative QK: about 20.0 ms @@ -429,7 +429,7 @@ LDS 36,864 28,672 bytes static instructions 11,524 11,358 ``` -The exact production-shape microbenchmark changed from 113.70 ms for the committed tiled path to 110.11 ms. The selected stage fell from about 45.5 ms to about 43.7 ms. The raw-prefix and reduction implementations are unchanged. +The 19,200-row microbenchmark changed from 113.70 ms for the committed tiled path to 110.11 ms. The selected stage fell from about 45.5 ms to about 43.7 ms. The raw-prefix and reduction implementations are unchanged. Two canonical 32k runs after this change measured: @@ -490,3 +490,98 @@ Logs: - `/tmp/dsv4-vulkan-32k-pmat-repeat.log` The next sparse-FA optimization should continue to target the selected PV/output half. The retained probability fragments remove redundant cooperative loads without increasing reported VGPR allocation. More invasive changes such as doubling the PV dimension tile can reduce barriers but must be designed around the eight available wave64 subgroups, 64 KiB LDS limit, and already high 192-VGPR allocation. Do not build a crossover matrix until the next kernel layout is settled because an optimization can shift the crossover. + +## Selected mask caching and deep-context micro matrix + +The sparse FA `kv` dimension is compressed K/V rows, not source-token context depth. For the tested DeepSeek V4 graph, a PP2048 batch uses 2,304 raw rows and approximately one compressed row per four source tokens. The useful synthetic mapping is: + +```text +Source context Sparse FA kv rows +32k 11,008 +64k 19,200 +128k 35,584 +256k 68,352 +512k 133,888 +``` + +The performance test registry now contains PP2048 cases at all five K extents with `n_kv_raw=2304` and `n_top_k=512`. Use this command template and replace `KV` with a value from the table: + +```bash +GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=KV,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0' +``` + +These are synthetic sparse-FA depth tests. They do not include the Lightning Indexer or the rest of the model graph and do not replace the canonical 32k llama-bench. + +The selected shader previously loaded the same selected-key mask value independently for all 32 heads during both softmax passes. The shader now loads each mask value once per 64-key block and reuses it from LDS. Q and probability storage also share one LDS allocation, while K and V staging share another because both pairs have disjoint lifetimes. + +Shader resources changed from the probability-reuse commit: + +```text + Before After +VGPRs 192 192 +VGPR spills 0 0 +LDS 28,672 24,576 bytes +static instructions 11,358 11,233 +``` + +A 128-dimension V staging experiment was correct but rejected. With operand aliasing it used exactly 32 KiB LDS. It was 2.3% slower at simulated 32k, approximately equal at 64k, and 1.6% faster at 128k. The retained 64-dimension layout is better for the canonical 32k target and is simpler. + +Selected-mask caching improved every synthetic depth relative to the same 64-dimension operand-alias layout: + +```text +Depth Selected before Selected after Split total before Split total after +32k 43.16 ms 41.23 ms 109.45 ms 108.07 ms +64k 43.22 ms 41.42 ms 110.45 ms 109.11 ms +128k 52.26 ms 48.74 ms 119.15 ms 116.92 ms +256k 51.64 ms 49.02 ms 118.06 ms 116.93 ms +512k 51.33 ms 49.20 ms 117.87 ms 116.27 ms +``` + +The 256k case exposed a separate path-selection cutoff. The raw-prefix ordinary FA dispatch encoded its mask row stride in 16 bits and disabled split sparse FA above 65,535 K rows. It now uses a flagged convention that carries the full 32-bit mask stride in the existing `split_kv` push constant for this one-partition partial-output dispatch. No push-constant structure was enlarged. At simulated 256k, this changes the selected path from unsplit cooperative sparse FA at 206.22 ms to split sparse FA at 118.06 ms before mask caching, a 42.8% reduction. The 512k case also uses the split path successfully. + +Canonical 32k PP2048 after operand aliasing and selected-mask caching: + +```text +Previous commit, two runs: +209.57-211.03 tok/s +Total Vulkan: 9.664-9.730 s +Split sparse FA: 1.128-1.136 s + raw-prefix FA: 0.192-0.193 s + selected sparse FA: 0.850-0.860 s + split reduction: 0.083-0.086 s + +Current result: +210.47 tok/s +Total Vulkan: 9.68921 s +Split sparse FA: 1.10160 s + raw-prefix FA: 0.192860 s + selected sparse FA: 0.821381 s + split reduction: 0.087354 s +Lightning Indexer: 1.102230 s +TOP_K: 0.072488 s +``` + +The selected stage improves 3.4-4.5% and complete split sparse FA improves 2.3-3.0% relative to the previous commit's two canonical runs. End-to-end throughput remains within model-wide benchmark noise. + +Correctness: + +- focused sparse top-K suite: 9/9 passed +- complete Vulkan Flash Attention suite: 13,297/13,297 passed +- no NaN or Inf failure +- simulated 256k and 512k performance cases both selected the split path and completed successfully + +Benchmark policy for this laptop APU: do not repeat llama-bench when the exact-shape microbenchmarks and the first canonical profiler block agree. Sustained load can lower GPU clocks and bias a repeat. Repeat only when the first result is anomalous or contradicts the microbenchmarks. Never build and benchmark concurrently. + +Logs: + +- `/tmp/dsv4-fa-mask-cache-32k.log` +- `/tmp/dsv4-fa-mask-cache-64k.log` +- `/tmp/dsv4-fa-mask-cache-128k.log` +- `/tmp/dsv4-fa-mask-cache-256k.log` +- `/tmp/dsv4-fa-mask-cache-512k.log` +- `/tmp/dsv4-fa-final-micro-32k.log` +- `/tmp/dsv4-fa-final-focused-correctness.log` +- `/tmp/dsv4-fa-mask-cache-all-correctness.log` +- `/tmp/dsv4-vulkan-32k-mask-cache.log` diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 845180c706bb..60ece80b50a6 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11160,7 +11160,7 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & static const char * top_k_split_env = getenv("GGML_VK_FA_TOPK_SPLIT"); const uint32_t mask_stride = (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)); - const bool try_split = use_cm && (!top_k_split_env || top_k_split_env[0] != '0') && n_kv_raw > 0 && top_k->ne[0] > 0 && mask_stride <= 0xffff; + const bool try_split = use_cm && (!top_k_split_env || top_k_split_env[0] != '0') && n_kv_raw > 0 && top_k->ne[0] > 0; if (try_split) { const uint32_t N = (uint32_t) q->ne[1]; const uint32_t D = 512; @@ -11206,7 +11206,8 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & } const uint32_t n_head_log2 = 64; - const uint32_t packed_gqa = (mask_stride << 16) | 1; + const uint32_t mask_stride_in_split_kv = 1u << 31; + const uint32_t packed_gqa = mask_stride_in_split_kv | 1u; const uint32_t packed_partitions = (partitions << 16) | 1; const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); const vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, k); @@ -11239,7 +11240,7 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], scale, 0.0f, 0.0f, n_head_log2, 1.0f, 1.0f, - packed_gqa, raw_kv, packed_partitions, + packed_gqa, mask_stride, packed_partitions, }; ggml_vk_dispatch_pipeline(ctx, subctx, raw_pipeline, diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 800a79a97d17..8fba7a450f7e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -211,8 +211,9 @@ void init_indices() q_stride = gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01; k_stride = p.nb11; v_stride = p.nb21; + const bool mask_stride_in_split_kv = (p.gqa_ratio & 0x80000000u) != 0; const uint32_t mask_stride_override = p.gqa_ratio >> 16; - m_stride = mask_stride_override != 0 ? mask_stride_override : (gqa_ratio > 1 ? 0 : KV); + m_stride = mask_stride_in_split_kv ? p.split_kv : (mask_stride_override != 0 ? mask_stride_override : (gqa_ratio > 1 ? 0 : KV)); } // Bias applied to softmax to stay in fp16 range. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp index d4bc4011d076..3d0024553e7a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp @@ -56,11 +56,10 @@ const uint PV_STRIDE = DIMS_PER_BLOCK / 4; const float MASK_NEG_INF = -65500.0; shared uint key_idx[KEYS_PER_BLOCK]; -shared f16vec4 q_sh[HEADS_PER_GROUP * QK_STRIDE]; -shared f16vec4 k_sh[KEYS_PER_BLOCK * QK_STRIDE]; +shared float key_mask[KEYS_PER_BLOCK]; +shared f16vec4 q_p_sh[HEADS_PER_GROUP * P_STRIDE]; +shared f16vec4 k_v_sh[KEYS_PER_BLOCK * V_STRIDE]; shared vec4 matrix_sh[KEYS_PER_BLOCK * SCORE_STRIDE]; -shared f16vec4 p_sh[HEADS_PER_GROUP * P_STRIDE]; -shared f16vec4 v_sh[KEYS_PER_BLOCK * V_STRIDE]; shared float old_scale_sh[HEADS_PER_GROUP]; shared float row_max_sh[HEADS_PER_GROUP]; shared float row_sum_sh[HEADS_PER_GROUP]; @@ -99,6 +98,7 @@ void main() { } } key_idx[tid] = key; + key_mask[tid] = key < p.n_kv ? float(data_m[mask_base + key]) : uintBitsToFloat(0xff800000); } barrier(); @@ -117,23 +117,23 @@ void main() { const uint offset = stream * p.nbk3 + key * p.nbk1 + d + d4 * 4; value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); } - k_sh[key_local * QK_STRIDE + d4] = value; + k_v_sh[key_local * QK_STRIDE + d4] = value; } if (tid < HEADS_PER_GROUP * (TILE / 4)) { const uint head_local = tid / (TILE / 4); const uint d4 = tid % (TILE / 4); const uint head = head_base + head_local; const uint offset = stream * p.nbq3 + head * p.nbq2 + token * p.nbq1 + d + d4 * 4; - q_sh[head_local * QK_STRIDE + d4] = f16vec4( + q_p_sh[head_local * QK_STRIDE + d4] = f16vec4( data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); } barrier(); const uint key_chunk = gl_SubgroupID % (KEYS_PER_BLOCK / TILE); const uint head_tile = gl_SubgroupID / (KEYS_PER_BLOCK / TILE); - coopMatLoad(kmat, k_sh, key_chunk * TILE * QK_STRIDE, + coopMatLoad(kmat, k_v_sh, key_chunk * TILE * QK_STRIDE, QK_STRIDE, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(qmat, q_sh, head_tile * TILE * QK_STRIDE, + coopMatLoad(qmat, q_p_sh, head_tile * TILE * QK_STRIDE, QK_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); scores = coopMatMulAdd(kmat, qmat, scores); barrier(); @@ -156,8 +156,7 @@ void main() { float block_max = uintBitsToFloat(0xff800000); [[unroll]] for (uint i = 0; i < KEYS_PER_BLOCK / (SUBGROUP_SIZE / 4); ++i) { const uint key_local = softmax_lane + i * (SUBGROUP_SIZE / 4); - const uint key = key_idx[key_local]; - const float mask = key < p.n_kv ? float(data_m[mask_base + key]) : uintBitsToFloat(0xff800000); + const float mask = key_mask[key_local]; const float score = float(matrix_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; block_max = mask < MASK_NEG_INF ? block_max : max(block_max, score); } @@ -172,15 +171,14 @@ void main() { float block_sum = 0.0; [[unroll]] for (uint i = 0; i < KEYS_PER_BLOCK / (SUBGROUP_SIZE / 4); ++i) { const uint key_local = softmax_lane + i * (SUBGROUP_SIZE / 4); - const uint key = key_idx[key_local]; - const float mask = key < p.n_kv ? float(data_m[mask_base + key]) : uintBitsToFloat(0xff800000); + const float mask = key_mask[key_local]; float weight = 0.0; if (mask >= MASK_NEG_INF) { const float score = float(matrix_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; weight = exp(score - new_max); block_sum += weight; } - p_sh[head_local * P_STRIDE + key_local / 4][key_local % 4] = float16_t(weight); + q_p_sh[head_local * P_STRIDE + key_local / 4][key_local % 4] = float16_t(weight); } [[unroll]] for (uint delta = 1; delta < SUBGROUP_SIZE / 4; delta *= 2) { block_sum += subgroupShuffleXor(block_sum, delta); @@ -207,7 +205,7 @@ void main() { coopmat pmats[KEYS_PER_BLOCK / TILE]; [[unroll]] for (uint key_chunk = 0; key_chunk < KEYS_PER_BLOCK / TILE; ++key_chunk) { const uint pv_head_tile = gl_SubgroupID / (DIMS_PER_BLOCK / TILE); - coopMatLoad(pmats[key_chunk], p_sh, + coopMatLoad(pmats[key_chunk], q_p_sh, pv_head_tile * TILE * P_STRIDE + key_chunk * (TILE / 4), P_STRIDE, gl_CooperativeMatrixLayoutRowMajor); } @@ -222,7 +220,7 @@ void main() { const uint offset = stream * p.nbk3 + key * p.nbk1 + dim_base + d4 * 4; value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); } - v_sh[key_local * V_STRIDE + d4] = value; + k_v_sh[key_local * V_STRIDE + d4] = value; } barrier(); @@ -233,7 +231,7 @@ void main() { const uint pv_head_tile = gl_SubgroupID / (DIMS_PER_BLOCK / TILE); const uint pv_dim_tile = gl_SubgroupID % (DIMS_PER_BLOCK / TILE); [[unroll]] for (uint key_chunk = 0; key_chunk < KEYS_PER_BLOCK / TILE; ++key_chunk) { - coopMatLoad(vmat, v_sh, + coopMatLoad(vmat, k_v_sh, key_chunk * TILE * V_STRIDE + pv_dim_tile * (TILE / 4), V_STRIDE, gl_CooperativeMatrixLayoutRowMajor); pv = coopMatMulAdd(pmats[key_chunk], vmat, pv); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 18ba88a3c335..66421146149d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10717,7 +10717,10 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, nb, 1024, 512, false)); } } - test_cases.emplace_back(new test_flash_attn_ext_top_k(19200, 2048, 2304, 512, false)); + // PP2048 compressed-K rows for source context depths 32k through 512k. + for (int kv : { 11008, 19200, 35584, 68352, 133888 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, 2048, 2304, 512, false)); + } return test_cases; } From 2c7937646a61a2ba5bbfe975dff4a3e6477a5aab Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 13 Aug 2026 14:32:06 +0000 Subject: [PATCH 087/109] vulkan: fix DeepSeek V4 sparse split attention with multiple sequences The raw/selected split path mis-indexed three things once q->ne[3] > 1. All three are inert at a single sequence, so the existing tests never reached them. - flash_attn_top_k_cm.comp wrote the L/M rows at a base that omitted the stream count. The split buffer is [O matrices][L/M rows] with both regions spanning ne3, so with more than one sequence the selected partition's L/M landed inside the O region and corrupted partition 0. - The mask's per-stream offset stepped by nem1 * KV. That assumes the mask row length equals KV, which the split path breaks: it overrides the mask stride to the full K range while KV covers only the raw prefix. Separate the two quantities - m_stride is the row-to-row step inside the tile (0 under GQA), m_row_len is the mask's real row length used to step between tokens and streams. They coincide everywhere except GQA and this path. - Query tiling conflicts with multiple streams: flash_attn_split_k_reduce takes the tile height as ne2 and derives the destination row from it, so with several tiles it writes stream s at s*tile_size instead of s*N. Tile only when there is one stream, which is the prefill case the tiling was added for. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 7 ++++++- .../vulkan-shaders/flash_attn_base.glsl | 18 +++++++++++++++--- .../vulkan-shaders/flash_attn_cm1.comp | 4 ++-- .../vulkan-shaders/flash_attn_top_k_cm.comp | 9 ++++++++- 4 files changed, 31 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 60ece80b50a6..91eda5481c63 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11168,7 +11168,12 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & const uint32_t NS = (uint32_t) q->ne[3]; const uint32_t raw_kv = (uint32_t) n_kv_raw; const uint32_t partitions = 2; - const uint32_t tile_size = std::min(N, 256u); + // Query tiling keeps the split scratch small, but flash_attn_split_k_reduce derives the + // destination row from ne2, which it is handed as the TILE height. With one tile that + // equals N and the stream stride is right; with several tiles and more than one stream + // it would write stream s at s*tile_size instead of s*N. Only tile when there is a + // single stream, which is the prefill case the tiling exists for. + const uint32_t tile_size = (NS == 1) ? std::min(N, 256u) : N; const uint32_t n_tiles = CEIL_DIV(N, tile_size); const bool f32acc = true; vk_fa_tuning_params tuning = get_fa_tuning_params(ctx->device, D, D, N, raw_kv, GGML_TYPE_F16, GGML_TYPE_F16, f32acc); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 8fba7a450f7e..e308c7214dd3 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -144,7 +144,7 @@ ACC_TYPE perElemOpGetSink(const in uint32_t r, const in uint32_t c, const in ACC uint32_t i, N, KV, split_k_index, Tr, start_j, end_j, gqa_iq1, iq2, iq3, rk2, rk3, rv2, rv3, ik2, ik3, iv2, iv3, - q_stride, k_stride, v_stride, m_stride, gqa_ratio, split_k_num, output_k_num; + q_stride, k_stride, v_stride, m_stride, m_row_len, gqa_ratio, split_k_num, output_k_num; bool partial_output; void init_indices() @@ -211,9 +211,21 @@ void init_indices() q_stride = gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01; k_stride = p.nb11; v_stride = p.nb21; + // Bit 31 of gqa_ratio means "the mask row stride is in split_kv", used by the DeepSeek V4 + // sparse split path where the mask spans the full K range but this dispatch only covers the + // raw prefix, so m_stride != KV. That path always sets split_k_num == 1, which is what keeps + // split_kv free to carry the stride; ggml_vk_flash_attn_top_k asserts the invariant. + // Otherwise: when using grouped query attention all rows share the same mask (stride 0). + // "p.gqa_ratio >> 16" is just a roundabout way of writing zero that prevents the compiler + // from folding the "&" through the select and breaking the alignment detection. const bool mask_stride_in_split_kv = (p.gqa_ratio & 0x80000000u) != 0; - const uint32_t mask_stride_override = p.gqa_ratio >> 16; - m_stride = mask_stride_in_split_kv ? p.split_kv : (mask_stride_override != 0 ? mask_stride_override : (gqa_ratio > 1 ? 0 : KV)); + m_stride = mask_stride_in_split_kv ? p.split_kv : ((gqa_ratio > 1) ? (p.gqa_ratio >> 16) : KV); + // Distinct from m_stride: m_stride is the row-to-row step INSIDE this tile (0 under GQA, + // where every row shares one mask row), while m_row_len is the mask tensor's actual row + // length, used to step between tokens and between streams. They differ under GQA and + // under the sparse split path, where this dispatch only covers the raw prefix (KV) but + // the mask rows span the whole K range. + m_row_len = mask_stride_in_split_kv ? p.split_kv : KV; } // Bias applied to softmax to stay in fp16 range. 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 e3ea909a59ce..38ae42c81f0c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -139,9 +139,9 @@ void main() { // FaBlockBytesK/V == 2 for f16 (sizeof f16) and == 16 for f32 (vec4) and == ggml block size for quants. uint32_t k_offset = (ik2*p.nb12 + ik3*p.nb13) / FaBlockBytesK; uint32_t v_offset = (iv2*p.nb22 + iv3*p.nb23) / FaBlockBytesV; - uint32_t m_offset = gqa_iq1*KV; + uint32_t m_offset = gqa_iq1*m_row_len; if (p.nem2 != 1 || p.nem3 != 1) { - m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV; + m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * m_row_len; mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp index 3d0024553e7a..e0628c65705b 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp @@ -276,8 +276,15 @@ void main() { const uint dim = out_idx % HEAD_SIZE; if (p.split_mode != 0) { const uint part_idx = 1; + // Split-K buffer layout, shared with flash_attn_cm1.comp and + // flash_attn_split_k_reduce.comp: all O matrices [HSV, ne1, k_num, ne2, ne3] + // first, then the L/M rows [ne1, k_num, ne2, ne3]. Here ne1 = n_head, + // ne2 = n_batch and ne3 = the stream count, so the L/M base must span every + // stream -- omitting gl_NumWorkGroups.z lands L/M inside the O region as soon + // as there is more than one sequence. + const uint n_streams = gl_NumWorkGroups.z; const uint matrix_base = HEAD_SIZE * p.n_head * (part_idx + 2 * (token + p.n_batch * stream)); - const uint lm_base = HEAD_SIZE * p.n_head * p.n_batch * 2 + p.n_head * 2 * (part_idx + 2 * (token + p.n_batch * stream)); + const uint lm_base = HEAD_SIZE * p.n_head * p.n_batch * n_streams * 2 + p.n_head * 2 * (part_idx + 2 * (token + p.n_batch * stream)); data_dst[matrix_base + (head_base + head_local) * HEAD_SIZE + dim] = accum[i]; if (dim == 0) { data_dst[lm_base + head_base + head_local] = row_sum_sh[head_local]; From 9c950c5441f7e417752473438a980dedad24d228 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 13 Aug 2026 14:32:23 +0000 Subject: [PATCH 088/109] vulkan: harden the sparse FA split path and drop its debug scaffolding No behaviour change on any path that runs today; this removes footguns the split path left in shared code and clears the leftovers. - flash_attn.comp and flash_attn_cm2.comp still read p.k_num and p.gqa_ratio raw, while the split path packs a partition count into the high half of k_num and a flag into bit 31 of gqa_ratio. Only the coopmat1 gate keeps them from ever seeing those values, and nothing said so. Use the decoded globals init_indices() already computes, so routing the split path at a sibling shader fails loudly instead of writing at wild offsets. - Drop profile_stage and GGML_VK_FA_TOPK_PROFILE. This also removes two continues from the sparse kernel's key loop. - Remove the unreachable mask-stride override branch (nothing sets bits 16..30 without bit 31) and restore the comment explaining why the GQA case writes zero the roundabout way: the compiler must not fold it, or stride alignment detection breaks. - Assert that the mask stride covers the raw KV range. The raw dispatch smuggles that stride through split_kv, which the shader also uses to derive its KV range, so a stride below KV would silently clip it. - Request descriptor sets only after the split scratch is known to fit, so the fallback path does not inherit n_tiles of unused requests. - n_head_log2 is unread with max_bias == 0; pass 0 rather than a number that looks computed. - Make the crossover depend on the path that will run. The coopmat sparse path wins as soon as any key is pruned, but the scalar fallback does not and regresses against dense FA until the pruned fraction is large, so it keeps its original 3x margin. - Record that the coopmat sparse shader requires 64-wide subgroups and eight of them, at the constants that encode it. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 30 ++++++++++++------- .../vulkan-shaders/flash_attn.comp | 20 ++++++------- .../vulkan-shaders/flash_attn_cm2.comp | 24 +++++++-------- .../vulkan-shaders/flash_attn_top_k.comp | 1 - .../vulkan-shaders/flash_attn_top_k_cm.comp | 15 ++++------ 5 files changed, 48 insertions(+), 42 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 91eda5481c63..c5451751e16c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1970,7 +1970,6 @@ struct vk_op_flash_attn_top_k_push_constants { uint32_t nb1, nb2, nb3; float scale; uint32_t has_sinks; - uint32_t profile_stage; uint32_t split_mode; }; static_assert(sizeof(vk_op_flash_attn_top_k_push_constants) <= 128); @@ -11128,7 +11127,14 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & return false; } const int64_t n_kv_active = n_kv_raw + top_k->ne[0]; - if (k->ne[1] <= n_kv_active) { + // Crossover against ordinary dense FA. The coopmat sparse path (and especially the + // raw/selected split below) is cheap enough to win as soon as any key is pruned; the + // scalar fallback shader is not, and regresses against dense FA until the pruned + // fraction is large, so it keeps its original 3x margin. Equality always uses dense + // FA because nothing is pruned. + const bool have_cm_sparse = ctx->device->pipeline_flash_attn_top_k_cm_f16 != nullptr; + const int64_t min_total_k = have_cm_sparse ? n_kv_active + 1 : 3 * n_kv_active; + if (k->ne[1] < min_total_k) { return false; } @@ -11147,15 +11153,13 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & (uint32_t) (dst->nb[1] / sizeof(float)), (uint32_t) (dst->nb[2] / sizeof(float)), (uint32_t) (dst->nb[3] / sizeof(float)), - scale, sinks != nullptr, 0, 0, + scale, sinks != nullptr, 0, }; const vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q); const vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; static const char * top_k_cm_env = getenv("GGML_VK_FA_TOPK_CM"); const bool use_cm = (!top_k_cm_env || top_k_cm_env[0] != '0') && ctx->device->pipeline_flash_attn_top_k_cm_f16; - static const char * top_k_profile_env = getenv("GGML_VK_FA_TOPK_PROFILE"); - pc.profile_stage = use_cm && top_k_profile_env ? atoi(top_k_profile_env) : 0; vk_pipeline pipeline = use_cm ? ctx->device->pipeline_flash_attn_top_k_cm_f16 : ctx->device->pipeline_flash_attn_top_k_f16; static const char * top_k_split_env = getenv("GGML_VK_FA_TOPK_SPLIT"); @@ -11195,13 +11199,12 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & pipelines[raw_state] = raw_pipeline = std::make_shared(); } } - ggml_pipeline_request_descriptor_sets(ctx, raw_pipeline, n_tiles); - ggml_pipeline_request_descriptor_sets(ctx, pipeline, n_tiles); - ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, n_tiles); - const uint64_t partition_size = ((uint64_t) D * NH + NH * 2) * sizeof(float) * tile_size * NS; const uint64_t split_size = partition_size * partitions; if (split_size <= ctx->device->properties.limits.maxStorageBufferRange) { + ggml_pipeline_request_descriptor_sets(ctx, raw_pipeline, n_tiles); + ggml_pipeline_request_descriptor_sets(ctx, pipeline, n_tiles); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, n_tiles); if (ctx->prealloc_size_split_k < split_size) { ctx->prealloc_size_split_k = split_size; ggml_vk_preallocate_buffers(ctx, subctx); @@ -11210,7 +11213,14 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & ggml_vk_sync_buffers(ctx, subctx); } - const uint32_t n_head_log2 = 64; + // ALiBi is disabled (max_bias == 0), so n_head_log2 is never read; keep it 0 + // rather than a value that looks computed. + const uint32_t n_head_log2 = 0; + // The raw dispatch smuggles the mask row stride through split_kv, which the FA + // shader also uses to derive its KV range as min(KV, (split_k_index+1)*split_kv). + // That is only safe while split_k_index == 0 and the stride covers the whole raw + // prefix -- both hold here, but assert rather than rely on it silently. + GGML_ASSERT(mask_stride >= raw_kv && "split_kv carries the mask stride; it must not clip the raw KV range"); const uint32_t mask_stride_in_split_kv = 1u << 31; const uint32_t packed_gqa = mask_stride_in_split_kv | 1u; const uint32_t packed_partitions = (partitions << 16) | 1; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp index 18a37add9bf8..5b19ff61093f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp @@ -186,9 +186,9 @@ void main() { // FaBlockBytesK/V == 2 for f16, 16 for f32, ggml block byte size for quants. uint32_t k_offset = (ik2*p.nb12 + ik3*p.nb13) / FaBlockBytesK; uint32_t v_offset = (iv2*p.nb22 + iv3*p.nb23) / FaBlockBytesV; - uint32_t m_offset = gqa_iq1*KV; + uint32_t m_offset = gqa_iq1*m_row_len; if (p.nem2 != 1 || p.nem3 != 1) { - m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV; + m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * m_row_len; mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; } @@ -210,7 +210,7 @@ void main() { } // Only load if the block is not all zeros if (mask_opt_bits != MASK_OPT_ALL_ZERO) { - bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; + bool nem1_bounds_check = !(gqa_ratio > 1) && (p.nem1 % Br) != 0; float max_mask = NEG_FLT_MAX_OVER_2; barrier(); @@ -660,10 +660,10 @@ void main() { // If there is split_k, then the split_k resolve shader does the final // division by L. Store the intermediate O value and per-row m and L values. - if (p.k_num > 1) { - if (p.gqa_ratio > 1) { + if (partial_output || split_k_num > 1) { + if (gqa_ratio > 1) { // note: O and Q have swapped coord 1,2. - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); @@ -674,7 +674,7 @@ void main() { } } - o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); + o_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)); [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); if (row < N) { @@ -688,7 +688,7 @@ void main() { const uint global_row = i * Br + row; if (global_row < N) { - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)) / 4; + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)) / 4; [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { data_ov4[o_offset + iq2 * HSV/4 + d * D_split + d_tid] = D_TYPEV4(Of[r][d]); @@ -696,7 +696,7 @@ void main() { } if (global_row < N && d_tid == 0 && col_tid == 0) { - uint32_t lm_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); + uint32_t lm_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)); data_o[lm_offset + iq2] = D_TYPE(Lf[r]); data_o[lm_offset + p.ne1 + iq2] = D_TYPE(Mf[r]); } @@ -742,7 +742,7 @@ void main() { uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4; - if (p.gqa_ratio > 1) { + if (gqa_ratio > 1) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); if (row < N) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp index 317411153087..54be1e6daa99 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -151,7 +151,7 @@ D_TYPE perElemOpNonGqaSplitKStore(const in uint32_t r, const in uint32_t c, cons uint32_t global_row = i * Br + r; if (global_row < N && c < HSV) { uint32_t o_off = HSV * p.ne1 - * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); + * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)); data_o[o_off + iq2 * HSV + c] = D_TYPE(elem); } return elem; @@ -161,8 +161,8 @@ D_TYPE perElemOpNonGqaSplitKStore(const in uint32_t r, const in uint32_t c, cons ACC_TYPE perElemOpNonGqaSplitKStoreCol0(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t lm_base, const in uint32_t iq2, const in uint32_t N) { uint32_t global_row = i * Br + r; if (global_row < N && c == 0) { - uint32_t lm_off = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 - + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); + uint32_t lm_off = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + + p.ne1 * 2 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)); data_o[lm_off + lm_base + iq2] = D_TYPE(elem); } return elem; @@ -242,9 +242,9 @@ void main() { // mo_offset will point to the tile starting at row i*Br and col 0 uint32_t mo_offset = mo_stride * i; - uint32_t m_offset = gqa_iq1*KV * 2 /*sizeof(float16_t)*/; + uint32_t m_offset = gqa_iq1*m_row_len * 2 /*sizeof(float16_t)*/; if (p.nem2 != 1 || p.nem3 != 1) { - m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV * 2 /*sizeof(float16_t)*/; + m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * m_row_len * 2 /*sizeof(float16_t)*/; mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; } @@ -268,7 +268,7 @@ void main() { } // Only load if the block is not all zeros if (mask_opt_bits != MASK_OPT_ALL_ZERO) { - bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; + bool nem1_bounds_check = !(gqa_ratio > 1) && (p.nem1 % Br) != 0; if (nem1_bounds_check) { tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutM = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); @@ -287,7 +287,7 @@ void main() { } else { tensorLayoutNV<2, Clamp> tensorLayoutM = createTensorLayoutNV(2, Clamp); // Don't clamp against nem1 when GQA is enabled - uint32_t m_height = p.gqa_ratio > 1 ? ~0 : p.nem1; + uint32_t m_height = gqa_ratio > 1 ? ~0 : p.nem1; tensorLayoutM = setTensorLayoutDimensionNV(tensorLayoutM, m_height, KV); tensorLayoutM = setTensorLayoutStrideNV(tensorLayoutM, m_stride, 1); @@ -406,15 +406,15 @@ void main() { // If there is split_k, then the split_k resolve shader does the final // division by L. Store the intermediate O value and per-row m and L values. - if (p.k_num > 1) { + if (partial_output || split_k_num > 1) { coopmat O_D = coopmat(O); - if (p.gqa_ratio > 1) { + if (gqa_ratio > 1) { // note: O and Q have swapped coord 1,2. - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)); coopMatPerElementNV(O_D, O_D, perElemOpGqaStore, o_offset, iq2, N); - o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); + o_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)); coopMatPerElementNV(L, L, perElemOpStoreCol0, o_offset, iq2, N); coopMatPerElementNV(M, M, perElemOpStoreCol0, o_offset + p.ne1, iq2, N); } else { @@ -473,7 +473,7 @@ void main() { uint32_t o_offset = gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV; - if (p.gqa_ratio > 1) { + if (gqa_ratio > 1) { coopMatPerElementNV(O_D, O_D, perElemOpGqaStore, o_offset, iq2, N); } else { tensorLayoutNV<3, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(3, gl_CooperativeMatrixClampModeConstantNV); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp index 62807471bdd1..930777b5f0ea 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp @@ -37,7 +37,6 @@ layout(push_constant) uniform Parameters { uint nb3; float scale; uint has_sinks; - uint profile_stage; uint split_mode; } p; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp index e0628c65705b..7a328938777a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp @@ -39,10 +39,15 @@ layout(push_constant) uniform Parameters { uint nb3; float scale; uint has_sinks; - uint profile_stage; uint split_mode; } p; +// This shader is hard-wired to a 512-thread workgroup of eight 64-wide subgroups: +// the softmax segments SUBGROUP_SIZE/4 = 16 lanes per head (so head_local stays < 32), +// and the QK/PV tiling assumes gl_NumSubgroups == 8 via gl_SubgroupID % 4 / gl_SubgroupID / 4. +// A different subgroup size indexes past row_max_sh/q_p_sh and corrupts the score tiles, so +// the pipeline is only created under `device->subgroup_size == 64` in ggml_vk_load_shaders. +// Keep that gate in sync with these constants. const uint TILE = 16; const uint HEAD_SIZE = 512; const uint HEADS_PER_GROUP = 32; @@ -146,10 +151,6 @@ void main() { SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); barrier(); - if (p.profile_stage == 1) { - continue; - } - { const uint head_local = tid / (SUBGROUP_SIZE / 4); const uint softmax_lane = tid % (SUBGROUP_SIZE / 4); @@ -198,10 +199,6 @@ void main() { accum[i] *= old_scale_sh[head_local]; } - if (p.profile_stage == 2) { - continue; - } - coopmat pmats[KEYS_PER_BLOCK / TILE]; [[unroll]] for (uint key_chunk = 0; key_chunk < KEYS_PER_BLOCK / TILE; ++key_chunk) { const uint pv_head_tile = gl_SubgroupID / (DIMS_PER_BLOCK / TILE); From 09f6a7dc6464a1acbdab388b8b455ef2c1aff94e Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 13 Aug 2026 14:32:59 +0000 Subject: [PATCH 089/109] test-backend-ops: cover sparse top-k FA with more than one sequence The sparse top-k fixture pinned every tensor's fourth dimension to 1, so no case reached the split path's stream indexing and three separate mis-indexings passed the suite. Add an ns parameter and three cases: two at ns=2 (with and without sinks, single tile) and one at ns=3 with 300 query tokens, which also crosses the 256-token tile boundary. The per-token selection is offset by the stream so a dropped stream stride reads another sequence's keys rather than the same ones. All three fail on the unfixed shaders and pass with them fixed. Co-Authored-By: Claude Opus 5 --- tests/test-backend-ops.cpp | 56 +++++++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 22 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 66421146149d..ce3fffa9d84d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7208,12 +7208,13 @@ struct test_flash_attn_ext_top_k : public test_case { const int64_t n_kv_raw; // dense prefix always attended const int64_t n_top_k; // selected keys per query token const bool sinks; + const int64_t ns; // sequences (ne3); >1 exercises the split-K stream stride static constexpr int64_t hs = 512; // V4 CSA head size, K == V latent static constexpr int64_t nh = 64; // V4 CSA query heads (MQA) std::string vars() override { - return VARS_TO_STR5(kv, nb, n_kv_raw, n_top_k, sinks); + return VARS_TO_STR6(kv, nb, n_kv_raw, n_top_k, sinks, ns); } double max_nmse_err() override { @@ -7224,27 +7225,27 @@ struct test_flash_attn_ext_top_k : public test_case { GGML_UNUSED(t); // only the active keys contribute compute on a sparse backend; count those so // perf mode reports the useful-work rate - return 2 * nh * nb * (hs + hs) * (n_kv_raw + n_top_k); + return 2 * nh * nb * ns * (hs + hs) * (n_kv_raw + n_top_k); } - test_flash_attn_ext_top_k(int64_t kv = 768, int64_t nb = 8, int64_t n_kv_raw = 64, int64_t n_top_k = 128, bool sinks = false) - : kv(kv), nb(nb), n_kv_raw(n_kv_raw), n_top_k(n_top_k), sinks(sinks) {} + test_flash_attn_ext_top_k(int64_t kv = 768, int64_t nb = 8, int64_t n_kv_raw = 64, int64_t n_top_k = 128, bool sinks = false, int64_t ns = 1) + : kv(kv), nb(nb), n_kv_raw(n_kv_raw), n_top_k(n_top_k), sinks(sinks), ns(ns) {} ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hs, nb, nh, 1); + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hs, nb, nh, ns); ggml_set_name(q, "q"); - ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, hs, kv, 1, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, hs, kv, 1, ns); ggml_set_name(k, "k"); // V4 CSA attends over the K latent itself: V is the same cache tensor - ggml_tensor * v = ggml_view_4d(ctx, k, hs, kv, 1, 1, k->nb[1], k->nb[2], k->nb[3], 0); + ggml_tensor * v = ggml_view_4d(ctx, k, hs, kv, 1, ns, k->nb[1], k->nb[2], k->nb[3], 0); ggml_set_name(v, "v"); - ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, 1); + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, ns); ggml_set_name(m, "m"); - ggml_tensor * t = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, nb, 1, 1); + ggml_tensor * t = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, nb, 1, ns); ggml_set_name(t, "top_k"); ggml_tensor * s = nullptr; @@ -7279,23 +7280,29 @@ struct test_flash_attn_ext_top_k : public test_case { // build a consistent (top_k, mask) pair: a deterministic per-token selection, // strided so adjacent tokens select overlapping-but-different keys, with one // deliberately invalid index (-1) whose mask slot stays -inf - std::vector top(n_top_k * nb); - std::vector mask(kv * nb); + std::vector top(n_top_k * nb * ns); + std::vector mask(kv * nb * ns); const ggml_fp16_t minus_inf = ggml_fp32_to_fp16(-INFINITY); const ggml_fp16_t zero = ggml_fp32_to_fp16(0.0f); - for (int64_t b = 0; b < nb; ++b) { - for (int64_t i = 0; i < kv; ++i) { - mask[b * kv + i] = i < n_kv_raw ? zero : minus_inf; - } - for (int64_t j = 0; j < n_top_k; ++j) { - int32_t idx = (int32_t) ((j * range) / n_top_k + b) % (int32_t) range; - if (j == n_top_k - 1 && b == 0) { - idx = -1; // exercise the ignore-invalid-index path - } else { - mask[b * kv + n_kv_raw + idx] = zero; + for (int64_t s = 0; s < ns; ++s) { + for (int64_t b = 0; b < nb; ++b) { + const int64_t mrow = (s * nb + b) * kv; + const int64_t trow = (s * nb + b) * n_top_k; + for (int64_t i = 0; i < kv; ++i) { + mask[mrow + i] = i < n_kv_raw ? zero : minus_inf; + } + for (int64_t j = 0; j < n_top_k; ++j) { + // offset the selection by the stream too, so a dropped stream stride + // reads another sequence's keys and shows up as a mismatch + int32_t idx = (int32_t) ((j * range) / n_top_k + b + s * 7) % (int32_t) range; + if (j == n_top_k - 1 && b == 0 && s == 0) { + idx = -1; // exercise the ignore-invalid-index path + } else { + mask[mrow + n_kv_raw + idx] = zero; + } + top[trow + j] = idx; } - top[b * n_top_k + j] = idx; } } @@ -10291,6 +10298,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k(1024, 64, 65, 128, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 128, 256, 512, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 257, 256, 512, false)); + // ns > 1: the split-K partial-output path indexes O and L/M by stream, so these cover + // the stream stride in both regions (single tile and multi-tile). + test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false, 2)); + test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, true, 2)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 300, 256, 512, false, 3)); return test_cases; } From 932245a60cb5aa4043e62db834830a3fbe17dd9f Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 13 Aug 2026 15:34:07 +0000 Subject: [PATCH 090/109] vulkan: let sparse FA query tiling and multiple sequences coexist flash_attn_split_k_reduce was using one value for two different things: ne2 is the split buffer's query count, which the sparse path hands the TILE height, but the destination's stream stride needs the FULL query count. With several tiles and more than one sequence it wrote stream s at s*tile_size instead of s*N, so the previous fix simply refused to tile whenever there was more than one stream. Give the reduce a separate dst_ne2 and use it for the destination index only. Ordinary FA split-K passes ne2 for it and is bit-identical. The sparse path passes the full batch, so tiling is unconditional again. This matters for memory, which is the binding constraint on this model. The scratch is capped at 256 queries per tile rather than scaling with the batch, so at ub2048 a 2-sequence context drops from 1028 MB to 128.5 MB and a 4-sequence one from 2056 MB to 257 MB. It also removes a cliff: 8 sequences at ub2048 would have exceeded maxStorageBufferRange and silently fallen back to the slower non-split kernel. Sparse FA timings are unchanged (every shape within 0.45%, inside run-to-run spread). 13,308 FLASH_ATTN_EXT cases pass, which covers the ordinary split-K path that shares this shader. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 18 ++++++++++-------- .../flash_attn_split_k_reduce.comp | 3 ++- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c5451751e16c..66b8fc7864d1 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2125,7 +2125,11 @@ struct vk_quantize_q8_1_push_constants { struct vk_op_flash_attn_split_k_reduce_push_constants { uint32_t D; uint32_t ne1; + // ne2 describes the SPLIT BUFFER (which may cover only a tile of queries); dst_ne2 is the + // destination's query count. They differ only when the caller tiles the split buffer, and + // the destination stride between streams must always use the full count. uint32_t ne2; + uint32_t dst_ne2; uint32_t ne3; uint32_t k_num; uint32_t sinks; @@ -11172,12 +11176,10 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & const uint32_t NS = (uint32_t) q->ne[3]; const uint32_t raw_kv = (uint32_t) n_kv_raw; const uint32_t partitions = 2; - // Query tiling keeps the split scratch small, but flash_attn_split_k_reduce derives the - // destination row from ne2, which it is handed as the TILE height. With one tile that - // equals N and the stream stride is right; with several tiles and more than one stream - // it would write stream s at s*tile_size instead of s*N. Only tile when there is a - // single stream, which is the prefill case the tiling exists for. - const uint32_t tile_size = (NS == 1) ? std::min(N, 256u) : N; + // Query tiling caps the split scratch at 256 queries regardless of batch or stream + // count. The reduce takes the tile height as ne2 and the full query count as dst_ne2, + // so the destination stream stride stays correct across tiles. + const uint32_t tile_size = std::min(N, 256u); const uint32_t n_tiles = CEIL_DIV(N, tile_size); const bool f32acc = true; vk_fa_tuning_params tuning = get_fa_tuning_params(ctx->device, D, D, N, raw_kv, GGML_TYPE_F16, GGML_TYPE_F16, f32acc); @@ -11272,7 +11274,7 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & ctx->prealloc_split_k_need_sync = true; ggml_vk_sync_buffers(ctx, subctx); - const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = {D, NH, tile_n, NS, partitions, sinks != nullptr}; + const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = {D, NH, tile_n, N, NS, partitions, sinks != nullptr}; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, {split_buf, sinks_buf, tile_dst}, reduce_pc, {NH, D, tile_n * NS}); ctx->prealloc_split_k_need_sync = true; @@ -11789,7 +11791,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx pc, { dispatch_x, workgroups_y, workgroups_z }); ggml_vk_sync_buffers(ctx, subctx); - const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, split_k, (sinks != nullptr) }; + const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne2, (uint32_t)ne3, split_k, (sinks != nullptr) }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, {split_k_buf, sinks_buf, dst_buf}, pc2, { (uint32_t)ne1, HSV, (uint32_t)(ne2 * ne3) }); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_split_k_reduce.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_split_k_reduce.comp index 68917fc0bb02..69342ba3ceea 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_split_k_reduce.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_split_k_reduce.comp @@ -14,6 +14,7 @@ layout (push_constant) uniform parameter { uint D; uint ne1; uint ne2; + uint dst_ne2; uint ne3; uint k_num; uint sinks; @@ -116,6 +117,6 @@ void main() { const float FLT_MAX = uintBitsToFloat(0x7F7FFFFF); O = clamp(O, -FLT_MAX, FLT_MAX); - data_d[(i3 * p.ne2 + i2) * p.ne1 * D + D * n + d] = O; + data_d[(i3 * p.dst_ne2 + i2) * p.ne1 * D + D * n + d] = O; } } From ff7ac3d701974fd47e60a23269287cc06ea1b34c Mon Sep 17 00:00:00 2001 From: Jaap Buurman Date: Thu, 13 Aug 2026 20:00:25 +0200 Subject: [PATCH 091/109] vulkan: parallelize DSV4 Lightning Indexer prefill Assisted-by: Codex --- .../DSV4-vulkan-lightning-indexer-progress.md | 158 ++++++++++++++++++ ggml/src/ggml-vulkan/ggml-vulkan.cpp | 18 +- .../vulkan-shaders/lightning_indexer_cm.comp | 94 +++++++---- .../vulkan-shaders/vulkan-shaders-gen.cpp | 3 +- tests/test-backend-ops.cpp | 9 + 5 files changed, 241 insertions(+), 41 deletions(-) create mode 100644 docs/development/DSV4-vulkan-lightning-indexer-progress.md diff --git a/docs/development/DSV4-vulkan-lightning-indexer-progress.md b/docs/development/DSV4-vulkan-lightning-indexer-progress.md new file mode 100644 index 000000000000..d19fbd6903aa --- /dev/null +++ b/docs/development/DSV4-vulkan-lightning-indexer-progress.md @@ -0,0 +1,158 @@ +# DeepSeek V4 Vulkan Lightning Indexer progress + +This is a restart note for the Strix Halo Lightning Indexer optimization. It is a development scratch pad and can be removed before the final PR. + +## Repository state + +- Main repository: `/home/jaap/Projects/git/llama.cpp` +- Optimization worktree: `/tmp/llama-strix-beta-bench` +- Branch: `strix-halo-vulkan-lightning-indexer` +- Base commit: `316c72ee9eab590f5891089d3b6bfc0d01d00d19` +- Base branch: Nathan's `strix-halo-vulkan-beta` +- Decode microbench work is stored in the main worktree as `stash@{0}: On strix-halo-vulkan: wip: DSV4 decode microbench depth matrix`. +- The Indexer changes are uncommitted. Do not commit without explicit user approval. An assisted commit needs an `Assisted-by:` trailer. +- Do not run builds and GPU benchmarks together. The APU shares its power and memory-bandwidth budget. +- GPU commands need sandbox escalation. + +## Objective and result + +After sparse prefill attention was flattened, the context-dependent Lightning Indexer became the next prefill bottleneck. The old cooperative-matrix shader used one wave64 subgroup per workgroup, processed one 16-key tile, and loaded one query head at a time. + +The new wide pipeline uses eight wave64 subgroups per workgroup. Each subgroup processes a separate 16-key tile, so one workgroup covers 128 keys. It stages four query heads and their weights together, reuses them across all eight subgroups, and uses subgroup-scoped synchronization between cooperative-matrix result stores. A workgroup barrier remains between four-head groups because all subgroups reuse the shared query storage. + +The optimized shader requires 512 workgroup invocations and 64 KiB shared memory. Pipeline creation is capability-based. Devices without those limits use a one-wave, one-head cooperative-matrix specialization. The scalar implementation remains the fallback when cooperative matrices are unavailable. The decode-specific cooperative-matrix pipeline is unchanged. + +At the 32k-equivalent prefill microbench shape: + +| Version | Time per layer | Throughput | +| --- | ---: | ---: | +| Baseline | 50.51 ms | 5.83 TFLOPS | +| Optimized | 31.81 ms | 9.25 TFLOPS | + +This is a 37.0% reduction in Lightning Indexer kernel time. + +The canonical 32k llama-bench improved from 209.45 to 216.32 tokens/s. Total Vulkan time fell from 9.73729 to 9.42448 seconds. Total Lightning Indexer time fell from 1.11860 to 0.697276 seconds. Sparse attention and top-K were effectively unchanged. + +## Changed files + +- `ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp`: parameterizes the shader, adds the eight-wave four-head implementation, and remains usable for the small fallback. +- `ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp`: generates wide `N_WAVES=8`, `HEADS_PER_TILE=4` and small `N_WAVES=1`, `HEADS_PER_TILE=1` variants. +- `ggml/src/ggml-vulkan/ggml-vulkan.cpp`: creates and selects the capability-gated wide pipeline and the small cooperative-matrix fallback. +- `tests/test-backend-ops.cpp`: adds 127, 128, and 129-key correctness boundaries and PP2048 performance shapes through 512k simulated source context. + +## Performance data + +The performance rows model the actual PP2048 Indexer shapes after the source context is filled in 2048-token batches. `kv=8704` is the measured shape near 32k source context. It differs from 32768 because the Indexer compresses source tokens into rows. + +| Source depth | `kv` | Baseline | Optimized | Reduction | +| ---: | ---: | ---: | ---: | ---: | +| 0 | 512 | 4.23 ms | 2.14 ms | 49.5% | +| 8k | 2560 | 17.10 ms | 10.09 ms | 41.0% | +| 16k | 4608 | 29.97 ms | 17.62 ms | 41.2% | +| 32k | 8704 | 50.51 ms | 32.94 ms | 34.8% | +| 64k | 16896 | 94.36 ms | 64.87 ms | 31.3% | +| 128k | 33280 | 174.96 ms | 128.76 ms | 26.4% | +| 256k | 66048 | 352.73 ms | 250.88 ms | 28.9% | +| 512k | 131584 | 696.79 ms | 495.90 ms | 28.8% | + +The final isolated 32k run after cleanup measured 31.81159 ms. Small matrix differences are normal laptop GPU clock variation. + +| Canonical 32k metric | Baseline | Optimized | +| --- | ---: | ---: | +| PP2048 | 209.45 tokens/s | 216.32 tokens/s | +| Total Vulkan | 9.73729 s | 9.42448 s | +| Lightning Indexer | 1.11860 s | 0.697276 s | +| Sparse FA raw | 0.198438 s | 0.195367 s | +| Sparse FA selected | 0.835110 s | 0.843385 s | +| Sparse FA reduce | 0.088463 s | 0.086678 s | +| TOP_K | 0.075473 s | 0.075350 s | + +Logs: + +- Baseline matrix: `/tmp/dsv4-lightning-prefill-baseline.log` +- Optimized matrix: `/tmp/dsv4-lightning-four-head-matrix.log` +- Final selected-pipeline 32k microbench: `/tmp/dsv4-lightning-final-selected-32k.log` +- Baseline canonical 32k llama-bench: `/tmp/dsv4-nathan-beta-32k-rerun-new-first.log` +- Optimized canonical 32k llama-bench: `/tmp/dsv4-lightning-final-32k-llama-bench.log` + +## Correctness and resources + +- Wide pipeline: all 20 focused F16 cases passed, including 127, 128, and 129-key boundaries. +- Small cooperative-matrix fallback: temporarily forced and all the same 20 cases passed. +- Wide shader on gfx1151: 168 VGPRs, 63,488 bytes LDS, no spills, eight subgroups per SIMD. +- The final pipeline-statistics run confirmed `lightning_indexer_cm_f16` was selected. +- No NaN, Inf, or comparison failures were reported. + +Correctness logs: + +- Wide: `/tmp/dsv4-lightning-consolidated-wide-correctness.log` +- Small fallback: `/tmp/dsv4-lightning-consolidated-small-correctness.log` + +## Experiments and decisions + +- Four waves improved the 32k shape about 3% and became worse at deep simulated contexts. +- Eight waves improved it about 8% before the other changes. +- Subgroup-scoped synchronization after cooperative-matrix stores improved the eight-wave version. +- Staging four query heads and weights produced the large gain by reducing redundant loads and barriers. +- Using only a subgroup barrier between head groups failed four boundary tests. One subgroup could overwrite shared query data while another still read it. A workgroup barrier is required there. +- A separate fallback shader source was avoided. Generator definitions create both variants from one file. + +## Commands + +Build only, with no GPU benchmark running: + +```sh +cd /tmp/llama-strix-beta-bench +git diff --check +cmake --build build --config Release --target test-backend-ops llama-bench -j "$(nproc)" +``` + +Focused correctness: + +```sh +cd /tmp/llama-strix-beta-bench +./build/bin/test-backend-ops test -b Vulkan0 -o LIGHTNING_INDEXER -p 'type_K=f16' > /tmp/dsv4-lightning-correctness.log 2>&1 +tail -n 30 /tmp/dsv4-lightning-correctness.log +``` + +Final 32k-equivalent microbench and pipeline selection: + +```sh +cd /tmp/llama-strix-beta-bench +GGML_VK_PIPELINE_STATS=lightning_indexer_cm_f16 ./build/bin/test-backend-ops perf -b Vulkan0 -o LIGHTNING_INDEXER -p 'kv=8704' > /tmp/dsv4-lightning-final-selected-32k.log 2>&1 +tail -n 16 /tmp/dsv4-lightning-final-selected-32k.log +``` + +Full Indexer depth matrix: + +```sh +cd /tmp/llama-strix-beta-bench +./build/bin/test-backend-ops perf -b Vulkan0 -o LIGHTNING_INDEXER -p 'nb=2048,nh=64,ns=1,nm=1,type_K=f16' > /tmp/dsv4-lightning-matrix.log 2>&1 +rg 'kv=(512|2560|4608|8704|16896|33280|66048|131584),nb=2048' /tmp/dsv4-lightning-matrix.log +``` + +Canonical 32k llama-bench. Run it only when needed, never while compiling, and inspect only the final block: + +```sh +cd /tmp/llama-strix-beta-bench +GGML_VK_PERF_LOGGER=1 ./build/bin/llama-bench -m /home/jaap/Projects/docker/localLLaMA/models/models--unsloth--DeepSeek-V4-Flash-0731-GGUF/snapshots/109848da2469efe1f1aab9e11acea08a065ccd4f/UD-IQ3_XXS/DeepSeek-V4-Flash-0731-UD-IQ3_XXS-00001-of-00004.gguf -r 1 -d 32768 -p 2048 -ub 2048 -fa 1 -n 0 > /tmp/dsv4-lightning-32k-llama-bench.log 2>&1 +last=$(grep -n 'Vulkan Timings:' /tmp/dsv4-lightning-32k-llama-bench.log | tail -n 1 | cut -d: -f1) +sed -n "${last},\$p" /tmp/dsv4-lightning-32k-llama-bench.log | tail -n 180 +``` + +Patch inspection: + +```sh +cd /tmp/llama-strix-beta-bench +git diff --check +git diff --stat +git diff +git status --short +``` + +## Next actions + +1. The user reviews and understands the four-file implementation and result summary. +2. Commit only after explicit user approval for that commit action. +3. Remove this scratch pad before a PR if it is not useful as permanent documentation. +4. Restore the separate decode microbench stash from the main worktree only if that work resumes. diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 66b8fc7864d1..095274cf1413 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1088,6 +1088,7 @@ struct vk_device_struct { vk_pipeline pipeline_gated_delta_net[4][2]; vk_pipeline pipeline_lightning_indexer_f16; vk_pipeline pipeline_lightning_indexer_cm_f16; + vk_pipeline pipeline_lightning_indexer_cm_small_f16; vk_pipeline pipeline_lightning_indexer_decode_cm_f16; vk_pipeline pipeline_flash_attn_top_k_f16; vk_pipeline pipeline_flash_attn_top_k_cm_f16; @@ -6076,10 +6077,18 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { device->subgroup_size); #if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) if (device->coopmat_support && device->coopmat_support_16x16x16_f32acc && device->subgroup_size_control) { - ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_f16, - "lightning_indexer_cm_f16", lightning_indexer_cm_f16_len, lightning_indexer_cm_f16_data, "main", 5, + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_small_f16, + "lightning_indexer_cm_small_f16", lightning_indexer_cm_small_f16_len, lightning_indexer_cm_small_f16_data, "main", 5, sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 16, 1}, {device->subgroup_size}, 1, true, true, device->subgroup_size); + if (device->properties.limits.maxComputeWorkGroupInvocations >= 512 && + device->properties.limits.maxComputeWorkGroupSize[0] >= 512 && + device->properties.limits.maxComputeSharedMemorySize >= 64 * 1024) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_f16, + "lightning_indexer_cm_f16", lightning_indexer_cm_f16_len, lightning_indexer_cm_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {128, 16, 1}, {device->subgroup_size}, 1, true, true, + device->subgroup_size); + } ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_decode_cm_f16, "lightning_indexer_decode_cm_f16", lightning_indexer_decode_cm_f16_len, lightning_indexer_decode_cm_f16_data, "main", 5, sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 1, 1}, {device->subgroup_size}, 1, true, true, @@ -12379,8 +12388,9 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const if (ctx->device->pipeline_lightning_indexer_decode_cm_f16 && src0->ne[2] == 1) { return ctx->device->pipeline_lightning_indexer_decode_cm_f16; } - return ctx->device->pipeline_lightning_indexer_cm_f16 && src0->ne[2] >= 16 ? - ctx->device->pipeline_lightning_indexer_cm_f16 : ctx->device->pipeline_lightning_indexer_f16; + vk_pipeline cm = ctx->device->pipeline_lightning_indexer_cm_f16 ? + ctx->device->pipeline_lightning_indexer_cm_f16 : ctx->device->pipeline_lightning_indexer_cm_small_f16; + return cm && src0->ne[2] >= 16 ? cm : ctx->device->pipeline_lightning_indexer_f16; } // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer() if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp index a0a3639d2546..c53eb74f8d4c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp @@ -5,9 +5,14 @@ #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require #extension GL_KHR_cooperative_matrix : require #extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : require layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +#if N_WAVES == 1 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; +#else +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; +#endif layout(binding = 0) readonly buffer QBuf { float data_q[]; }; layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; @@ -40,13 +45,16 @@ const uint VEC_PER_HEAD = HEAD_SIZE / 4; const uint TILE_STRIDE = VEC_PER_HEAD + 2; const uint SCORE_STRIDE = TILE / 4 + 1; -shared f16vec4 q_sh[TILE * TILE_STRIDE]; -shared f16vec4 k_sh[TILE * TILE_STRIDE]; -shared vec4 score_sh[TILE * SCORE_STRIDE]; +shared f16vec4 q_sh[HEADS_PER_TILE][TILE * TILE_STRIDE]; +shared f16vec4 k_sh[N_WAVES][TILE * TILE_STRIDE]; +shared vec4 score_sh[N_WAVES][TILE * SCORE_STRIDE]; +shared float weight_sh[HEADS_PER_TILE][TILE]; void main() { const uint tid = gl_LocalInvocationIndex; - const uint kv_base = gl_WorkGroupID.x * TILE; + const uint lane = gl_SubgroupInvocationID; + const uint wave = gl_SubgroupID; + const uint kv_base = (gl_WorkGroupID.x * N_WAVES + wave) * TILE; const uint token_base = gl_WorkGroupID.y * TILE; const uint stream = gl_WorkGroupID.z; @@ -55,63 +63,77 @@ void main() { totals[i] = 0.0; } - for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { - const uint key = idx / VEC_PER_HEAD; - const uint d4 = idx % VEC_PER_HEAD; - const uint kv = kv_base + key; + for (uint idx = tid; idx < N_WAVES * TILE * VEC_PER_HEAD; idx += gl_WorkGroupSize.x) { + const uint load_wave = idx / (TILE * VEC_PER_HEAD); + const uint wave_idx = idx % (TILE * VEC_PER_HEAD); + const uint key = wave_idx / VEC_PER_HEAD; + const uint d4 = wave_idx % VEC_PER_HEAD; + const uint kv = (gl_WorkGroupID.x * N_WAVES + load_wave) * TILE + key; f16vec4 value = f16vec4(0.0); if (kv < p.n_kv) { const uint offset = stream * p.nbk3 + kv * p.nbk2 + d4 * 4; value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); } - k_sh[key * TILE_STRIDE + d4] = value; + k_sh[load_wave][key * TILE_STRIDE + d4] = value; } barrier(); - for (uint head = 0; head < N_HEAD; ++head) { - for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { - const uint token_local = idx / VEC_PER_HEAD; - const uint d4 = idx % VEC_PER_HEAD; + for (uint head_base = 0; head_base < N_HEAD; head_base += HEADS_PER_TILE) { + for (uint idx = tid; idx < HEADS_PER_TILE * TILE * VEC_PER_HEAD; idx += gl_WorkGroupSize.x) { + const uint head_local = idx / (TILE * VEC_PER_HEAD); + const uint head_idx = idx % (TILE * VEC_PER_HEAD); + const uint token_local = head_idx / VEC_PER_HEAD; + const uint d4 = head_idx % VEC_PER_HEAD; const uint token = token_base + token_local; f16vec4 value = f16vec4(0.0); if (token < p.n_batch) { - const uint offset = stream * p.nbq3 + token * p.nbq2 + head * p.nbq1 + d4 * 4; + const uint offset = stream * p.nbq3 + token * p.nbq2 + (head_base + head_local) * p.nbq1 + d4 * 4; value = f16vec4(data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); } - q_sh[token_local * TILE_STRIDE + d4] = value; + q_sh[head_local][token_local * TILE_STRIDE + d4] = value; + } + if (tid < HEADS_PER_TILE * TILE) { + const uint head_local = tid / TILE; + const uint token_local = tid % TILE; + const uint token = token_base + token_local; + weight_sh[head_local][token_local] = token < p.n_batch ? data_w[stream * p.nbw3 + token * p.nbw1 + head_base + head_local] : 0.0; } barrier(); - coopmat scores = - coopmat(0.0); - coopmat kmat; - coopmat qmat; + [[unroll]] for (uint head_local = 0; head_local < HEADS_PER_TILE; ++head_local) { + coopmat scores = + coopmat(0.0); + coopmat kmat; + coopmat qmat; - [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { - coopMatLoad(kmat, k_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); - coopMatLoad(qmat, q_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); - scores = coopMatMulAdd(kmat, qmat, scores); - } + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(kmat, k_sh[wave], d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(qmat, q_sh[head_local], d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + scores = coopMatMulAdd(kmat, qmat, scores); + } - coopMatStore(scores, score_sh, 0, SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); - barrier(); + coopMatStore(scores, score_sh[wave], 0, SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease); - [[unroll]] for (uint i = 0; i < 4; ++i) { - const uint idx = tid + i * SUBGROUP_SIZE; - const uint key = idx / TILE; - const uint token_local = idx % TILE; - const uint token = token_base + token_local; - if (token < p.n_batch && kv_base + key < p.n_kv) { - const float score = score_sh[key * SCORE_STRIDE + token_local / 4][token_local % 4]; - const float weight = data_w[stream * p.nbw3 + token * p.nbw1 + head]; - totals[i] += max(score, 0.0) * weight; + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint idx = lane + i * SUBGROUP_SIZE; + const uint key = idx / TILE; + const uint token_local = idx % TILE; + const uint token = token_base + token_local; + if (token < p.n_batch && kv_base + key < p.n_kv) { + const float score = score_sh[wave][key * SCORE_STRIDE + token_local / 4][token_local % 4]; + totals[i] += max(score, 0.0) * weight_sh[head_local][token_local]; + } + } + if (head_local + 1 < HEADS_PER_TILE) { + controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease); } } barrier(); } [[unroll]] for (uint i = 0; i < 4; ++i) { - const uint idx = tid + i * SUBGROUP_SIZE; + const uint idx = lane + i * SUBGROUP_SIZE; const uint key = idx / TILE; const uint token_local = idx % TILE; const uint kv = kv_base + key; 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 845720ad8b1e..8b0916a18a2f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -803,7 +803,8 @@ void process_shaders() { string_to_spv("lightning_indexer_f16", "lightning_indexer_scalar64.comp", {}); #if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) - string_to_spv("lightning_indexer_cm_f16", "lightning_indexer_cm.comp", {}); + string_to_spv("lightning_indexer_cm_f16", "lightning_indexer_cm.comp", {{"N_WAVES", "8"}, {"HEADS_PER_TILE", "4"}}); + string_to_spv("lightning_indexer_cm_small_f16", "lightning_indexer_cm.comp", {{"N_WAVES", "1"}, {"HEADS_PER_TILE", "1"}}); string_to_spv("lightning_indexer_decode_cm_f16", "lightning_indexer_decode_cm.comp", {}); #endif string_to_spv("flash_attn_top_k_f16", "flash_attn_top_k.comp", {}); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index ce3fffa9d84d..b2e65f10de47 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10304,6 +10304,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, true, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 300, 256, 512, false, 3)); + for (int kv : { 127, 128, 129 }) { + test_cases.emplace_back(new test_lightning_indexer(128, 64, kv, 32, 4, 1, GGML_TYPE_F16)); + } + return test_cases; } #ifdef _MSC_VER @@ -10718,6 +10722,11 @@ static std::vector> make_test_cases_perf() { } } } + // DSV4 PP2048 indexer rows after filling source contexts from 8k through 512k. + // The zero-depth kv=512 shape is covered above. + for (int kv : { 2560, 4608, 8704, 16896, 33280, 66048, 131584 }) { + test_cases.emplace_back(new test_lightning_indexer(128, 64, kv, 2048, 1, 1, GGML_TYPE_F16)); + } // sparse top-k FA at V4 decode/prefill shapes — the A/B instrument for the // gather-to-compact work (n_active = n_kv_raw + n_top_k stays fixed as kv grows). From 93f18c376d53914ec9637b3f5560a741764d8aca Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 13 Aug 2026 16:00:20 +0000 Subject: [PATCH 092/109] vulkan: extend DeepSeek V4 gather-to-compact to small batches The gap and its consequence were diagnosed by Jaap Buurman: the sparse prefill path gates on batch >= 64 and gather-to-compact gated on batch == 1, so batch 2..63 fell through to dense attention over the whole compressed KV, at a cost that grows with context. That is where a speculative draft lands (n_max 3-5), which is why token generation dropped off sharply with DSpark enabled. This implements the fix for that diagnosis. Give each query token its own gathered top-k block rather than deduplicating into a union. No dedup pass, no atomics, and the size is bounded by n_kv_raw + n_batch*n_top_k regardless of depth. Cross-token rows are neutralised through the mask, which already encodes each token's selection: token t reads -inf on any block that is not its own, so the softmax cannot double count. The compact mask is token-major [n_batch][kv_c], which is what the GQA mask path already expects (m_stride 0, rows stepped by gqa_iq1 * m_row_len). Measured on gfx1151, test-backend-ops perf, medians of 2 launches, n_kv_raw=2304 n_top_k=512. Gather cost is flat in depth (894 us at batch 4 at every depth); dense is not: kv rows batch dense gathered speedup 11008 2 2194 us 680 us 3.23x 11008 4 2204 us 895 us 2.46x 11008 8 2217 us 2217 us 1.00x (gate declines: kv < 2*kv_c) 35584 2 7065 us 680 us 10.38x 35584 4 7099 us 895 us 7.93x 35584 8 7114 us 1316 us 5.41x 133888 2 7552 us 680 us 11.10x 133888 4 7971 us 894 us 8.92x 133888 8 8348 us 1320 us 6.32x Batch 1 is unchanged in behaviour and stays on the same code path. Not validated end to end. A deduplicated union would shrink the gathered set further wherever adjacent draft tokens select overlapping keys, and would lift the batch ceiling documented in the following commit. Suggested-by: Jaap Buurman (@Mushoz) Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 19 +++++-- .../vulkan-shaders/flash_attn_gather.comp | 52 ++++++++++++++----- tests/test-backend-ops.cpp | 17 ++++++ 3 files changed, 70 insertions(+), 18 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 095274cf1413..f58ff247b1e7 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1958,7 +1958,7 @@ static_assert(sizeof(vk_op_dsv4_hc_post_push_constants) <= 128); struct vk_op_flash_attn_gather_push_constants { uint32_t n_kv, n_kv_raw, n_top_k, kv_c; - uint32_t nbk1, nbk3, nbt3, nbm3, nem3; + uint32_t nbk1, nbk3, nbt1, nbt3, nbm1, nbm3, nem3, n_batch; }; static_assert(sizeof(vk_op_flash_attn_gather_push_constants) <= 128); @@ -11306,6 +11306,7 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & struct vk_fa_compact_state { bool active = false; uint32_t kv_c = 0; + uint32_t n_batch = 1; vk_subbuffer kc_buf, mc_buf; }; @@ -11323,7 +11324,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ static const char * gather_env = getenv("GGML_VK_FA_TOPK_GATHER"); if ((gather_env && gather_env[0] == '0') || !top_k || !ctx->device->pipeline_flash_attn_gather_f16 || - q->ne[1] != 1 || // single-token decode only; batched queries need a union gather + q->ne[1] < 1 || q->ne[1] >= 64 || // 1..63: >=64 goes to the sparse prefill path q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || q->ne[0] != 512 || k->ne[0] != 512 || v->ne[0] != 512 || q->ne[2] != 64 || @@ -11347,7 +11348,10 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ return false; } - const uint32_t kv_c = GGML_PAD((uint32_t)(n_kv_raw + top_k->ne[0]), 256u); + const uint32_t n_batch = (uint32_t) q->ne[1]; + // every token gets its own top-k block; bounded by n_kv_raw + n_batch*n_top_k regardless + // of context depth, which is the whole point at decode + const uint32_t kv_c = GGML_PAD((uint32_t)(n_kv_raw + (int64_t) n_batch * top_k->ne[0]), 256u); // the gather writes then re-reads ~the active bytes; dense reads the source KV once, // so compaction only pays when the source is comfortably larger than the active set if ((uint64_t) k->ne[1] < 2ull * kv_c) { @@ -11356,7 +11360,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ const uint32_t ns = (uint32_t) q->ne[3]; const size_t kc_sz = (size_t) ns * kv_c * 512 * sizeof(ggml_fp16_t); - const size_t mc_sz = (size_t) ns * kv_c * sizeof(ggml_fp16_t); + const size_t mc_sz = (size_t) ns * n_batch * kv_c * sizeof(ggml_fp16_t); if (ctx->prealloc_size_y < kc_sz + mc_sz) { ctx->prealloc_size_y = kc_sz + mc_sz; @@ -11373,9 +11377,12 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ (uint32_t) k->ne[1], (uint32_t) n_kv_raw, (uint32_t) top_k->ne[0], kv_c, (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), (uint32_t) (top_k->nb[3] / sizeof(int32_t)), + (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (mask->nb[3] / sizeof(ggml_fp16_t)), (uint32_t) mask->ne[3], + n_batch, }; st.kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); @@ -11389,6 +11396,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ st.active = true; st.kv_c = kv_c; + st.n_batch = n_batch; return true; } @@ -11453,6 +11461,9 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx if (ggml_vk_flash_attn_gather_compact(ctx, subctx, q, k, v, mask, dst, fa_compact)) { KV = fa_compact.kv_c; nem0 = fa_compact.kv_c; + nem1 = fa_compact.n_batch; + nem2 = 1; + nem3 = (uint32_t) q->ne[3]; nem1 = N; nem2 = 1; nem3 = (uint32_t) q->ne[3]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp index 4e3dc1c42623..3c9ac68e2c69 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp @@ -4,9 +4,15 @@ #extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require -// Gathers the active KV rows of a top-k sparse attention (DeepSeek V4 CSA decode) into a -// compact contiguous scratch: rows [0, n_kv_raw) of the source (the dense prefix), then the -// n_top_k selected rows, then zero padding up to kv_c. The gathered mask row keeps the +// Gathers the active KV rows of a top-k sparse attention (DeepSeek V4 CSA) into a compact +// contiguous scratch: rows [0, n_kv_raw) of the source (the dense prefix), then, for each of +// the n_batch query tokens, that token's n_top_k selected rows, then zero padding up to kv_c. +// +// n_batch > 1 (small-batch decode, e.g. speculative drafts) gives every token its own block +// rather than deduplicating into a union. That costs n_batch*n_top_k gathered rows instead of +// |union|, but needs no dedup pass and no atomics, and the size is bounded by the worst case +// the caller already allocates for. Cross-token rows are neutralised through the mask: token t +// sees -inf on any block that is not its own, so nothing is double counted in the softmax. The gathered mask row keeps the // per-key mask values so causality/validity survive compaction; invalid top-k indices and // padding get -inf mask and zeroed K (softmax-neutral either way, zeroed so no NaN*0). // One workgroup per compact row; V is the K latent (V==K), so a single gather serves both. @@ -26,9 +32,12 @@ layout(push_constant) uniform Parameters { uint kv_c; // padded compact row count == dispatch row range uint nbk1; // K source row stride, elements uint nbk3; // K source stream stride, elements + uint nbt1; // top_k row (per query token) stride, elements uint nbt3; // top_k stream stride, elements + uint nbm1; // mask source row (per query token) stride, elements uint nbm3; // mask source stream stride, elements uint nem3; // mask ne[3], for stream broadcast + uint n_batch; // query tokens sharing this gather; <= LANES } p; const uint HEAD_SIZE = 512; @@ -39,14 +48,23 @@ void main() { const uint stream = gl_WorkGroupID.z; const uint tid = gl_LocalInvocationIndex; - // map compact row -> source row; p.n_kv is the invalid sentinel - uint src = p.n_kv; + // map compact row -> source row; p.n_kv is the invalid sentinel. + // owner is the token whose block this row belongs to, or ALL_TOKENS for the shared prefix. + const uint ALL_TOKENS = 0xffffffffu; + uint src = p.n_kv; + uint owner = ALL_TOKENS; if (row < p.n_kv_raw) { src = row; - } else if (row < p.n_kv_raw + p.n_top_k) { - const int idx = data_top[stream * p.nbt3 + (row - p.n_kv_raw)]; - if (idx >= 0 && uint(idx) < p.n_kv - p.n_kv_raw) { - src = p.n_kv_raw + uint(idx); + } else { + const uint off = row - p.n_kv_raw; + const uint tok = off / p.n_top_k; + const uint slot = off - tok * p.n_top_k; + if (tok < p.n_batch) { + owner = tok; + const int idx = data_top[stream * p.nbt3 + tok * p.nbt1 + slot]; + if (idx >= 0 && uint(idx) < p.n_kv - p.n_kv_raw) { + src = p.n_kv_raw + uint(idx); + } } } @@ -56,15 +74,21 @@ void main() { [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { data_kc[dst_base + tid + i * LANES] = data_k[src_base + tid + i * LANES]; } - if (tid == 0) { - data_mc[stream * p.kv_c + row] = data_m[(stream % p.nem3) * p.nbm3 + src]; - } } else { [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { data_kc[dst_base + tid + i * LANES] = float16_t(0.0); } - if (tid == 0) { - data_mc[stream * p.kv_c + row] = float16_t(uintBitsToFloat(0xff800000)); + } + + // Compact mask is token-major [n_batch][kv_c], which is what the GQA mask path expects + // (m_stride 0, rows stepped by gqa_iq1 * m_row_len). One lane per token; n_batch <= LANES. + const float NEG_INF = uintBitsToFloat(0xff800000); + if (tid < p.n_batch) { + const uint mc_idx = (stream * p.n_batch + tid) * p.kv_c + row; + float mv = NEG_INF; + if (src < p.n_kv && (owner == ALL_TOKENS || owner == tid)) { + mv = float(data_m[(stream % p.nem3) * p.nbm3 + tid * p.nbm1 + src]); } + data_mc[mc_idx] = float16_t(mv); } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index b2e65f10de47..4d3176fa90a2 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10300,6 +10300,16 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 257, 256, 512, false)); // ns > 1: the split-K partial-output path indexes O and L/M by stream, so these cover // the stream stride in both regions (single tile and multi-tile). + // small-batch decode (speculative drafts): each token gets its own gathered top-k block, + // so cross-token rows must be masked out or the softmax double counts. kv must be large + // enough that compaction is worth it (the gather gates on kv >= 2*kv_c). + test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 2, 1024, 512, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 3, 1024, 512, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 4, 1024, 512, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 8, 1024, 512, true)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(32768, 16, 2304, 512, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(65536, 63, 2304, 512, false)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 4, 1024, 512, false, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, true, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 300, 256, 512, false, 3)); @@ -10742,6 +10752,13 @@ static std::vector> make_test_cases_perf() { for (int kv : { 11008, 19200, 35584, 68352, 133888 }) { test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, 2048, 2304, 512, false)); } + // small-batch decode at depth: the speculative-draft regime (batch 2-8), where the old + // path fell through to dense attention over the whole compressed KV. + for (int kv : { 11008, 35584, 133888 }) { + for (int nb : { 1, 2, 4, 8 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, nb, 2304, 512, false)); + } + } return test_cases; } From 4c358646881590784e7387e788697c2961df5a40 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 13 Aug 2026 22:20:56 +0000 Subject: [PATCH 093/109] vulkan: record the resource limits behind the two DSV4 prefill kernels Both additions sit close to a limit that is invisible at the call site. The parallel Lightning Indexer uses 62720 of 65536 bytes of shared memory at N_WAVES=8 / HEADS_PER_TILE=4, so exactly one workgroup fits per CU. That is the intended trade, but raising either constant overruns the budget and the pipeline then fails to create and silently falls back to the small variant. Write the arithmetic down next to the arrays. The small-batch gather's compact set is independent of context depth but grows with batch, so its attention work is quadratic in batch against dense's linear. The kv >= 2*kv_c gate already caps this at about batch 6 at 32k depth and ~30 at 128k, and beyond the cap dense runs instead, so it is never slower. But the measured speedups do not show that ceiling and it is the main argument for building the deduplicated union later. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 10 ++++++++-- .../vulkan-shaders/lightning_indexer_cm.comp | 10 ++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f58ff247b1e7..d8387b14023d 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11349,8 +11349,14 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ } const uint32_t n_batch = (uint32_t) q->ne[1]; - // every token gets its own top-k block; bounded by n_kv_raw + n_batch*n_top_k regardless - // of context depth, which is the whole point at decode + // Every token gets its own top-k block, so the compact set is n_kv_raw + n_batch*n_top_k: + // independent of context depth, which is the point, but GROWING WITH BATCH. Attention work + // is then n_batch * (n_kv_raw + n_batch*n_top_k), i.e. quadratic in batch, against dense's + // n_batch * n_kv. Break-even is n_batch = (n_kv - n_kv_raw) / n_top_k, and the + // kv >= 2*kv_c gate below caps the useful batch at (n_kv/2 - n_kv_raw) / n_top_k -- + // about 6 at 32k depth, ~30 at 128k, batch-capped at 512k. Beyond that the gate declines + // and dense runs, so this can never be slower; it just stops helping. A deduplicated union + // would lift that ceiling wherever draft tokens select overlapping keys. const uint32_t kv_c = GGML_PAD((uint32_t)(n_kv_raw + (int64_t) n_batch * top_k->ne[0]), 256u); // the gather writes then re-reads ~the active bytes; dense reads the source KV once, // so compaction only pays when the source is comfortably larger than the active set diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp index c53eb74f8d4c..d4379e8899d0 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp @@ -45,6 +45,16 @@ const uint VEC_PER_HEAD = HEAD_SIZE / 4; const uint TILE_STRIDE = VEC_PER_HEAD + 2; const uint SCORE_STRIDE = TILE / 4 + 1; +// Shared-memory budget, N_WAVES=8 / HEADS_PER_TILE=4 (TILE 16, HEAD_SIZE 128): +// k_sh 8 * 16 * 34 * 8 B = 34816 +// q_sh 4 * 16 * 34 * 8 B = 17408 +// score_sh 8 * 16 * 5 * 16 B = 10240 +// weight_sh 4 * 16 * 4 B = 256 +// total = 62720 of 65536 (95.7%) +// Only one workgroup fits per CU at that size, which is the intended trade. Raising either +// constant overruns: HEADS_PER_TILE=8 needs 80128 B and the pipeline then fails to create, +// silently falling back to the small variant. The host gates on +// maxComputeSharedMemorySize >= 64 KiB; keep that gate and this arithmetic in sync. shared f16vec4 q_sh[HEADS_PER_TILE][TILE * TILE_STRIDE]; shared f16vec4 k_sh[N_WAVES][TILE * TILE_STRIDE]; shared vec4 score_sh[N_WAVES][TILE * SCORE_STRIDE]; From 7532a73164d6cc16b356a2def64997619542617b Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Fri, 14 Aug 2026 00:01:20 +0000 Subject: [PATCH 094/109] vulkan: deduplicated union for DeepSeek V4 small-batch decode Replaces the per-token top-k blocks with one row per DISTINCT selected key, so the compact set stops growing linearly with batch. Measured on the real model, adjacent tokens share 60% of their selections over 4 tokens and 76% over 8, so the union is materially smaller than n_batch*n_top_k. The size is only known on the GPU, so flash-attention now reads its KV bound from a buffer instead of the push constant, behind a new DYNAMIC_KV pipeline flag. Every other pipeline folds the flag away at compile time and is unchanged (13308 FLASH_ATTN_EXT cases pass identically with the union off). No indirect dispatch is needed: FA workgroup counts come from neq1/neq2/neq3 and never from KV, so only the loop bound moves. Padding the count to 256 keeps KV % Bc == 0, which lets the aligned pipeline variant still apply. Dedup marks a bitmap from the top-k lists and compacts by scanning bitmap WORDS. An earlier version scanned the mask row by row: simpler, but it made dedup cost scale with depth and measured 0.71x at kv=133888, ie a loss. The bitmap form is O(n_batch*n_top_k) to mark and R/32 to compact, and is depth-independent. test-backend-ops perf, medians of 2, n_kv_raw=2304 n_top_k=512, fixture overlap 60% (the default generator produces near-zero overlap and would make the union look worthless by construction): kv batch per-token union speedup 35584 2 680 us 632 us 1.08x 35584 4 892 us 720 us 1.24x 35584 8 1315 us 900 us 1.46x 133888 2 679 us 638 us 1.06x 133888 4 892 us 726 us 1.23x 133888 8 1314 us 909 us 1.45x Union cost is flat in depth (720 vs 726 us at batch 4 across a 3.8x depth range), which the per-token form was not. Opt-in via GGML_VK_FA_TOPK_UNION=1, single stream only, and it falls back to the per-token blocks when the compressed region exceeds the shared bitmap. The kv >= 2*kv_c gate still uses the worst case, so the batch ceiling is NOT yet lifted; doing that needs a GPU-side fallback for an oversized union. Suggested-by: Jaap Buurman (@Mushoz) Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 209 +++++++++++++++++- .../vulkan-shaders/flash_attn_base.glsl | 8 +- .../flash_attn_gather_union.comp | 82 +++++++ .../vulkan-shaders/flash_attn_union.comp | 122 ++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 2 + tests/test-backend-ops.cpp | 20 +- 6 files changed, 428 insertions(+), 15 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index d8387b14023d..4ae3cc7ed80b 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -64,6 +64,7 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { #include #include #include +#include #include #include #include @@ -1093,6 +1094,8 @@ struct vk_device_struct { vk_pipeline pipeline_flash_attn_top_k_f16; vk_pipeline pipeline_flash_attn_top_k_cm_f16; vk_pipeline pipeline_flash_attn_gather_f16; + vk_pipeline pipeline_flash_attn_union_f16; + vk_pipeline pipeline_flash_attn_gather_union_f16; vk_pipeline pipeline_dsv4_hc_pre_f32; vk_pipeline pipeline_dsv4_hc_comb_f32; vk_pipeline pipeline_dsv4_hc_post_f32; @@ -1956,6 +1959,12 @@ struct vk_op_dsv4_hc_post_push_constants { }; static_assert(sizeof(vk_op_dsv4_hc_post_push_constants) <= 128); +struct vk_op_flash_attn_union_push_constants { + uint32_t n_kv, n_kv_raw, n_batch, n_top_k, max_union, nbt1, max_words, pad_to; +}; +struct vk_op_flash_attn_gather_union_push_constants { + uint32_t n_kv, n_kv_raw, kv_c_max, nbk1, nbm1, n_batch; +}; struct vk_op_flash_attn_gather_push_constants { uint32_t n_kv, n_kv_raw, n_top_k, kv_c; uint32_t nbk1, nbk3, nbt1, nbt3, nbm1, nbm3, nem3, n_batch; @@ -4021,14 +4030,16 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_ } 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, - bool use_mask, bool use_mask_opt, bool use_logit_softcap, ggml_type k_type, ggml_type v_type) { + bool use_mask, bool use_mask_opt, bool use_logit_softcap, ggml_type k_type, ggml_type v_type, + bool use_dynamic_kv = false) { const bool old_amd_windows = device->vendor_id == VK_VENDOR_ID_AMD && device->driver_id == vk::DriverId::eAmdProprietary && (device->architecture == AMD_GCN || device->architecture == AMD_RDNA1 || device->architecture == AMD_RDNA2); uint32_t flags = (use_mask_opt ? 1 : 0) | (use_mask ? 2 : 0) | (use_logit_softcap ? 4 : 0) | - (old_amd_windows ? 8 : 0); + (old_amd_windows ? 8 : 0) | + (use_dynamic_kv ? 16 : 0); const uint32_t subgroup_size = params.disable_subgroups ? 0 : params.subgroup_size; @@ -4650,7 +4661,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16"; } - ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7, + ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 8, sizeof(vk_flash_attn_push_constants), {Br, 1, 1}, get_fa_spec_constants(fa.first), aligned ? Bc : 1, true, !fa_ds, !fa_ds ? fa_sgs : 0); @@ -4686,7 +4697,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { else { spv_data = flash_attn_f32_f16_f16acc_cm1_data; spv_size = flash_attn_f32_f16_f16acc_cm1_len; } name = aligned ? "flash_attn_f32_f16_aligned_cm1" : "flash_attn_f32_f16_cm1"; } - ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7, + ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 8, sizeof(vk_flash_attn_push_constants), {Br, 1, 1}, get_fa_spec_constants(fa.first), aligned ? Bc : 1, true, !fa_ds, !fa_ds ? fa_sgs : 0); @@ -4723,7 +4734,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { if (f32acc) { spv_data = flash_attn_f32_f16_cm2_data; spv_size = flash_attn_f32_f16_cm2_len; name = "flash_attn_f32_f16_f32acc_cm2"; } else { spv_data = flash_attn_f32_f16_f16acc_cm2_data; spv_size = flash_attn_f32_f16_f16acc_cm2_len; name = "flash_attn_f32_f16_f16acc_cm2"; } } - ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7, + ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 8, sizeof(vk_flash_attn_push_constants), {Br, 1, 1}, get_fa_spec_constants(fa.first), aligned ? Bc : 1, true, false, 0); } @@ -6107,6 +6118,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { "flash_attn_gather_f16", flash_attn_gather_f16_len, flash_attn_gather_f16_data, "main", 5, sizeof(vk_op_flash_attn_gather_push_constants), {1, 1, 1}, {}, 1, true, true, device->subgroup_size); + if (device->subgroup_arithmetic) { + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_union_f16, + "flash_attn_union_f16", flash_attn_union_f16_len, flash_attn_union_f16_data, "main", 3, + sizeof(vk_op_flash_attn_union_push_constants), {1, 1, 1}, {}, 1, true, true, + device->subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_union_f16, + "flash_attn_gather_union_f16", flash_attn_gather_union_f16_len, flash_attn_gather_union_f16_data, "main", 6, + sizeof(vk_op_flash_attn_gather_union_push_constants), {1, 1, 1}, {}, 1, true, true, + device->subgroup_size); + } } // DSv4 fused hyper-connection ops: plain f32 compute, no subgroup/coopmat requirements @@ -11151,6 +11172,56 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & return false; } + // ---- diagnostic: adjacent-token top-k overlap (GGML_VK_TOPK_OVERLAP=1) --------------- + // Decides whether a deduplicated union is worth building for the small-batch path. A + // speculative draft attends at adjacent POSITIONS, so overlap between adjacent prefill + // tokens is the same quantity and a plain deep prefill samples it thousands of times + // without needing a draft model. Reports |union| / (W * n_top_k) for window sizes W. + // Sampled, host-side, off by default; the readback would be far too costly otherwise. + { + static const char * ov_env = getenv("GGML_VK_TOPK_OVERLAP"); + if (ov_env && ov_env[0] == '1') { + static std::mutex ov_mu; + static uint64_t ov_seen = 0; + static std::map> ov_acc; // W -> {sum ratio, n} + std::lock_guard lock(ov_mu); + if ((ov_seen++ % 32) == 0) { // sample: this is a multi-MB readback + const uint32_t nb = (uint32_t) q->ne[1]; + const uint32_t tk = (uint32_t) top_k->ne[0]; + const int32_t rng = (int32_t) (k->ne[1] - n_kv_raw); + std::vector idx((size_t) nb * tk); + vk_subbuffer sb = ggml_vk_tensor_subbuffer(ctx, top_k); + ggml_vk_buffer_read(sb.buffer, sb.offset, idx.data(), idx.size() * sizeof(int32_t)); + for (uint32_t W : {2u, 4u, 8u}) { + if (nb < W) continue; + double sum = 0.0; uint64_t n = 0; + for (uint32_t t0 = 0; t0 + W <= nb; t0 += W) { // disjoint windows + std::unordered_set u; + uint64_t valid = 0; + for (uint32_t t = t0; t < t0 + W; ++t) { + for (uint32_t j = 0; j < tk; ++j) { + const int32_t v = idx[(size_t) t * tk + j]; + if (v >= 0 && v < rng) { u.insert(v); ++valid; } + } + } + if (valid) { sum += (double) u.size() / (double) valid; ++n; } + } + if (n) { auto & a = ov_acc[W]; a.first += sum; a.second += n; } + } + fprintf(stderr, "[topk-overlap] sample %llu (n_batch=%u, kv=%lld)\n", + (unsigned long long) ov_seen, nb, (long long) k->ne[1]); + for (const auto & e : ov_acc) { + const double r = e.second.first / (double) e.second.second; + const double now = (double) n_kv_raw + (double) e.first * tk; + const double dedup = (double) n_kv_raw + r * (double) e.first * tk; + fprintf(stderr, "[topk-overlap] W=%u union/selected=%.3f (overlap %.1f%%) " + "kv_c %.0f -> %.0f projected small-batch gain %.1f%%\n", + e.first, r, 100.0 * (1.0 - r), now, dedup, 100.0 * (1.0 - dedup / now)); + } + } + } + } + vk_op_flash_attn_top_k_push_constants pc = { (uint32_t) q->ne[1], (uint32_t) k->ne[1], (uint32_t) n_kv_raw, (uint32_t) top_k->ne[0], (uint32_t) q->ne[2], @@ -11270,7 +11341,7 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & }; ggml_vk_dispatch_pipeline(ctx, subctx, raw_pipeline, - {tile_q, k_buf, k_buf, tile_mask, tile_q, split_buf, tile_q}, + {tile_q, k_buf, k_buf, tile_mask, tile_q, split_buf, tile_q, tile_q /* dyn-KV: unused */}, raw_pc, {tile_n, NH, NS}); ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_RAW (sub-op)"); @@ -11305,8 +11376,10 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & struct vk_fa_compact_state { bool active = false; - uint32_t kv_c = 0; + bool dynamic_kv = false; // KV row count lives in kv_buf, not the push constant + uint32_t kv_c = 0; // upper bound; the real count is runtime when dynamic_kv uint32_t n_batch = 1; + vk_subbuffer kv_buf; vk_subbuffer kc_buf, mc_buf; }; @@ -11364,6 +11437,74 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ return false; } + // ---- deduplicated union (GGML_VK_FA_TOPK_UNION=1) ---------------------------------- + // Same compact layout, but one row per DISTINCT selected key instead of one block per + // token. Measured adjacent-token overlap on the real model is 60% at 4 tokens and 76% at + // 8, so the union is materially smaller. Its size is only known on the GPU, so the FA + // reads its KV bound from a buffer (the DYNAMIC_KV pipeline flag) rather than a push + // constant; padding the count to 256 keeps KV % Bc == 0 so the aligned variant still + // applies. Single stream only: the FA takes one KV for all streams. + static const char * union_env = getenv("GGML_VK_FA_TOPK_UNION"); + if (union_env && union_env[0] == '1' && q->ne[3] == 1 && n_batch > 1 && + ctx->device->pipeline_flash_attn_union_f16 && ctx->device->pipeline_flash_attn_gather_union_f16) { + const uint32_t max_union = (uint32_t) ((int64_t) n_batch * top_k->ne[0]); + // shared bitmap capacity in flash_attn_union.comp + const uint32_t max_words = 12288; + const uint32_t need_words = (uint32_t) (((k->ne[1] - n_kv_raw) + 31) / 32); + if (need_words > max_words) { + goto union_unavailable; // fall through to the per-token block form + } + const size_t ukc_sz = (size_t) kv_c * 512 * sizeof(ggml_fp16_t); + const size_t umc_sz = (size_t) n_batch * kv_c * sizeof(ggml_fp16_t); + const size_t ul_sz = (size_t) max_union * sizeof(uint32_t); + const size_t uc_sz = 2 * sizeof(uint32_t); + const size_t need = ukc_sz + umc_sz + ul_sz + uc_sz; + if (ctx->prealloc_size_y < need) { + ctx->prealloc_size_y = need; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_y_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + const vk_subbuffer kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); + const vk_subbuffer mc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, ukc_sz); + const vk_subbuffer ul_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, ukc_sz + umc_sz); + const vk_subbuffer uc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, ukc_sz + umc_sz + ul_sz); + + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, 1); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_gather_union_f16, 1); + + const vk_op_flash_attn_union_push_constants upc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, n_batch, (uint32_t) top_k->ne[0], max_union, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_union_f16, + { ggml_vk_tensor_subbuffer(ctx, top_k), ul_buf, uc_buf }, upc, { 1, 1, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + + const vk_op_flash_attn_gather_union_push_constants gpc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, kv_c, + (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), + n_batch, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_gather_union_f16, + { ggml_vk_tensor_subbuffer(ctx, k), ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), + kc_buf, mc_buf, uc_buf }, gpc, { kv_c, 1, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + ctx->prealloc_y_need_sync = true; + + st.active = true; + st.dynamic_kv = true; + st.kv_c = kv_c; + st.n_batch = n_batch; + st.kc_buf = kc_buf; + st.mc_buf = mc_buf; + st.kv_buf = uc_buf; + return true; + } +union_unavailable:; + const uint32_t ns = (uint32_t) q->ne[3]; const size_t kc_sz = (size_t) ns * kv_c * 512 * sizeof(ggml_fp16_t); const size_t mc_sz = (size_t) ns * n_batch * kv_c * sizeof(ggml_fp16_t); @@ -11400,6 +11541,53 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ ggml_vk_sync_buffers(ctx, subctx); ctx->prealloc_y_need_sync = true; + // ---- diagnostic: measure real top-k overlap between the query tokens ---------------- + // GGML_VK_TOPK_OVERLAP=1. Off by default and never touched on the hot path. Answers the + // one question that decides whether a deduplicated union is worth building: this path + // gathers n_batch*n_top_k rows, a union would gather |union| rows, and the op's cost is + // linear in that count (measured: 204.4 / 205.7 / 205.5 us per row at batch 2 / 4 / 8). + // Needs a real model - synthetic top-k selections say nothing about real overlap, which + // is exactly how the mul_mat_id large-tile probe misled once before. + static const char * overlap_env = getenv("GGML_VK_TOPK_OVERLAP"); + if (overlap_env && overlap_env[0] == '1' && n_batch > 1) { + static std::mutex ov_mutex; + static uint64_t ov_calls = 0, ov_selected = 0, ov_union = 0; + static std::map> ov_by_batch; // n_batch -> {selected, union} + const size_t n_idx = (size_t) n_batch * top_k->ne[0]; + std::vector idx(n_idx); + vk_subbuffer top_sb = ggml_vk_tensor_subbuffer(ctx, top_k); + ggml_vk_buffer_read(top_sb.buffer, top_sb.offset, idx.data(), n_idx * sizeof(int32_t)); + + const int32_t range = (int32_t) (k->ne[1] - n_kv_raw); + std::unordered_set uni; + uint64_t valid = 0; + for (size_t i = 0; i < n_idx; ++i) { + const int32_t v = idx[i]; + if (v >= 0 && v < range) { uni.insert(v); ++valid; } + } + std::lock_guard lock(ov_mutex); + ov_calls++; ov_selected += valid; ov_union += uni.size(); + auto & e = ov_by_batch[n_batch]; + e.first += valid; e.second += uni.size(); + if ((ov_calls % 256) == 0) { + fprintf(stderr, "[topk-overlap] calls=%llu selected=%llu union=%llu " + "union/selected=%.3f => a dedup union would gather %.1f%% fewer compressed rows\n", + (unsigned long long) ov_calls, (unsigned long long) ov_selected, + (unsigned long long) ov_union, + ov_selected ? (double) ov_union / (double) ov_selected : 0.0, + ov_selected ? 100.0 * (1.0 - (double) ov_union / (double) ov_selected) : 0.0); + for (const auto & kv : ov_by_batch) { + const double ratio = kv.second.first ? (double) kv.second.second / (double) kv.second.first : 0.0; + // projected op speedup uses the measured linear cost model on kv_c + const double now = (double) n_kv_raw + (double) kv.first * (double) top_k->ne[0]; + const double dedup = (double) n_kv_raw + ratio * (double) kv.first * (double) top_k->ne[0]; + fprintf(stderr, "[topk-overlap] n_batch=%u union/selected=%.3f " + "kv_c %.0f -> %.0f projected op gain %.1f%%\n", + kv.first, ratio, now, dedup, 100.0 * (1.0 - dedup / now)); + } + } + } + st.active = true; st.kv_c = kv_c; st.n_batch = n_batch; @@ -11611,7 +11799,8 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16 && (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256); vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc, - mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff); + mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff, + fa_compact.dynamic_kv); vk_pipeline pipeline = nullptr; @@ -11813,7 +12002,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx vk_subbuffer split_k_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, mask_opt_buf}, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, mask_opt_buf, fa_compact.dynamic_kv ? fa_compact.kv_buf : q_buf}, pc, { dispatch_x, workgroups_y, workgroups_z }); ggml_vk_sync_buffers(ctx, subctx); @@ -11828,7 +12017,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx workgroups_x *= pipeline->wg_denoms[0]; } ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf}, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf, fa_compact.dynamic_kv ? fa_compact.kv_buf : q_buf}, pc, { workgroups_x, workgroups_y, workgroups_z }); } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index e308c7214dd3..b562c5d78749 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -24,6 +24,11 @@ const bool USE_MASK_OPT = (Flags & 1) != 0; const bool MASK_ENABLE = (Flags & 2) != 0; const bool LOGIT_SOFTCAP = (Flags & 4) != 0; const bool OLD_AMD_WINDOWS = (Flags & 8) != 0; +// KV comes from a buffer instead of the push constant. Used by paths that compact K/V on the +// GPU, where the row count is only known after a dedup pass and so cannot be pushed. The +// workgroup counts derive from neq1/neq2/neq3 and never from KV, so no indirect dispatch is +// needed: only this loop bound changes. Folds away for every other pipeline. +const bool DYNAMIC_KV = (Flags & 16) != 0; // Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths const uint32_t HSK_pad = (HSK + 15) & ~15; @@ -81,6 +86,7 @@ layout (binding = 5) writeonly buffer O {D_TYPE data_o[];}; layout (binding = 5) writeonly buffer OV4 {D_TYPEV4 data_ov4[];}; layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; +layout (binding = 7) readonly buffer KVB {uint32_t data_kv_dyn[];}; #define MASK_OPT_ALL_NEG_INF 1 #define MASK_OPT_ALL_ZERO 2 @@ -150,7 +156,7 @@ bool partial_output; void init_indices() { N = p.N; - KV = p.KV; + KV = DYNAMIC_KV ? data_kv_dyn[0] : p.KV; gqa_ratio = p.gqa_ratio & 0xffff; split_k_num = p.k_num & 0xffff; output_k_num = p.k_num >> 16; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp new file mode 100644 index 000000000000..775b1bc96c62 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp @@ -0,0 +1,82 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +// Gathers K/V and the mask for the DeepSeek V4 small-batch decode path, using the deduplicated +// union produced by flash_attn_union.comp: rows [0, n_kv_raw) of the source, then one row per +// distinct selected compressed row, then padding. +// +// Simpler than the block-per-token form it replaces, because a union row appears exactly once: +// there is no owner to track and no cross-token masking to apply. Each token just reads its own +// mask value for the gathered source row, which is already -inf where that token did not select +// it, so the softmax still cannot double count. +// +// Row count is a runtime value in data_c[0]; rows past the union are zeroed K and -inf mask, +// which is softmax-neutral and keeps the padded tail harmless. + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 1) readonly buffer UBuf { uint data_u[]; }; +layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; +layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; +layout(binding = 5) readonly buffer CBuf { uint data_c[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_kv_raw; + uint kv_c_max; + uint nbk1; + uint nbm1; + uint n_batch; +} p; + +const uint HEAD_SIZE = 512; +const uint LANES = 64; + +void main() { + const uint row = gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationIndex; + + const uint kv_c = data_c[0]; // padded compact rows, the FA's runtime KV + const uint n_uni = data_c[1]; // unpadded union size + + if (row >= kv_c) { + return; + } + + uint src = p.n_kv; // sentinel: invalid + if (row < p.n_kv_raw) { + src = row; + } else if (row - p.n_kv_raw < n_uni) { + src = p.n_kv_raw + data_u[row - p.n_kv_raw]; + } + + const uint dst_base = row * HEAD_SIZE; + if (src < p.n_kv) { + const uint src_base = src * p.nbk1; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + data_kc[dst_base + tid + i * LANES] = data_k[src_base + tid + i * LANES]; + } + } else { + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + data_kc[dst_base + tid + i * LANES] = float16_t(0.0); + } + } + + // Compact mask is token-major [n_batch][kv_c]. The stride must be the RUNTIME kv_c, not + // kv_c_max: the FA derives m_row_len from KV, which is now that same runtime value, and a + // mismatch would step the mask by the wrong amount for every token past the first. The + // buffer is allocated for kv_c_max, so a smaller stride simply leaves a tail unused. + const float NEG_INF = uintBitsToFloat(0xff800000); + if (tid < p.n_batch) { + float mv = NEG_INF; + if (src < p.n_kv) { + mv = float(data_m[tid * p.nbm1 + src]); + } + data_mc[tid * kv_c + row] = float16_t(mv); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp new file mode 100644 index 000000000000..cd779e0a0f99 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp @@ -0,0 +1,122 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_basic : require + +// Builds the deduplicated union of the compressed rows selected by any of the n_batch query +// tokens, for the DeepSeek V4 small-batch decode path. +// +// Marks a bitmap from the top-k index lists, then compacts by scanning bitmap WORDS. The +// marking is O(n_batch * n_top_k), independent of context depth, and the compaction touches +// R/32 words instead of R rows. An earlier version scanned the mask row by row instead: that +// is simpler, but it made dedup cost scale with depth and measured 0.71x (ie a loss) at +// kv=133888, which defeats the purpose. Do not go back to it. +// +// Ascending source order falls out of the bitmap scan, so the result is deterministic. +// +// Emits the index list plus the PADDED compact row count. Padding to a multiple of pad_to +// (a multiple of every FA block width) is what lets flash-attention keep its "aligned" +// pipeline variant even though the row count is now a runtime value. +// +// One workgroup: the bitmap and the running offset both live in shared memory. + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer TBuf { int data_top[]; }; +layout(binding = 1) writeonly buffer UBuf { uint data_u[]; }; +layout(binding = 2) writeonly buffer CBuf { uint data_c[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_kv_raw; + uint n_batch; + uint n_top_k; + uint max_union; + uint nbt1; + uint max_words; // capacity of the shared bitmap, host-checked + uint pad_to; +} p; + +// 12288 words = 393216 compressed rows, ~48 KiB of shared memory. +const uint MAX_WORDS = 12288; + +shared uint bitmap[MAX_WORDS]; +shared uint wave_totals[16]; +shared uint base_sh; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint wave = gl_SubgroupID; + const uint nwave = gl_NumSubgroups; + const uint R = p.n_kv - p.n_kv_raw; + const uint words = min((R + 31) / 32, p.max_words); + + // phase 1: clear + for (uint w = tid; w < words; w += gl_WorkGroupSize.x) { + bitmap[w] = 0; + } + if (tid == 0) { + base_sh = 0; + } + barrier(); + + // phase 2: mark. Depth-independent: one pass over the top-k lists. + const uint n_cand = p.n_batch * p.n_top_k; + for (uint c = tid; c < n_cand; c += gl_WorkGroupSize.x) { + const uint t = c / p.n_top_k; + const uint j = c - t * p.n_top_k; + const int idx = data_top[t * p.nbt1 + j]; + if (idx >= 0 && uint(idx) < R) { + atomicOr(bitmap[uint(idx) >> 5], 1u << (uint(idx) & 31u)); + } + } + barrier(); + + // phase 3: compact. R/32 iterations, ascending, exact running offset in shared memory. + for (uint chunk = 0; chunk < words; chunk += gl_WorkGroupSize.x) { + const uint w = chunk + tid; + const uint bits = w < words ? bitmap[w] : 0u; + const uint cnt = bitCount(bits); + + const uint wave_off = subgroupExclusiveAdd(cnt); + const uint wave_tot = subgroupAdd(cnt); + if (lane == 0) { + wave_totals[wave] = wave_tot; + } + barrier(); + + uint prefix = 0; + for (uint i = 0; i < wave; ++i) { + prefix += wave_totals[i]; + } + uint total = 0; + for (uint i = 0; i < nwave; ++i) { + total += wave_totals[i]; + } + + uint slot = base_sh + prefix + wave_off; + uint rem = bits; + while (rem != 0) { + const uint b = findLSB(rem); + rem &= rem - 1; + if (slot < p.max_union) { + data_u[slot] = w * 32 + b; + } + ++slot; + } + barrier(); + if (tid == 0) { + base_sh += total; + } + barrier(); + } + + if (tid == 0) { + const uint u = min(base_sh, p.max_union); + const uint rows = p.n_kv_raw + u; + data_c[0] = ((rows + p.pad_to - 1) / p.pad_to) * p.pad_to; + data_c[1] = u; + } +} 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 8b0916a18a2f..f5b53b7fcf40 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -812,6 +812,8 @@ void process_shaders() { string_to_spv("flash_attn_top_k_cm_f16", "flash_attn_top_k_cm.comp", {}); #endif string_to_spv("flash_attn_gather_f16", "flash_attn_gather.comp", {}); + string_to_spv("flash_attn_union_f16", "flash_attn_union.comp", {}); + string_to_spv("flash_attn_gather_union_f16", "flash_attn_gather_union.comp", {}); string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {}); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4d3176fa90a2..df257f346d2f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7209,12 +7209,13 @@ struct test_flash_attn_ext_top_k : public test_case { const int64_t n_top_k; // selected keys per query token const bool sinks; const int64_t ns; // sequences (ne3); >1 exercises the split-K stream stride + const int64_t ov; // % of each token's picks shared with its neighbours (dedup-union realism) static constexpr int64_t hs = 512; // V4 CSA head size, K == V latent static constexpr int64_t nh = 64; // V4 CSA query heads (MQA) std::string vars() override { - return VARS_TO_STR6(kv, nb, n_kv_raw, n_top_k, sinks, ns); + return VARS_TO_STR7(kv, nb, n_kv_raw, n_top_k, sinks, ns, ov); } double max_nmse_err() override { @@ -7228,8 +7229,8 @@ struct test_flash_attn_ext_top_k : public test_case { return 2 * nh * nb * ns * (hs + hs) * (n_kv_raw + n_top_k); } - test_flash_attn_ext_top_k(int64_t kv = 768, int64_t nb = 8, int64_t n_kv_raw = 64, int64_t n_top_k = 128, bool sinks = false, int64_t ns = 1) - : kv(kv), nb(nb), n_kv_raw(n_kv_raw), n_top_k(n_top_k), sinks(sinks), ns(ns) {} + test_flash_attn_ext_top_k(int64_t kv = 768, int64_t nb = 8, int64_t n_kv_raw = 64, int64_t n_top_k = 128, bool sinks = false, int64_t ns = 1, int64_t ov = 0) + : kv(kv), nb(nb), n_kv_raw(n_kv_raw), n_top_k(n_top_k), sinks(sinks), ns(ns), ov(ov) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hs, nb, nh, ns); @@ -7295,7 +7296,10 @@ struct test_flash_attn_ext_top_k : public test_case { for (int64_t j = 0; j < n_top_k; ++j) { // offset the selection by the stream too, so a dropped stream stride // reads another sequence's keys and shows up as a mismatch - int32_t idx = (int32_t) ((j * range) / n_top_k + b + s * 7) % (int32_t) range; + const bool shared = (int64_t) j * 100 < n_top_k * ov; + int32_t idx = shared + ? (int32_t) ((j * range) / n_top_k + s * 7) % (int32_t) range + : (int32_t) ((j * range) / n_top_k + b + s * 7) % (int32_t) range; if (j == n_top_k - 1 && b == 0 && s == 0) { idx = -1; // exercise the ignore-invalid-index path } else { @@ -10759,6 +10763,14 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, nb, 2304, 512, false)); } } + // Same shapes with realistic adjacent-token overlap. Measured on DeepSeek-V4-Flash the + // real overlap is 60% over 4 adjacent tokens and 76% over 8; the default generator is + // near 0%, which would make a deduplicated union look worthless by construction. + for (int kv : { 35584, 133888 }) { + for (int nb : { 2, 4, 8 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, nb, 2304, 512, false, 1, 60)); + } + } return test_cases; } From 0ff9e372238a6e9e766d3d9c0f88841e169f63b9 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Fri, 14 Aug 2026 01:21:31 +0000 Subject: [PATCH 095/109] vulkan: gate the DeepSeek V4 small-batch union on the measured union size The union shrinks the compact set but the gate could not see it. kv_c is the worst case n_kv_raw + n_batch*n_top_k, so at batch 8 and 32k depth the gate priced 6400 rows against 11008 source rows, declined, and dense attention ran over the whole compressed KV -- even though those selections deduplicate to about 3300 rows. That is the batch ceiling the previous commit documented and did not lift. The union size is only known on the device, and the note on the previous commit assumed lifting the ceiling therefore needed a GPU-side fallback for an oversized union. It does not. The union is bounded by the source, so there is no dispatch to recover from; the missing piece was only a measurement. The union shader now writes its count into a small host-visible buffer and the host prices the next step from it. The read is deliberately unsynchronised and one graph stale -- overlap is a property of the model and the draft, not of one op -- and a wrong estimate costs part of one step, never correctness, because every allocation and dispatch bound is still the worst case. Estimates are held per batch size. A speculative decode varies the batch with the accept count, and the overlap itself varies with the batch (0.64 at 2 tokens, 0.40 at 4, 0.24 at 8), so a single slot would be invalidated on nearly every step. Each estimate tracks the latest measurement, lightly smoothed, and a declined step dispatches the scan in a new count_only mode that writes no index list. Both of those are deliberate and were measured the other way round first. Holding a decaying peak to stay conservative, and sampling the probe every 256 declines to stay cheap, together produced 1992 us on a shape the union runs in 898: the asymmetry actually runs the other way, because an estimate that is too high declines compaction and forgoes 2-3x for as long as it stays high, while one that is too low costs a single step and is corrected by the count that step produces. Sampling makes it worse still, since a decline is exactly the state in which the compact path stops refreshing the estimate, so the stale value stays latched for the whole sampling period. The probe is one workgroup against the ~2.2 ms dense op it rides along with, and it costs 0.2% of a declined step. The same measurement settles the reverse case. Where selections do not overlap, the union is the same size as the per-token blocks and its scan is pure cost, so those now stay on the per-token path rather than being admitted whenever the worst-case gate happened to allow them. test-backend-ops perf, medians of 3 launches counterbalanced A B B A A B, spreads at or under 0.4%, n_kv_raw=2304 n_top_k=512, at kv=11008 (~32k source) which is the shape where the gate used to decline. The fixture's ov is a per-token share, not the union/selected ratio the model was measured by: at nb tokens it yields (ov + (1-ov)*nb)/nb of the selections, so ov=86 is the setting that reproduces the 0.243 measured on DeepSeek-V4-Flash over 8 adjacent tokens, and ov=60 is deliberately more pessimistic than the model. batch overlap before after 8 ov=86 2215.2 us 704.8 us 3.14x 16 ov=86 2227.3 us 841.3 us 2.65x 8 ov=60 2214.5 us 897.9 us 2.47x 16 ov=60 2226.2 us 2231.2 us 1.00x union does not fit, declined 8 ov=0 2213.5 us 2219.0 us 1.00x nothing to deduplicate, declined The union cost is depth-independent as before, so the lifted cells now sit alongside the deeper ones: 897.9 / 900.7 / 906.6 us at batch 8 ov=60 across kv 11008 / 35584 / 133888. Every cell at kv 35584 and 133888 is 1.00x: this changes which shapes are admitted, not how the union performs once admitted. Both declines above are the estimate working rather than failing. At ov=60 and batch 16 the union really is 5888 rows against 11008 source rows and compaction would not pay; at ov=0 there is no overlap to exploit at all. Verified from the executed graph with GGML_VK_FA_UNION_STATS=1, added here, which reports the measured union/candidate ratio and the resulting decision rather than leaving engagement to be inferred from a timing. It reads 0.246 at ov=86 and batch 8, against the 0.243 measured on the model. 13310 FLASH_ATTN_EXT cases pass with the union off and on, including three new overlap cases at the shapes the gate now admits. End to end on a model that has no top-k tensor at all, so the compaction path is never reached and the change should be structurally inert: Qwen3-Coder-30B-A3B UD-Q4_K_XL, same counterbalanced order, medians of 3. pp512 1578.20 -> 1591.67 t/s and tg32 96.89 -> 97.28 at d0; pp512 700.64 -> 699.32 and tg32 53.92 -> 54.45 at d16384. All within 1%. The union path itself cannot be checked this way: it needs batch 2..63 decode, which only a speculative draft produces. Still measured against a synthetic fixture rather than real draft tokens: that needs a runnable target and draft pair, which this box cannot host. Suggested-by: Jaap Buurman (@Mushoz) Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 214 +++++++++++++++--- .../vulkan-shaders/flash_attn_union.comp | 14 +- tests/test-backend-ops.cpp | 21 +- 3 files changed, 216 insertions(+), 33 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 4ae3cc7ed80b..7467f9ec1f63 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1960,7 +1960,7 @@ struct vk_op_dsv4_hc_post_push_constants { static_assert(sizeof(vk_op_dsv4_hc_post_push_constants) <= 128); struct vk_op_flash_attn_union_push_constants { - uint32_t n_kv, n_kv_raw, n_batch, n_top_k, max_union, nbt1, max_words, pad_to; + uint32_t n_kv, n_kv_raw, n_batch, n_top_k, max_union, nbt1, max_words, pad_to, count_only; }; struct vk_op_flash_attn_gather_union_push_constants { uint32_t n_kv, n_kv_raw, kv_c_max, nbk1, nbm1, n_batch; @@ -2482,6 +2482,14 @@ struct ggml_backend_vk_context { uint64_t fa_dequant_gate_sz; bool fa_dequant_gate_fits; bool fa_dequant_gate_logged; + // DeepSeek V4 small-batch union: the compact row count is produced on the device, so the + // host prices compaction from the last count it wrote. See ggml_vk_fa_union_estimate. + // Per batch size, because the overlap depends on it (measured 0.64 at 2 tokens, 0.40 at 4, + // 0.24 at 8) and because a speculative decode varies the batch with the accept count, so a + // single slot would be invalidated on nearly every step. This path caps the batch at 64. + vk_buffer fa_union_stat; + float fa_union_est_ratio[64]; // union / candidates, decaying peak; 0 = unseeded + uint64_t fa_union_declines; vk::Fence fence, almost_ready_fence; bool submit_pending {}; bool almost_ready_fence_pending {}; @@ -7857,6 +7865,8 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { ctx->fa_dequant_gate_sz = 0; ctx->fa_dequant_gate_fits = false; ctx->fa_dequant_gate_logged = false; + memset(ctx->fa_union_est_ratio, 0, sizeof(ctx->fa_union_est_ratio)); + ctx->fa_union_declines = 0; // Fixed size of 1KB, for deterministic behavior ctx->prealloc_size_add_rms_partials = 1024; @@ -11383,6 +11393,80 @@ struct vk_fa_compact_state { vk_subbuffer kc_buf, mc_buf; }; +// Small host-visible buffer holding the last union count the device produced: +// [0] padded compact rows (also read by the gather and by the FA), [1] raw union size, +// [2] the candidate count it came from, [3] the batch it came from. Host-visible is a +// requirement rather than a preference here - the whole point is that the host can read it +// without submitting. +static bool ggml_vk_fa_union_stat_init(ggml_backend_vk_context * ctx) { + if (ctx->fa_union_stat) { + return ctx->fa_union_stat->ptr != nullptr; + } + try { + ctx->fa_union_stat = ggml_vk_create_buffer(ctx->device, 64, + {vk::MemoryPropertyFlagBits::eDeviceLocal | vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent}); + } catch (const vk::SystemError &) { + return false; + } + if (ctx->fa_union_stat->ptr == nullptr) { + return false; + } + memset(ctx->fa_union_stat->ptr, 0, 64); + return true; +} + +// Host-coherent memory still needs the write made available to the host domain. This is the +// only barrier in the path that names eHost, and it is cheap: nothing waits on it, it just +// lets the next graph's host-side read see the count this one produced. +static void ggml_vk_fa_union_stat_host_barrier(vk_context & subctx) { + subctx->s->buffer->buf.pipelineBarrier( + vk::PipelineStageFlagBits::eComputeShader, + vk::PipelineStageFlagBits::eHost, + {}, + { { vk::AccessFlagBits::eShaderWrite, vk::AccessFlagBits::eHostRead } }, + {}, {}); +} + +// Price the union without synchronising for it. The count read here belongs to whichever call +// last wrote the slot, which under a decode is the final flash-attention op of the previous +// graph - one token stale, and that is fine: selection overlap is a property of the model and +// the draft, not of an individual op, and a wrong estimate costs part of one step rather than +// correctness (the compact buffers are still sized for the worst case). +// +// Tracks the latest measurement, lightly smoothed. The asymmetry runs the other way from what +// a conservative estimator would assume: an estimate that is too HIGH declines compaction and +// forgoes 2-3x for as long as it stays high, while one that is too low costs a single step at +// roughly dense cost and is corrected by the count that step produces. An earlier version held +// a decaying peak instead and measured 1992 us where the union delivers 900, because a spell of +// genuinely low overlap pinned the estimate and 0.999 per read took hundreds of steps to relax. +// +// The words are read without ordering against the device write, so they can come from +// different calls. The sample is filed under the batch the device reported rather than the +// batch being priced, so a torn read costs one mispriced step for that batch and then +// corrects, which is the same failure the estimate already tolerates. +// +// Returns the padded compact row count to gate on, or 0 when this batch is unseeded. +static uint32_t ggml_vk_fa_union_estimate(ggml_backend_vk_context * ctx, uint32_t n_kv_raw, + uint32_t n_batch, uint32_t n_cand) { + const volatile uint32_t * stat = (const volatile uint32_t *) ctx->fa_union_stat->ptr; + const uint32_t u = stat[1]; + const uint32_t cand = stat[2]; + const uint32_t nb_obs = stat[3]; + + if (u > 0 && cand > 0 && nb_obs > 0 && nb_obs < 64) { + const float r = std::min(1.0f, (float) u / (float) cand); + float & e = ctx->fa_union_est_ratio[nb_obs]; + e = e > 0.0f ? 0.5f * r + 0.5f * e : r; + } + + const float ratio = ctx->fa_union_est_ratio[n_batch]; + if (ratio <= 0.0f) { + return 0; + } + return GGML_PAD(n_kv_raw + (uint32_t) ceilf(ratio * (float) n_cand), 256u); +} + // V4 sparse decode (gather-to-compact): the sparse prefill shader above gates on // q->ne[1] >= 64, so single-token decode otherwise attends densely over the whole // compressed KV, at a cost that grows with context. Instead, gather the active rows @@ -11422,20 +11506,17 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ } const uint32_t n_batch = (uint32_t) q->ne[1]; - // Every token gets its own top-k block, so the compact set is n_kv_raw + n_batch*n_top_k: - // independent of context depth, which is the point, but GROWING WITH BATCH. Attention work - // is then n_batch * (n_kv_raw + n_batch*n_top_k), i.e. quadratic in batch, against dense's - // n_batch * n_kv. Break-even is n_batch = (n_kv - n_kv_raw) / n_top_k, and the - // kv >= 2*kv_c gate below caps the useful batch at (n_kv/2 - n_kv_raw) / n_top_k -- - // about 6 at 32k depth, ~30 at 128k, batch-capped at 512k. Beyond that the gate declines - // and dense runs, so this can never be slower; it just stops helping. A deduplicated union - // would lift that ceiling wherever draft tokens select overlapping keys. - const uint32_t kv_c = GGML_PAD((uint32_t)(n_kv_raw + (int64_t) n_batch * top_k->ne[0]), 256u); - // the gather writes then re-reads ~the active bytes; dense reads the source KV once, - // so compaction only pays when the source is comfortably larger than the active set - if ((uint64_t) k->ne[1] < 2ull * kv_c) { - return false; - } + const uint32_t n_cand = (uint32_t) ((int64_t) n_batch * top_k->ne[0]); + // Worst case: every token gets its own top-k block, so the compact set is + // n_kv_raw + n_batch*n_top_k -- independent of context depth, which is the point, but + // GROWING WITH BATCH. Attention work is then n_batch * (n_kv_raw + n_batch*n_top_k), i.e. + // quadratic in batch, against dense's n_batch * n_kv. Break-even is + // n_batch = (n_kv - n_kv_raw) / n_top_k, and a kv >= 2*kv_c gate on this worst case caps + // the useful batch at (n_kv/2 - n_kv_raw) / n_top_k -- about 6 at 32k depth, ~30 at 128k. + // Beyond that the gate declines and dense runs, so this can never be slower; it just stops + // helping. The union below lifts that ceiling by gating on what the selections actually + // deduplicate to; kv_c remains the bound for every allocation and dispatch count. + const uint32_t kv_c = GGML_PAD((uint32_t) (n_kv_raw + (int64_t) n_cand), 256u); // ---- deduplicated union (GGML_VK_FA_TOPK_UNION=1) ---------------------------------- // Same compact layout, but one row per DISTINCT selected key instead of one block per @@ -11444,39 +11525,101 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // reads its KV bound from a buffer (the DYNAMIC_KV pipeline flag) rather than a push // constant; padding the count to 256 keeps KV % Bc == 0 so the aligned variant still // applies. Single stream only: the FA takes one KV for all streams. + // + // Gated on the ESTIMATED union rather than on kv_c, which is why this sits above the + // worst-case gate: at batch 8 and 32k depth the worst case is 6400 rows against 11008 + // source rows and would decline, while the union measures around 3300 and is well worth + // compacting. kv_c stays the worst case for every allocation and dispatch bound, so a + // wrong estimate is a slow step, never a wrong answer. + const uint32_t max_words = 12288; // shared bitmap capacity in flash_attn_union.comp + const bool bitmap_fits = (uint64_t) ((k->ne[1] - n_kv_raw) + 31) / 32 <= max_words; + static const char * union_env = getenv("GGML_VK_FA_TOPK_UNION"); - if (union_env && union_env[0] == '1' && q->ne[3] == 1 && n_batch > 1 && - ctx->device->pipeline_flash_attn_union_f16 && ctx->device->pipeline_flash_attn_gather_union_f16) { - const uint32_t max_union = (uint32_t) ((int64_t) n_batch * top_k->ne[0]); - // shared bitmap capacity in flash_attn_union.comp - const uint32_t max_words = 12288; - const uint32_t need_words = (uint32_t) (((k->ne[1] - n_kv_raw) + 31) / 32); - if (need_words > max_words) { - goto union_unavailable; // fall through to the per-token block form + if (union_env && union_env[0] == '1' && q->ne[3] == 1 && n_batch > 1 && bitmap_fits && + ctx->device->pipeline_flash_attn_union_f16 && ctx->device->pipeline_flash_attn_gather_union_f16 && + ggml_vk_fa_union_stat_init(ctx)) { + const uint32_t max_union = n_cand; + const uint32_t kv_c_est = ggml_vk_fa_union_estimate(ctx, (uint32_t) n_kv_raw, n_batch, n_cand); + // Two separate questions. Does the compact set fit under the gate at all, and does + // deduplicating actually shrink it: with no overlap to exploit the union is the same + // size as the per-token blocks and the scan is pure cost, measured at 1.2% of the op + // at 512k depth. The worst-case bound on the source keeps a collapse in overlap to + // roughly dense cost for the one step it takes the estimate to catch up. + const bool worth_it = kv_c_est != 0 && kv_c_est < kv_c && + (uint64_t) k->ne[1] >= 2ull * kv_c_est && + (uint64_t) k->ne[1] >= (uint64_t) kv_c; + + // GGML_VK_FA_UNION_STATS=1: what the gate actually decided and on what measurement. + // The alternative is inferring engagement from a timing, which is how a sparse path + // gets credited for a run it never took. + static const char * stats_env = getenv("GGML_VK_FA_UNION_STATS"); + if (stats_env && stats_env[0] == '1') { + static uint64_t calls = 0; + if ((calls++ % 256) == 0) { + fprintf(stderr, "[fa-union] n_kv=%lld n_kv_raw=%d n_batch=%u cand=%u " + "union/cand=%.3f kv_c %u -> est %u %s\n", + (long long) k->ne[1], n_kv_raw, n_batch, n_cand, (double) ctx->fa_union_est_ratio[n_batch], + kv_c, kv_c_est, worth_it ? "UNION" : "declined"); + } + } + + if (!worth_it) { + // Nothing is known about this batch shape, or the last count says compaction does + // not pay. Either way the count is what settles it, so produce one: the scan is a + // single workgroup and depth-independent, and count_only needs no index list. Dense + // runs this step; the next one decides on a measurement instead of a bound. + // + // Every declined step, not a sample: a decline is exactly the state in which the + // estimate stops being refreshed by the compact path, so sampling it leaves a stale + // estimate latched for as many steps as the sampling period. The dispatch is one + // workgroup against the ~2.2 ms dense op it is riding along with. + ctx->fa_union_declines++; + { + const vk_op_flash_attn_union_push_constants ppc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, n_batch, (uint32_t) top_k->ne[0], max_union, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, 1u, + }; + const vk_subbuffer stat_buf = ggml_vk_subbuffer(ctx, ctx->fa_union_stat); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, 1); + // The probe writes the same slot a taken union path reads back as its row + // count, and a graph can contain both: the estimate is refreshed from the + // device as the graph is recorded, so an op late in the graph can be admitted + // after an earlier one was declined. Order it explicitly rather than rely on + // the compact path's own sync, which this path does not go through. + ggml_vk_sync_buffers(ctx, subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_union_f16, + { ggml_vk_tensor_subbuffer(ctx, top_k), stat_buf, stat_buf }, ppc, { 1, 1, 1 }); + ggml_vk_fa_union_stat_host_barrier(subctx); + } + goto union_unavailable; } + const size_t ukc_sz = (size_t) kv_c * 512 * sizeof(ggml_fp16_t); const size_t umc_sz = (size_t) n_batch * kv_c * sizeof(ggml_fp16_t); const size_t ul_sz = (size_t) max_union * sizeof(uint32_t); - const size_t uc_sz = 2 * sizeof(uint32_t); - const size_t need = ukc_sz + umc_sz + ul_sz + uc_sz; + const size_t need = ukc_sz + umc_sz + ul_sz; if (ctx->prealloc_size_y < need) { ctx->prealloc_size_y = need; ggml_vk_preallocate_buffers(ctx, subctx); } - if (ctx->prealloc_y_need_sync) { - ggml_vk_sync_buffers(ctx, subctx); - } + // Unconditional, not gated on prealloc_y_need_sync: this also orders the count slot + // against a probe dispatched earlier in the same graph, which leaves that flag clear. + ggml_vk_sync_buffers(ctx, subctx); const vk_subbuffer kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); const vk_subbuffer mc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, ukc_sz); const vk_subbuffer ul_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, ukc_sz + umc_sz); - const vk_subbuffer uc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, ukc_sz + umc_sz + ul_sz); + // The count lives in the stat buffer rather than in prealloc_y so that the same write + // that the gather and the FA consume is also the one the host prices the next step from. + // Both are single-slot and rewritten by every layer, so the WAR hazard is unchanged: the + // prealloc_y sync above is a global barrier and orders the previous layer's read. + const vk_subbuffer uc_buf = ggml_vk_subbuffer(ctx, ctx->fa_union_stat); ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, 1); ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_gather_union_f16, 1); const vk_op_flash_attn_union_push_constants upc = { (uint32_t) k->ne[1], (uint32_t) n_kv_raw, n_batch, (uint32_t) top_k->ne[0], max_union, - (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, 0u, }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_union_f16, { ggml_vk_tensor_subbuffer(ctx, top_k), ul_buf, uc_buf }, upc, { 1, 1, 1 }); @@ -11492,6 +11635,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ { ggml_vk_tensor_subbuffer(ctx, k), ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), kc_buf, mc_buf, uc_buf }, gpc, { kv_c, 1, 1 }); ggml_vk_sync_buffers(ctx, subctx); + ggml_vk_fa_union_stat_host_barrier(subctx); ctx->prealloc_y_need_sync = true; st.active = true; @@ -11505,6 +11649,13 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ } union_unavailable:; + // Per-token blocks have no dedup, so this form really does cost kv_c: the gather writes + // then re-reads ~the active bytes while dense reads the source KV once, so compaction only + // pays when the source is comfortably larger than the active set. + if ((uint64_t) k->ne[1] < 2ull * kv_c) { + return false; + } + const uint32_t ns = (uint32_t) q->ne[3]; const size_t kc_sz = (size_t) ns * kv_c * 512 * sizeof(ggml_fp16_t); const size_t mc_sz = (size_t) ns * n_batch * kv_c * sizeof(ggml_fp16_t); @@ -17158,8 +17309,11 @@ static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) { ggml_vk_destroy_buffer(ctx->prealloc_y); ggml_vk_destroy_buffer(ctx->prealloc_split_k); ggml_vk_destroy_buffer(ctx->prealloc_add_rms_partials); + ggml_vk_destroy_buffer(ctx->fa_union_stat); ggml_vk_destroy_buffer(ctx->sync_staging); + memset(ctx->fa_union_est_ratio, 0, sizeof(ctx->fa_union_est_ratio)); + ctx->prealloc_y_last_pipeline_used = nullptr; ctx->prealloc_y_last_tensor_used = nullptr; ctx->prealloc_y_last_k_padded = false; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp index cd779e0a0f99..513375709300 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp @@ -19,6 +19,15 @@ // (a multiple of every FA block width) is what lets flash-attention keep its "aligned" // pipeline variant even though the row count is now a runtime value. // +// Also emits the raw union size, the candidate count and the batch it came from. The host +// reads those back to decide whether compaction is worth it at all, and files the sample under +// the batch reported here rather than the one it is pricing: it cannot tell which call last +// wrote the slot, and the overlap it is measuring depends on the batch. +// +// count_only skips the list writes. The scan still has to run to produce the count, but with +// nothing to write the index buffer need not exist, which is what lets the host price the +// union on a call it is not going to compact. +// // One workgroup: the bitmap and the running offset both live in shared memory. layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; @@ -36,6 +45,7 @@ layout(push_constant) uniform Parameters { uint nbt1; uint max_words; // capacity of the shared bitmap, host-checked uint pad_to; + uint count_only; // 1: produce the count, write no index list } p; // 12288 words = 393216 compressed rows, ~48 KiB of shared memory. @@ -97,7 +107,7 @@ void main() { } uint slot = base_sh + prefix + wave_off; - uint rem = bits; + uint rem = p.count_only != 0 ? 0u : bits; while (rem != 0) { const uint b = findLSB(rem); rem &= rem - 1; @@ -118,5 +128,7 @@ void main() { const uint rows = p.n_kv_raw + u; data_c[0] = ((rows + p.pad_to - 1) / p.pad_to) * p.pad_to; data_c[1] = u; + data_c[2] = n_cand; + data_c[3] = p.n_batch; } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index df257f346d2f..cf6bb39ebbf2 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10313,6 +10313,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 8, 1024, 512, true)); test_cases.emplace_back(new test_flash_attn_ext_top_k(32768, 16, 2304, 512, false)); test_cases.emplace_back(new test_flash_attn_ext_top_k(65536, 63, 2304, 512, false)); + // overlapping selections at the shapes where the compaction gate is tightest: with a + // deduplicated union these are admitted on the estimated union size rather than the + // worst case, so they cover the estimator's gate as well as the union itself. + test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 8, 2304, 512, false, 1, 60)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 16, 2304, 512, false, 1, 60)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 16, 2304, 512, false, 1, 86)); test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 4, 1024, 512, false, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, true, 2)); @@ -10766,11 +10772,22 @@ static std::vector> make_test_cases_perf() { // Same shapes with realistic adjacent-token overlap. Measured on DeepSeek-V4-Flash the // real overlap is 60% over 4 adjacent tokens and 76% over 8; the default generator is // near 0%, which would make a deduplicated union look worthless by construction. - for (int kv : { 35584, 133888 }) { - for (int nb : { 2, 4, 8 }) { + // kv=11008 (~32k source) and nb=16 are where the compaction gate is tightest: the + // worst-case compact set 2304 + nb*512 crosses kv/2 at nb=6, so those cells measure + // whether the gate can be opened by the union rather than by the worst case. + for (int kv : { 11008, 35584, 133888 }) { + for (int nb : { 2, 4, 8, 16 }) { test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, nb, 2304, 512, false, 1, 60)); } } + // ov is a per-token share, not the union/selected ratio the model was measured by: at nb + // tokens it gives a union of (ov + (1-ov)*nb)/nb of the selections, so ov=60 is 0.475 at + // nb=8 where the model measured 0.243. ov=86 is the setting that reproduces the model, and + // at kv=11008 it is the difference between a union that fits under the gate and one that + // does not. + for (int nb : { 8, 16 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, nb, 2304, 512, false, 1, 86)); + } return test_cases; } From ce9fef7d87dc2dfc3244f65fdf232c49e188a51a Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Fri, 14 Aug 2026 02:27:30 +0000 Subject: [PATCH 096/109] vulkan: measure the V4 union on real draft tokens, and correct the fixture calibration The previous commit calibrated its fixture against a PROXY, because a proxy was all that was reachable at the time: adjacent PREFILL tokens standing in for draft tokens. With the draft now drivable, the proxy turns out to have been optimistic, and the correction runs against that commit's own message. DeepSeek-V4-Flash UD-IQ3_XXS with the DSpark draft, Vulkan attention and CPU experts, 25538 tokens of context. The gate reports union/candidate 0.560 and 0.609 on the target layers and 0.484 to 0.651 on a second shape (n_kv_raw 256, most consistent with the draft's own sparse attention, also GPU-resident). Mean about 0.56 at batch 4, against 0.397 from the proxy at a window of 4. Draft tokens diverge from each other more than adjacent prefill tokens do, which is what a draft exploring rather than committed text should look like. The fixture's ov maps to (ov + (1-ov)*nb)/nb, so ov=60 predicts 0.55 at batch 4 against 0.56 measured: ov=60 reproduces real drafting almost exactly. The previous commit says ov=86 is the setting that reproduces the model and ov=60 is deliberately pessimistic. That is backwards for real draft tokens, and the realistic headline at batch 8 is 2.47x rather than 3.14x. The kernel numbers in that table are unchanged and still correct for the overlap each column states; only which column describes reality has moved. Batch 8 overlap itself is still not measured on real drafts, for the reason below, so 2.47x is an extrapolation from batch 4. The path does engage: 56 of 57 sampled gate decisions took the union, and the single decline is the unseeded first call, which then probes. That is the designed bootstrap, observed on the real model rather than a fixture. Scope, stated plainly because the previous commit implies more: DSpark drafts 3 tokens, so the batch is 4 and --spec-draft-n-max does not raise it, an MTP head having a fixed width. At batch 4 and this depth kv_c is 4352 against n_kv 8704, so the OLD worst-case gate admits too and the gate change makes no difference there; the smaller compact set comes from the union commit. For this target and draft the gate change buys a lower depth threshold instead: roughly 21.0k tokens rather than 25.5k on the target layers, and 11.8k rather than 17.7k on the second shape. The 2.47x and 2.65x cells need batch 8 or 16, which this pair does not generate. GGML_VK_FA_UNION_STATS takes a period now rather than being fixed at 256, and reports running totals. The fixed period is why this took two runs: a 64-token speculative decode over 25k of context emitted exactly ONE line, which reported only that the first call was unseeded. A diagnostic whose sampling rate can hide the thing it exists to measure is not a diagnostic. Measurements and provenance: ~/strix-results/derived-dsv4-union-gate-20260814.md section 9. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 7467f9ec1f63..13ae9340e9b5 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11549,17 +11549,25 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ (uint64_t) k->ne[1] >= 2ull * kv_c_est && (uint64_t) k->ne[1] >= (uint64_t) kv_c; - // GGML_VK_FA_UNION_STATS=1: what the gate actually decided and on what measurement. - // The alternative is inferring engagement from a timing, which is how a sparse path - // gets credited for a run it never took. + // GGML_VK_FA_UNION_STATS=N: report every Nth call (N=1 means every call) what the gate + // decided and on what measurement. The alternative is inferring engagement from a + // timing, which is how a sparse path gets credited for a run it never took. + // + // The period is a parameter because a fixed one is a way to miss the answer: a 64-token + // speculative decode over 25k of context produced ONE line at a period of 256, which + // said only that the first call was unseeded. Running totals rather than instants, so a + // single late line still reports whether the path engaged. static const char * stats_env = getenv("GGML_VK_FA_UNION_STATS"); - if (stats_env && stats_env[0] == '1') { - static uint64_t calls = 0; - if ((calls++ % 256) == 0) { + if (stats_env && stats_env[0] != '\0' && stats_env[0] != '0') { + static uint64_t calls = 0, taken = 0; + const uint64_t period = std::max(1ull, (unsigned long long) atoll(stats_env)); + taken += worth_it ? 1 : 0; + if ((calls++ % period) == 0) { fprintf(stderr, "[fa-union] n_kv=%lld n_kv_raw=%d n_batch=%u cand=%u " - "union/cand=%.3f kv_c %u -> est %u %s\n", + "union/cand=%.3f kv_c %u -> est %u %s (%llu/%llu taken)\n", (long long) k->ne[1], n_kv_raw, n_batch, n_cand, (double) ctx->fa_union_est_ratio[n_batch], - kv_c, kv_c_est, worth_it ? "UNION" : "declined"); + kv_c, kv_c_est, worth_it ? "UNION" : "declined", + (unsigned long long) taken, (unsigned long long) calls); } } From 7dfe8c8961dfa9091363f9ac91ff6e060a4b5824 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Fri, 14 Aug 2026 03:11:13 +0000 Subject: [PATCH 097/109] vulkan: default the DeepSeek V4 small-batch union on Every other gate in this family is on by default with =0 to disable: GGML_VK_FA_TOPK, GGML_VK_FA_TOPK_CM, GGML_VK_FA_TOPK_SPLIT. GGML_VK_FA_TOPK_UNION=1 was the odd one out because it began as a prototype, and it is no longer one. What makes it safe to default is that the path declines itself rather than needing a user to know when to avoid it. Where the selections do not deduplicate, the estimate says so and the per-token form runs instead; where the compact set would not fit under the gate, dense runs. Both were measured, at 1.00x, rather than assumed. Kept as one commit of its own rather than folded into the gate change, because it is a policy change and not a mechanism one: reverting these four lines restores opt-in behaviour without touching the gate logic, which is what a bisect would want if a regression turns up. 13310 FLASH_ATTN_EXT cases pass with the variable UNSET, which is the configuration this changes and the one that ships. That is the same code path as the previously validated GGML_VK_FA_TOPK_UNION=1, but it was run again rather than argued from equivalence. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 13ae9340e9b5..99d6509fc27a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11518,10 +11518,10 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // deduplicate to; kv_c remains the bound for every allocation and dispatch count. const uint32_t kv_c = GGML_PAD((uint32_t) (n_kv_raw + (int64_t) n_cand), 256u); - // ---- deduplicated union (GGML_VK_FA_TOPK_UNION=1) ---------------------------------- + // ---- deduplicated union (default on, GGML_VK_FA_TOPK_UNION=0 disables) -------------- // Same compact layout, but one row per DISTINCT selected key instead of one block per - // token. Measured adjacent-token overlap on the real model is 60% at 4 tokens and 76% at - // 8, so the union is materially smaller. Its size is only known on the GPU, so the FA + // token. Measured on real draft tokens the union is 0.56 of the selections at batch 4, + // so it is materially smaller. Its size is only known on the GPU, so the FA // reads its KV bound from a buffer (the DYNAMIC_KV pipeline flag) rather than a push // constant; padding the count to 256 keeps KV % Bc == 0 so the aligned variant still // applies. Single stream only: the FA takes one KV for all streams. @@ -11535,7 +11535,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ const bool bitmap_fits = (uint64_t) ((k->ne[1] - n_kv_raw) + 31) / 32 <= max_words; static const char * union_env = getenv("GGML_VK_FA_TOPK_UNION"); - if (union_env && union_env[0] == '1' && q->ne[3] == 1 && n_batch > 1 && bitmap_fits && + if ((!union_env || union_env[0] != '0') && q->ne[3] == 1 && n_batch > 1 && bitmap_fits && ctx->device->pipeline_flash_attn_union_f16 && ctx->device->pipeline_flash_attn_gather_union_f16 && ggml_vk_fa_union_stat_init(ctx)) { const uint32_t max_union = n_cand; From 8851c3705f0fe3a368dd6d04243cff68e22ff2d7 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Fri, 14 Aug 2026 06:27:04 +0000 Subject: [PATCH 098/109] vulkan: let the DeepSeek V4 small-batch gather serve quantised K/V The gather rejected anything but f16 K/V, so a DSv4 run with -ctk q8_0 -ctv q8_0 took the dense fallback for every batch 2..63 decode - the whole small-batch path, silently off. Two community test sets were run that way before anyone noticed, which is what prompted this. The restriction was incidental rather than essential. The gather RELOCATES rows; it never reads a value out of one. Addressing K as raw 4-byte words instead of f16 elements makes the same shader serve any type whose row is a whole number of words, which every block-quantised type here satisfies at head size 512. The union shader needed no change at all - it only ever touched top_k indices. Zeroing an unused row now writes zero BYTES. That decodes to zero for the block-quantised types as well as for f16, because a zero scale zeroes the block; those rows are -inf in the mask regardless, so zeroing only keeps a garbage dot product out of the softmax as a NaN. The type gate asks ggml_vk_fa_kv_native rather than just "is it quantised". While the compact scratch is active the dequant/contiguize pass is disabled by construction, and the assert guarding that combination aborts rather than falling back, so admitting a type flash-attention has no native shader for would crash instead of degrade. test-backend-ops perf, kv=11008 n_kv_raw=2304 n_top_k=512 ov=60, gather off vs on: batch f16 dense -> gathered q8_0 dense -> gathered 2 2192 -> 645 us 3.40x 3917 -> 1110 us 3.53x 4 2203 -> 726 us 3.03x 3924 -> 1257 us 3.12x 8 2213 -> 910 us 2.43x 3936 -> 1567 us 2.51x 16 2219 -> 2232 us declined 3930 -> 3952 us declined Two things worth stating because they contradict what I expected going in. Compaction does NOT pay more on quantised K/V. The reasoning was that a gathered set is contiguous and would dodge the strided-read and channel-aliasing taxes that hurt quantised reads more; the measured ratios differ by 3 to 4%, inside noise. It pays the same proportion from a worse starting point. That worse starting point is the real result: q8_0 attention is about 1.78x SLOWER than f16 here, dense (3917 vs 2192) and gathered (1110 vs 645) alike. Halving the bytes does not pay for the dequant work in the inner loop at this head size. So this change makes q8_0 K/V survivable for DSv4 rather than advisable - f16 is still the faster cache by a wide margin, and that inverts the usual q8_0-KV-buys-speed rule of thumb for this model. 13318 FLASH_ATTN_EXT cases pass, including 8 new q8_0/q4_0 top-k cases at the shapes a DSv4 decode actually hits. test_flash_attn_ext_top_k takes a type_K parameter for them. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 55 ++++++++++++++----- .../vulkan-shaders/flash_attn_gather.comp | 25 +++++---- .../flash_attn_gather_union.comp | 25 ++++++--- tests/test-backend-ops.cpp | 24 ++++++-- 4 files changed, 91 insertions(+), 38 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 99d6509fc27a..46f312579880 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1962,12 +1962,14 @@ static_assert(sizeof(vk_op_dsv4_hc_post_push_constants) <= 128); struct vk_op_flash_attn_union_push_constants { uint32_t n_kv, n_kv_raw, n_batch, n_top_k, max_union, nbt1, max_words, pad_to, count_only; }; +// nbk1/nbk3 are in 4-byte WORDS, not elements: the gather relocates K rows verbatim and never +// interprets what is in them, so it works for any type whose row is a whole number of words. struct vk_op_flash_attn_gather_union_push_constants { - uint32_t n_kv, n_kv_raw, kv_c_max, nbk1, nbm1, n_batch; + uint32_t n_kv, n_kv_raw, kv_c_max, nbk1, nbm1, n_batch, row_words; }; struct vk_op_flash_attn_gather_push_constants { uint32_t n_kv, n_kv_raw, n_top_k, kv_c; - uint32_t nbk1, nbk3, nbt1, nbt3, nbm1, nbm3, nem3, n_batch; + uint32_t nbk1, nbk3, nbt1, nbt3, nbm1, nbm3, nem3, n_batch, row_words; }; static_assert(sizeof(vk_op_flash_attn_gather_push_constants) <= 128); @@ -11389,6 +11391,8 @@ struct vk_fa_compact_state { bool dynamic_kv = false; // KV row count lives in kv_buf, not the push constant uint32_t kv_c = 0; // upper bound; the real count is runtime when dynamic_kv uint32_t n_batch = 1; + uint32_t row_bytes = 0; // bytes per compact K row; K may be quantised + uint32_t row_elems = 0; // K row stride in ELEMENTS/blocks, for the FA push constant vk_subbuffer kv_buf; vk_subbuffer kc_buf, mc_buf; }; @@ -11479,10 +11483,22 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ const ggml_tensor * mask, ggml_tensor * dst, vk_fa_compact_state & st) { const ggml_tensor * top_k = dst->src[5]; static const char * gather_env = getenv("GGML_VK_FA_TOPK_GATHER"); + // K/V type: the gather relocates rows verbatim and never reads a value out of them, so it + // does not have to be f16 - it has to be a type whose row is a whole number of 4-byte words, + // and one flash-attention has a NATIVE shader for. The dequant/contiguize pass is disabled + // while this scratch is active and asserts if it turns out to be needed, so admitting a + // non-native type here would abort rather than fall back. + const bool kv_word_addressable = + k->type == v->type && + ggml_vk_fa_kv_native(k->type, ctx->device->coopmat2) && + k->ne[0] % ggml_blck_size(k->type) == 0 && + ggml_row_size(k->type, k->ne[0]) % 4 == 0 && + k->nb[1] % 4 == 0 && k->nb[3] % 4 == 0; + if ((gather_env && gather_env[0] == '0') || !top_k || !ctx->device->pipeline_flash_attn_gather_f16 || q->ne[1] < 1 || q->ne[1] >= 64 || // 1..63: >=64 goes to the sparse prefill path - q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || + q->type != GGML_TYPE_F32 || !kv_word_addressable || !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || q->ne[0] != 512 || k->ne[0] != 512 || v->ne[0] != 512 || q->ne[2] != 64 || k->ne[2] != 1 || v->ne[2] != 1 || @@ -11492,6 +11508,9 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ return false; } + const uint32_t k_row_bytes = (uint32_t) ggml_row_size(k->type, k->ne[0]); + const uint32_t k_row_words = k_row_bytes / 4; + float max_bias = 0.0f; float logit_softcap = 0.0f; memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); @@ -11602,7 +11621,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ goto union_unavailable; } - const size_t ukc_sz = (size_t) kv_c * 512 * sizeof(ggml_fp16_t); + const size_t ukc_sz = (size_t) kv_c * k_row_bytes; const size_t umc_sz = (size_t) n_batch * kv_c * sizeof(ggml_fp16_t); const size_t ul_sz = (size_t) max_union * sizeof(uint32_t); const size_t need = ukc_sz + umc_sz + ul_sz; @@ -11635,9 +11654,9 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ const vk_op_flash_attn_gather_union_push_constants gpc = { (uint32_t) k->ne[1], (uint32_t) n_kv_raw, kv_c, - (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (k->nb[1] / 4), (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), - n_batch, + n_batch, k_row_words, }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_gather_union_f16, { ggml_vk_tensor_subbuffer(ctx, k), ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), @@ -11650,6 +11669,8 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ st.dynamic_kv = true; st.kv_c = kv_c; st.n_batch = n_batch; + st.row_bytes = k_row_bytes; + st.row_elems = (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); st.kc_buf = kc_buf; st.mc_buf = mc_buf; st.kv_buf = uc_buf; @@ -11665,7 +11686,7 @@ union_unavailable:; } const uint32_t ns = (uint32_t) q->ne[3]; - const size_t kc_sz = (size_t) ns * kv_c * 512 * sizeof(ggml_fp16_t); + const size_t kc_sz = (size_t) ns * kv_c * k_row_bytes; const size_t mc_sz = (size_t) ns * n_batch * kv_c * sizeof(ggml_fp16_t); if (ctx->prealloc_size_y < kc_sz + mc_sz) { @@ -11681,14 +11702,14 @@ union_unavailable:; const vk_op_flash_attn_gather_push_constants pc = { (uint32_t) k->ne[1], (uint32_t) n_kv_raw, (uint32_t) top_k->ne[0], kv_c, - (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), - (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (k->nb[1] / 4), + (uint32_t) (k->nb[3] / 4), (uint32_t) (top_k->nb[1] / sizeof(int32_t)), (uint32_t) (top_k->nb[3] / sizeof(int32_t)), (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (mask->nb[3] / sizeof(ggml_fp16_t)), (uint32_t) mask->ne[3], - n_batch, + n_batch, k_row_words, }; st.kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); @@ -11750,6 +11771,8 @@ union_unavailable:; st.active = true; st.kv_c = kv_c; st.n_batch = n_batch; + st.row_bytes = k_row_bytes; + st.row_elems = (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); return true; } @@ -11909,8 +11932,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type)); uint32_t v_stride = (uint32_t)(nbv1 / ggml_type_size(v->type)); if (fa_compact.active) { - k_stride = 512; - v_stride = 512; + // rows are tightly packed in the compact scratch; for a quantised K this is the block + // count per row, which is what nbk1 / ggml_type_size would have given for the source + k_stride = fa_compact.row_elems; + v_stride = fa_compact.row_elems; } // For F32, the shader treats it as a block of size 4 (for vec4 loads) @@ -12123,9 +12148,9 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx ggml_vk_sync_buffers(ctx, subctx); } - // compact scratch layout: [512, kv_c, 1, ns] f16, tightly packed - const uint32_t eff_nbk2 = fa_compact.active ? fa_compact.kv_c * 512 * (uint32_t)sizeof(ggml_fp16_t) : nbk2_eff; - const uint32_t eff_nbk3 = fa_compact.active ? fa_compact.kv_c * 512 * (uint32_t)sizeof(ggml_fp16_t) : nbk3_eff; + // compact scratch layout: [512, kv_c, 1, ns] tightly packed, in K's own type + const uint32_t eff_nbk2 = fa_compact.active ? fa_compact.kv_c * fa_compact.row_bytes : nbk2_eff; + const uint32_t eff_nbk3 = fa_compact.active ? fa_compact.kv_c * fa_compact.row_bytes : nbk3_eff; const uint32_t eff_nbv2 = fa_compact.active ? eff_nbk2 : nbv2_eff; const uint32_t eff_nbv3 = fa_compact.active ? eff_nbk3 : nbv3_eff; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp index 3c9ac68e2c69..6dd24d31b517 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp @@ -19,10 +19,12 @@ layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; -layout(binding = 0) readonly buffer KBuf { float16_t data_k[]; }; +// K is addressed as raw 4-byte WORDS, never as values: a gather relocates rows and does not +// need to know whether they hold f16 elements or quantised blocks. Only the MASK is typed. +layout(binding = 0) readonly buffer KBuf { uint data_k[]; }; layout(binding = 1) readonly buffer TopBuf { int data_top[]; }; layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; -layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; +layout(binding = 3) writeonly buffer KcBuf { uint data_kc[]; }; layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; layout(push_constant) uniform Parameters { @@ -30,17 +32,17 @@ layout(push_constant) uniform Parameters { uint n_kv_raw; // dense prefix length uint n_top_k; // selected rows for the (single) query token uint kv_c; // padded compact row count == dispatch row range - uint nbk1; // K source row stride, elements - uint nbk3; // K source stream stride, elements + uint nbk1; // K source row stride, WORDS + uint nbk3; // K source stream stride, WORDS uint nbt1; // top_k row (per query token) stride, elements uint nbt3; // top_k stream stride, elements uint nbm1; // mask source row (per query token) stride, elements uint nbm3; // mask source stream stride, elements uint nem3; // mask ne[3], for stream broadcast uint n_batch; // query tokens sharing this gather; <= LANES + uint row_words; // bytes per K row / 4 } p; -const uint HEAD_SIZE = 512; const uint LANES = 64; void main() { @@ -68,15 +70,18 @@ void main() { } } - const uint dst_base = (stream * p.kv_c + row) * HEAD_SIZE; + // Zeroing writes zero BYTES, which decode to zero for every block-quantised type here (a + // zero scale zeroes the block) as well as for f16. The row is -inf in the mask either way; + // zeroing only keeps a garbage dot product from reaching the softmax as a NaN. + const uint dst_base = (stream * p.kv_c + row) * p.row_words; if (src < p.n_kv) { const uint src_base = stream * p.nbk3 + src * p.nbk1; - [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { - data_kc[dst_base + tid + i * LANES] = data_k[src_base + tid + i * LANES]; + for (uint i = tid; i < p.row_words; i += LANES) { + data_kc[dst_base + i] = data_k[src_base + i]; } } else { - [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { - data_kc[dst_base + tid + i * LANES] = float16_t(0.0); + for (uint i = tid; i < p.row_words; i += LANES) { + data_kc[dst_base + i] = 0u; } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp index 775b1bc96c62..4669e0d1c75e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp @@ -18,10 +18,14 @@ layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; -layout(binding = 0) readonly buffer KBuf { float16_t data_k[]; }; +// K is addressed as raw 4-byte WORDS, never as values. A gather relocates rows; it does not +// need to know whether they hold f16 elements or quantised blocks, so the same shader serves +// every K type whose row is a whole number of words. Only the MASK is typed, and a mask is +// always f16. +layout(binding = 0) readonly buffer KBuf { uint data_k[]; }; layout(binding = 1) readonly buffer UBuf { uint data_u[]; }; layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; -layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; +layout(binding = 3) writeonly buffer KcBuf { uint data_kc[]; }; layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; layout(binding = 5) readonly buffer CBuf { uint data_c[]; }; @@ -29,12 +33,12 @@ layout(push_constant) uniform Parameters { uint n_kv; uint n_kv_raw; uint kv_c_max; - uint nbk1; + uint nbk1; // K source row stride, WORDS uint nbm1; uint n_batch; + uint row_words; // bytes per K row / 4 } p; -const uint HEAD_SIZE = 512; const uint LANES = 64; void main() { @@ -55,15 +59,18 @@ void main() { src = p.n_kv_raw + data_u[row - p.n_kv_raw]; } - const uint dst_base = row * HEAD_SIZE; + // Zeroing an unused row writes zero BYTES, which decode to zero for every block-quantised + // type here (a zero scale zeroes the block) as well as for f16. The row is -inf in the mask + // either way; zeroing only keeps a garbage dot product from reaching the softmax as a NaN. + const uint dst_base = row * p.row_words; if (src < p.n_kv) { const uint src_base = src * p.nbk1; - [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { - data_kc[dst_base + tid + i * LANES] = data_k[src_base + tid + i * LANES]; + for (uint i = tid; i < p.row_words; i += LANES) { + data_kc[dst_base + i] = data_k[src_base + i]; } } else { - [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { - data_kc[dst_base + tid + i * LANES] = float16_t(0.0); + for (uint i = tid; i < p.row_words; i += LANES) { + data_kc[dst_base + i] = 0u; } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index cf6bb39ebbf2..f01bf2ca7d14 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7210,12 +7210,13 @@ struct test_flash_attn_ext_top_k : public test_case { const bool sinks; const int64_t ns; // sequences (ne3); >1 exercises the split-K stream stride const int64_t ov; // % of each token's picks shared with its neighbours (dedup-union realism) + const ggml_type type_K; // K/V cache type; V is the same tensor, so one type covers both static constexpr int64_t hs = 512; // V4 CSA head size, K == V latent static constexpr int64_t nh = 64; // V4 CSA query heads (MQA) std::string vars() override { - return VARS_TO_STR7(kv, nb, n_kv_raw, n_top_k, sinks, ns, ov); + return VARS_TO_STR8(kv, nb, n_kv_raw, n_top_k, sinks, ns, ov, type_K); } double max_nmse_err() override { @@ -7229,14 +7230,15 @@ struct test_flash_attn_ext_top_k : public test_case { return 2 * nh * nb * ns * (hs + hs) * (n_kv_raw + n_top_k); } - test_flash_attn_ext_top_k(int64_t kv = 768, int64_t nb = 8, int64_t n_kv_raw = 64, int64_t n_top_k = 128, bool sinks = false, int64_t ns = 1, int64_t ov = 0) - : kv(kv), nb(nb), n_kv_raw(n_kv_raw), n_top_k(n_top_k), sinks(sinks), ns(ns), ov(ov) {} + test_flash_attn_ext_top_k(int64_t kv = 768, int64_t nb = 8, int64_t n_kv_raw = 64, int64_t n_top_k = 128, bool sinks = false, int64_t ns = 1, int64_t ov = 0, + ggml_type type_K = GGML_TYPE_F16) + : kv(kv), nb(nb), n_kv_raw(n_kv_raw), n_top_k(n_top_k), sinks(sinks), ns(ns), ov(ov), type_K(type_K) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hs, nb, nh, ns); ggml_set_name(q, "q"); - ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, hs, kv, 1, ns); + ggml_tensor * k = ggml_new_tensor_4d(ctx, type_K, hs, kv, 1, ns); ggml_set_name(k, "k"); // V4 CSA attends over the K latent itself: V is the same cache tensor @@ -10319,6 +10321,15 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 8, 2304, 512, false, 1, 60)); test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 16, 2304, 512, false, 1, 60)); test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 16, 2304, 512, false, 1, 86)); + // quantised K/V: the gather relocates rows verbatim, so it should serve any type whose row + // is a whole number of 4-byte words. These are the shapes a DSv4 decode with -ctk q8_0 hits, + // which took the dense fallback entirely before the gather learned to address rows as bytes. + for (ggml_type tk : { GGML_TYPE_Q8_0, GGML_TYPE_Q4_0 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 4, 1024, 512, false, 1, 0, tk)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 8, 2304, 512, false, 1, 60, tk)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 16, 2304, 512, false, 1, 86, tk)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 1, 1024, 512, false, 1, 0, tk)); + } test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 4, 1024, 512, false, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, true, 2)); @@ -10788,6 +10799,11 @@ static std::vector> make_test_cases_perf() { for (int nb : { 8, 16 }) { test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, nb, 2304, 512, false, 1, 86)); } + // q8_0 K/V at the same shapes: DSv4 with -ctk q8_0 took the dense fallback before the + // gather became type-agnostic, so this is the cell that says whether it now pays there. + for (int nb : { 2, 4, 8, 16 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, nb, 2304, 512, false, 1, 60, GGML_TYPE_Q8_0)); + } return test_cases; } From 6b7c998655f7af15d984831d52adb04268abf440 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Fri, 14 Aug 2026 06:52:40 +0000 Subject: [PATCH 099/109] vulkan: dequantise q8_0 K/V inside the DeepSeek V4 small-batch gather Relocating quantised rows verbatim leaves flash attention to decode them in its inner loop, which it does once per query block that reads a row rather than once per row. Measured, that costs 0.15 us per KV row attended - 0.1481 / 0.1510 / 0.1512 us/row at batch 8 / 4 / 2 and 0.1565 dense, so a constant across a 3.6x span of row counts and across both regimes. It is what made q8_0 attention 1.78x slower than f16 while reading half the bytes. The gather is the natural place to pay it instead: it already touches every selected row exactly once. Decoding there makes the compact scratch f16, flash attention takes its f16 path, and the redundancy disappears. No extra pass - the pass already existed. test-backend-ops perf, kv=11008 n_kv_raw=2304 n_top_k=512 ov=60, q8_0 K/V: batch dense verbatim gather decoded gather f16 gather 2 3917 us 1110 us 630 us 640 us 4 3924 us 1257 us 715 us 727 us 8 3936 us 1567 us 891 us 908 us q8_0 now edges f16 by 1.5 to 1.8% rather than trailing it by 78%: the FA does identical f16 work either way, and the gather's scattered read side moves half the bytes. Against the dense fallback a q8_0 cache took before any of this work, batch 8 is 4.41x. Flash attention picks its pipeline from what the SCRATCH holds, not from the source tensor, so k_type_eff/v_type_eff now also key off the compact state - the same mechanism use_dequant_kv already used, which is why that variable existed in this form. q8_0 only, deliberately. It is the type both community test sets used and the one worth having; every other quantised type keeps the verbatim path, which is correct for them and validated. Generalising means one shader variant per type through dequant_funcs.glsl, which is only worth the shader-permutation cost if this shape of win reproduces elsewhere. Not covered: the per-token gather still relocates verbatim, so a batch where the union declines but the per-token form fits keeps paying the inline decode. Same fix applies, and the same measurement will say whether it is worth a second shader. 13318 FLASH_ATTN_EXT cases pass. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 33 +++++-- .../flash_attn_gather_union_dq.comp | 99 +++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 1 + 3 files changed, 124 insertions(+), 9 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 46f312579880..47cd75141a5b 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1096,6 +1096,7 @@ struct vk_device_struct { vk_pipeline pipeline_flash_attn_gather_f16; vk_pipeline pipeline_flash_attn_union_f16; vk_pipeline pipeline_flash_attn_gather_union_f16; + vk_pipeline pipeline_flash_attn_gather_union_dq_q8_0; vk_pipeline pipeline_dsv4_hc_pre_f32; vk_pipeline pipeline_dsv4_hc_comb_f32; vk_pipeline pipeline_dsv4_hc_post_f32; @@ -6137,6 +6138,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { "flash_attn_gather_union_f16", flash_attn_gather_union_f16_len, flash_attn_gather_union_f16_data, "main", 6, sizeof(vk_op_flash_attn_gather_union_push_constants), {1, 1, 1}, {}, 1, true, true, device->subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_union_dq_q8_0, + "flash_attn_gather_union_dq_q8_0", flash_attn_gather_union_dq_q8_0_len, flash_attn_gather_union_dq_q8_0_data, "main", 6, + sizeof(vk_op_flash_attn_gather_union_push_constants), {1, 1, 1}, {}, 1, true, true, + device->subgroup_size); } } @@ -11393,6 +11398,7 @@ struct vk_fa_compact_state { uint32_t n_batch = 1; uint32_t row_bytes = 0; // bytes per compact K row; K may be quantised uint32_t row_elems = 0; // K row stride in ELEMENTS/blocks, for the FA push constant + bool dequantized = false; // scratch holds f16 because the gather decoded on the way in vk_subbuffer kv_buf; vk_subbuffer kc_buf, mc_buf; }; @@ -11621,7 +11627,12 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ goto union_unavailable; } - const size_t ukc_sz = (size_t) kv_c * k_row_bytes; + // Decoding on the way in makes the scratch f16 and hands flash attention its f16 path, + // which is worth far more than the extra scratch bytes: the inline decode it replaces + // costs a measured 0.15 us per KV row attended, every step. + const bool dq = k->type == GGML_TYPE_Q8_0 && ctx->device->pipeline_flash_attn_gather_union_dq_q8_0; + const uint32_t u_row_by = dq ? (uint32_t) (k->ne[0] * sizeof(ggml_fp16_t)) : k_row_bytes; + const size_t ukc_sz = (size_t) kv_c * u_row_by; const size_t umc_sz = (size_t) n_batch * kv_c * sizeof(ggml_fp16_t); const size_t ul_sz = (size_t) max_union * sizeof(uint32_t); const size_t need = ukc_sz + umc_sz + ul_sz; @@ -11641,8 +11652,10 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // prealloc_y sync above is a global barrier and orders the previous layer's read. const vk_subbuffer uc_buf = ggml_vk_subbuffer(ctx, ctx->fa_union_stat); + vk_pipeline gather_pipe = dq ? ctx->device->pipeline_flash_attn_gather_union_dq_q8_0 + : ctx->device->pipeline_flash_attn_gather_union_f16; ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, 1); - ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_gather_union_f16, 1); + ggml_pipeline_request_descriptor_sets(ctx, gather_pipe, 1); const vk_op_flash_attn_union_push_constants upc = { (uint32_t) k->ne[1], (uint32_t) n_kv_raw, n_batch, (uint32_t) top_k->ne[0], max_union, @@ -11652,13 +11665,14 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ { ggml_vk_tensor_subbuffer(ctx, top_k), ul_buf, uc_buf }, upc, { 1, 1, 1 }); ggml_vk_sync_buffers(ctx, subctx); + // the fused decoder steps K in BLOCKS and writes elements; the verbatim one does words const vk_op_flash_attn_gather_union_push_constants gpc = { (uint32_t) k->ne[1], (uint32_t) n_kv_raw, kv_c, - (uint32_t) (k->nb[1] / 4), + dq ? (uint32_t) (k->nb[1] / ggml_type_size(k->type)) : (uint32_t) (k->nb[1] / 4), (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), - n_batch, k_row_words, + n_batch, dq ? (uint32_t) k->ne[0] : k_row_words, }; - ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_gather_union_f16, + ggml_vk_dispatch_pipeline(ctx, subctx, gather_pipe, { ggml_vk_tensor_subbuffer(ctx, k), ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), kc_buf, mc_buf, uc_buf }, gpc, { kv_c, 1, 1 }); ggml_vk_sync_buffers(ctx, subctx); @@ -11669,8 +11683,9 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ st.dynamic_kv = true; st.kv_c = kv_c; st.n_batch = n_batch; - st.row_bytes = k_row_bytes; - st.row_elems = (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); + st.row_bytes = u_row_by; + st.row_elems = dq ? (uint32_t) k->ne[0] : (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); + st.dequantized = dq; st.kc_buf = kc_buf; st.mc_buf = mc_buf; st.kv_buf = uc_buf; @@ -11908,8 +11923,8 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx // If this fires, supports_op admitted a non-native K/V type the gate then rejected; the // native shader would return garbage rather than fail, so abort instead. GGML_ASSERT(use_dequant_kv || !kv_needs_dequant); - 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; + const ggml_type k_type_eff = (use_dequant_kv || fa_compact.dequantized) ? GGML_TYPE_F16 : k->type; + const ggml_type v_type_eff = (use_dequant_kv || fa_compact.dequantized) ? GGML_TYPE_F16 : v->type; // 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). diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp new file mode 100644 index 000000000000..fc2a0f4541a5 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp @@ -0,0 +1,99 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : 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 + +#include "types.glsl" + +// Gather + DEQUANTISE in one pass, for the DeepSeek V4 small-batch decode path with a +// quantised KV cache. +// +// The sibling flash_attn_gather_union.comp relocates rows verbatim, which leaves flash +// attention to dequantise them in its inner loop. That costs a MEASURED 0.15 us per KV row +// attended - constant across batch 2/4/8 and across the dense and gathered regimes alike, and +// large enough to make q8_0 attention 1.78x slower than f16 despite reading half the bytes. +// It is redundant work: the FA re-decodes the same row for every query block that reads it, +// whereas the gather already touches each selected row exactly once. +// +// So dequantise here instead. The compact scratch becomes f16, flash attention takes its f16 +// path, and the decode is paid once per row rather than once per use. The extra cost is +// writing f16 instead of q8_0 into the scratch; the pass itself already existed. +// +// q8_0 only. Other quantised types keep the verbatim path, which is correct for them and was +// validated; generalising means one variant per type through dequant_funcs.glsl, which is only +// worth doing if this shape of win holds up. + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer KBuf { block_q8_0_packed16 data_k[]; }; +layout(binding = 1) readonly buffer UBuf { uint data_u[]; }; +layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; +layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; +layout(binding = 5) readonly buffer CBuf { uint data_c[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_kv_raw; + uint kv_c_max; + uint nbk1; // K source row stride, BLOCKS + uint nbm1; + uint n_batch; + uint row_elems; // elements per K row (== head size) +} p; + +const uint LANES = 64; +const uint QK8_0 = 32; + +void main() { + const uint row = gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationIndex; + + const uint kv_c = data_c[0]; // padded compact rows, the FA's runtime KV + const uint n_uni = data_c[1]; // unpadded union size + + if (row >= kv_c) { + return; + } + + uint src = p.n_kv; // sentinel: invalid + if (row < p.n_kv_raw) { + src = row; + } else if (row - p.n_kv_raw < n_uni) { + src = p.n_kv_raw + data_u[row - p.n_kv_raw]; + } + + // 64 lanes x 8 elements covers a 512-element row exactly; a q8_0 block is 32 elements, so + // each lane stays inside one block and reads four int16 pairs from it. + const uint dst_base = row * p.row_elems; + const uint e0 = tid * 8; + if (src < p.n_kv) { + const uint ib = src * p.nbk1 + e0 / QK8_0; + const uint iqs = e0 % QK8_0; + const float d = float(data_k[ib].d); + [[unroll]] for (uint j = 0; j < 4; ++j) { + const i8vec2 v = unpack8(int32_t(data_k[ib].qs[iqs / 2 + j])).xy; + data_kc[dst_base + e0 + 2 * j] = float16_t(d * float(v.x)); + data_kc[dst_base + e0 + 2 * j + 1] = float16_t(d * float(v.y)); + } + } else { + [[unroll]] for (uint j = 0; j < 8; ++j) { + data_kc[dst_base + e0 + j] = float16_t(0.0); + } + } + + // Compact mask is token-major [n_batch][kv_c], stride the RUNTIME kv_c: the FA derives + // m_row_len from KV, which is that same runtime value. + const float NEG_INF = uintBitsToFloat(0xff800000); + if (tid < p.n_batch) { + float mv = NEG_INF; + if (src < p.n_kv) { + mv = float(data_m[tid * p.nbm1 + src]); + } + data_mc[tid * kv_c + row] = float16_t(mv); + } +} 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 f5b53b7fcf40..0aace33feadb 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -814,6 +814,7 @@ void process_shaders() { string_to_spv("flash_attn_gather_f16", "flash_attn_gather.comp", {}); string_to_spv("flash_attn_union_f16", "flash_attn_union.comp", {}); string_to_spv("flash_attn_gather_union_f16", "flash_attn_gather_union.comp", {}); + string_to_spv("flash_attn_gather_union_dq_q8_0", "flash_attn_gather_union_dq.comp", {}); string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {}); From 2cd621b90f40f27c711f4c41814b6e7d400812c8 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Fri, 14 Aug 2026 07:29:56 +0000 Subject: [PATCH 100/109] vulkan: decode q4_0 in the V4 gather too, with the real element mapping q4_0 turns out to want this more than q8_0 did. Its inline-decode penalty is the same shape - a constant per KV row attended, flat across batch and across dense vs gathered - but larger: 0.172 to 0.179 us/row against q8_0's 0.148 to 0.153. Unpacking nibbles costs more ALU than the halved bytes save, which is also why q4_0 was the SLOWEST of the three types before this. kv=11008 n_kv_raw=2304 n_top_k=512 ov=60, batch 2 / 4 / 8 type dense verbatim gather decoded gather f16 2192/2203/2213 - 640/726/902 q8_0 3917/3924/3936 1110/1257/1567 628/712/888 q4_0 4134/4141/4145 1189/1342/1671 626/711/886 All three converge, because after the gather they run the same f16 attention; what is left is the scattered-read side, where fewer bytes now wins. Against the dense fallback each type took before any of this work, batch 8 is 2.45x for f16, 4.43x for q8_0 and 4.68x for q4_0. The interesting part is what the first attempt got wrong. Generalising via dequantize4() from dequant_funcs.glsl looked obvious and passed on q8_0, then failed q4_0 with NMSE 1.17. That helper is permutation AGNOSTIC by design: mul_mat_vec only ever feeds it into a dot product, where any consistent element order gives the same sum, so for q4_0 it returns nibbles in packed order rather than element order. q8_0 passed only because its layout happens to be contiguous. Materialising a row into memory is precisely the use where the order matters, so each type's true mapping is now written out - for q4_0, byte j carries element j in its low nibble and element j+16 in its high one, exactly as dequant_q4_0.comp materialises it. That is also why this stops at two types rather than looping over the whole native list: every type needs its mapping stated and tested, an #error guards anything else reaching the shader, and q8_0 and q4_0 are the two actually used as a KV cache. q4_1/q5_0/q5_1/iq4_nl keep the verbatim gather, which is correct for them. 13318 FLASH_ATTN_EXT cases pass, including q4_0 and q8_0 top-k cases at the shapes a DSv4 decode hits, plus q4_0 perf cells alongside the q8_0 ones. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 17 +++-- .../flash_attn_gather_union_dq.comp | 70 +++++++++++++------ .../vulkan-shaders/vulkan-shaders-gen.cpp | 9 ++- tests/test-backend-ops.cpp | 6 +- 4 files changed, 70 insertions(+), 32 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 47cd75141a5b..f81a23556ba3 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1096,7 +1096,7 @@ struct vk_device_struct { vk_pipeline pipeline_flash_attn_gather_f16; vk_pipeline pipeline_flash_attn_union_f16; vk_pipeline pipeline_flash_attn_gather_union_f16; - vk_pipeline pipeline_flash_attn_gather_union_dq_q8_0; + vk_pipeline pipeline_flash_attn_gather_union_dq[GGML_TYPE_COUNT]; vk_pipeline pipeline_dsv4_hc_pre_f32; vk_pipeline pipeline_dsv4_hc_comb_f32; vk_pipeline pipeline_dsv4_hc_post_f32; @@ -6138,10 +6138,15 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { "flash_attn_gather_union_f16", flash_attn_gather_union_f16_len, flash_attn_gather_union_f16_data, "main", 6, sizeof(vk_op_flash_attn_gather_union_push_constants), {1, 1, 1}, {}, 1, true, true, device->subgroup_size); - ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_union_dq_q8_0, - "flash_attn_gather_union_dq_q8_0", flash_attn_gather_union_dq_q8_0_len, flash_attn_gather_union_dq_q8_0_data, "main", 6, - sizeof(vk_op_flash_attn_gather_union_push_constants), {1, 1, 1}, {}, 1, true, true, +#define CREATE_FA_GATHER_DQ(TYPE, NAMED) \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_union_dq[TYPE], \ + "flash_attn_gather_union_dq_" #NAMED, flash_attn_gather_union_dq_ ## NAMED ## _len, \ + flash_attn_gather_union_dq_ ## NAMED ## _data, "main", 6, \ + sizeof(vk_op_flash_attn_gather_union_push_constants), {1, 1, 1}, {}, 1, true, true, \ device->subgroup_size); + CREATE_FA_GATHER_DQ(GGML_TYPE_Q4_0, q4_0) + CREATE_FA_GATHER_DQ(GGML_TYPE_Q8_0, q8_0) +#undef CREATE_FA_GATHER_DQ } } @@ -11630,7 +11635,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // Decoding on the way in makes the scratch f16 and hands flash attention its f16 path, // which is worth far more than the extra scratch bytes: the inline decode it replaces // costs a measured 0.15 us per KV row attended, every step. - const bool dq = k->type == GGML_TYPE_Q8_0 && ctx->device->pipeline_flash_attn_gather_union_dq_q8_0; + const bool dq = ggml_is_quantized(k->type) && ctx->device->pipeline_flash_attn_gather_union_dq[k->type]; const uint32_t u_row_by = dq ? (uint32_t) (k->ne[0] * sizeof(ggml_fp16_t)) : k_row_bytes; const size_t ukc_sz = (size_t) kv_c * u_row_by; const size_t umc_sz = (size_t) n_batch * kv_c * sizeof(ggml_fp16_t); @@ -11652,7 +11657,7 @@ static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_ // prealloc_y sync above is a global barrier and orders the previous layer's read. const vk_subbuffer uc_buf = ggml_vk_subbuffer(ctx, ctx->fa_union_stat); - vk_pipeline gather_pipe = dq ? ctx->device->pipeline_flash_attn_gather_union_dq_q8_0 + vk_pipeline gather_pipe = dq ? ctx->device->pipeline_flash_attn_gather_union_dq[k->type] : ctx->device->pipeline_flash_attn_gather_union_f16; ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, 1); ggml_pipeline_request_descriptor_sets(ctx, gather_pipe, 1); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp index fc2a0f4541a5..fd3d0708b2f0 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp @@ -7,29 +7,32 @@ #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require #extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +// iq4_nl's value table lives in shared memory and is typed FLOAT_TYPE +#define FLOAT_TYPE float #include "types.glsl" // Gather + DEQUANTISE in one pass, for the DeepSeek V4 small-batch decode path with a -// quantised KV cache. +// quantised KV cache. One variant per K type, via DATA_A_*. // // The sibling flash_attn_gather_union.comp relocates rows verbatim, which leaves flash -// attention to dequantise them in its inner loop. That costs a MEASURED 0.15 us per KV row -// attended - constant across batch 2/4/8 and across the dense and gathered regimes alike, and -// large enough to make q8_0 attention 1.78x slower than f16 despite reading half the bytes. -// It is redundant work: the FA re-decodes the same row for every query block that reads it, -// whereas the gather already touches each selected row exactly once. +// attention to decode them in its inner loop - once per query block that reads a row, rather +// than once per row. Measured, that is a constant per KV row attended: // -// So dequantise here instead. The compact scratch becomes f16, flash attention takes its f16 -// path, and the decode is paid once per row rather than once per use. The extra cost is -// writing f16 instead of q8_0 into the scratch; the pass itself already existed. +// q8_0 0.148 - 0.153 us/row q4_0 0.172 - 0.179 us/row // -// q8_0 only. Other quantised types keep the verbatim path, which is correct for them and was -// validated; generalising means one variant per type through dequant_funcs.glsl, which is only -// worth doing if this shape of win holds up. +// flat across batch 2/4/8 and across the dense and gathered regimes alike. It is what made +// quantised attention slower than f16 despite reading fewer bytes, and q4_0 worse than q8_0 +// despite reading fewer bytes still: unpacking nibbles costs more ALU than the bytes save. +// +// The gather already touches each selected row exactly once, so decoding here converts per-use +// work into per-row work. The compact scratch becomes f16 and flash attention takes its f16 +// path. The pass itself already existed, so the only added cost is writing f16 rather than +// blocks into the scratch. layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; -layout(binding = 0) readonly buffer KBuf { block_q8_0_packed16 data_k[]; }; +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; + layout(binding = 1) readonly buffer UBuf { uint data_u[]; }; layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; @@ -47,9 +50,12 @@ layout(push_constant) uniform Parameters { } p; const uint LANES = 64; -const uint QK8_0 = 32; void main() { +#ifdef NEEDS_INIT_IQ_SHMEM + // barrier inside; must run before any divergent return below + init_iq_shmem(gl_WorkGroupSize); +#endif const uint row = gl_WorkGroupID.x; const uint tid = gl_LocalInvocationIndex; @@ -67,19 +73,37 @@ void main() { src = p.n_kv_raw + data_u[row - p.n_kv_raw]; } - // 64 lanes x 8 elements covers a 512-element row exactly; a q8_0 block is 32 elements, so - // each lane stays inside one block and reads four int16 pairs from it. + // 64 lanes x 8 elements covers a 512-element row exactly, and QUANT_K is 32, so a lane's 8 + // elements sit wholly inside one block AND wholly inside one nibble half. + // + // Deliberately NOT dequantize4() from dequant_funcs.glsl. That helper is permutation + // AGNOSTIC: mul_mat_vec only ever feeds it into a dot product, where any consistent element + // order gives the same answer, so for q4_0 it returns nibbles in packed order rather than + // element order. Materialising a row to memory is the one use where the order matters, and + // using it here produced NMSE 1.17 on q4_0 while q8_0 passed - q8_0's layout just happens to + // be contiguous. Each type's true element mapping is written out below instead. const uint dst_base = row * p.row_elems; const uint e0 = tid * 8; if (src < p.n_kv) { - const uint ib = src * p.nbk1 + e0 / QK8_0; - const uint iqs = e0 % QK8_0; - const float d = float(data_k[ib].d); - [[unroll]] for (uint j = 0; j < 4; ++j) { - const i8vec2 v = unpack8(int32_t(data_k[ib].qs[iqs / 2 + j])).xy; - data_kc[dst_base + e0 + 2 * j] = float16_t(d * float(v.x)); - data_kc[dst_base + e0 + 2 * j + 1] = float16_t(d * float(v.y)); + const uint ib = src * p.nbk1 + e0 / QUANT_K; + const uint iqs = e0 % QUANT_K; + const float d = float(data_a[ib].d); +#if defined(DATA_A_Q8_0) + // element e is qs[e] + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_kc[dst_base + e0 + l] = float16_t(d * float(data_a[ib].qs[iqs + l])); + } +#elif defined(DATA_A_Q4_0) + // byte j carries element j in its low nibble and element j+16 in its high one + const uint shift = (iqs >> 4) * 4; // 0 for elements 0..15, 4 for 16..31 + const uint byte0 = iqs & 0xF; + [[unroll]] for (uint l = 0; l < 8; ++l) { + const float q = float((data_a[ib].qs[byte0 + l] >> shift) & 0xF) - 8.0f; + data_kc[dst_base + e0 + l] = float16_t(d * q); } +#else +#error "no element mapping written for this K type" +#endif } else { [[unroll]] for (uint j = 0; j < 8; ++j) { data_kc[dst_base + e0 + j] = float16_t(0.0); 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 0aace33feadb..4ef16f6165a2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -814,7 +814,14 @@ void process_shaders() { string_to_spv("flash_attn_gather_f16", "flash_attn_gather.comp", {}); string_to_spv("flash_attn_union_f16", "flash_attn_union.comp", {}); string_to_spv("flash_attn_gather_union_f16", "flash_attn_gather_union.comp", {}); - string_to_spv("flash_attn_gather_union_dq_q8_0", "flash_attn_gather_union_dq.comp", {}); + // one decoder per quantised K type flash-attention supports natively; f16/bf16/f32 need no + // decode and take the verbatim gather + // q8_0 and q4_0 only: each needs its true element mapping written out (see the shader), and + // these are the two types actually used as a KV cache. The rest take the verbatim gather. + for (const auto& tname : {"q4_0", "q8_0"}) { + string_to_spv("flash_attn_gather_union_dq_" + std::string(tname), "flash_attn_gather_union_dq.comp", + {{"DATA_A_" + to_uppercase(tname), "1"}}); + } string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {}); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f01bf2ca7d14..6cf435fc0c3c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10801,8 +10801,10 @@ static std::vector> make_test_cases_perf() { } // q8_0 K/V at the same shapes: DSv4 with -ctk q8_0 took the dense fallback before the // gather became type-agnostic, so this is the cell that says whether it now pays there. - for (int nb : { 2, 4, 8, 16 }) { - test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, nb, 2304, 512, false, 1, 60, GGML_TYPE_Q8_0)); + for (ggml_type tk : { GGML_TYPE_Q8_0, GGML_TYPE_Q4_0 }) { + for (int nb : { 2, 4, 8, 16 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, nb, 2304, 512, false, 1, 60, tk)); + } } return test_cases; From 7d453235f9e537eb3447e1fc95fe522c5d035ea2 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Mon, 17 Aug 2026 09:31:21 +0000 Subject: [PATCH 101/109] vulkan : restore the unrolled row copy in the DSV4 gather shaders Generalising the gather from f16 elements to raw words (6b2cade31) replaced [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) with a loop bounded by the row_words push constant, which cannot be unrolled and pays a bounds check per iteration. Half as many, twice as wide accesses were still 6% slower on the smallest op the gather serves. Bisected over the 19 beta3 commits with the tip's test file overlaid at every point, so the instrument stayed fixed while beta3 grew its own perf cases. nb=1 at kv=8192: 54.58 us at the pre-beta3 base, flat through commit 16 (54.84), 58.24 at 6b2cade31, 58.23 at the tip. Step 4*LANES instead. Each access stays contiguous across the lanes, so coalescing is unchanged, and an f16 row at head size 512 is 256 words == 4*LANES: one iteration, no loop overhead. The tail loop carries the quantised row sizes, which are not multiples of 4*LANES (q8_0 is 136 words, q4_0 is 72). nb=1 against the pre-beta3 base, f16 K/V: -0.07% at kv=8192, -2.83% at kv=35584, mean -1.29% across the six shapes, so the regression is gone rather than reduced. Against beta3 as merged: -4.8% to -6.2%. The nb>=2 gains are untouched, worst cell -0.51%. 13318 FLASH_ATTN_EXT cases pass. Worth noting the gather itself was never the cost at nb=1: with GGML_VK_FA_TOPK_GATHER=0 that cell is 249.84 us against 58.14 us with it on. Assisted-by: Claude Opus 5 --- .../vulkan-shaders/flash_attn_gather.comp | 22 +++++++++++++++++-- .../flash_attn_gather_union.comp | 20 +++++++++++++++-- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp index 6dd24d31b517..4e16b2a9cece 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp @@ -74,13 +74,31 @@ void main() { // zero scale zeroes the block) as well as for f16. The row is -inf in the mask either way; // zeroing only keeps a garbage dot product from reaching the softmax as a NaN. const uint dst_base = (stream * p.kv_c + row) * p.row_words; + // row_words is a push constant, so this cannot be [[unroll]]ed the way it was when the row + // was a fixed 512 f16 elements. Stepping 4*LANES recovers that: each access stays contiguous + // across the lanes, and an f16 row (256 words == 4*LANES) is one iteration with no loop + // overhead. The tail carries the quantised row sizes, which are not multiples of 4*LANES. if (src < p.n_kv) { const uint src_base = stream * p.nbk3 + src * p.nbk1; - for (uint i = tid; i < p.row_words; i += LANES) { + uint i = tid; + for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { + data_kc[dst_base + i] = data_k[src_base + i]; + data_kc[dst_base + i + LANES] = data_k[src_base + i + LANES]; + data_kc[dst_base + i + 2 * LANES] = data_k[src_base + i + 2 * LANES]; + data_kc[dst_base + i + 3 * LANES] = data_k[src_base + i + 3 * LANES]; + } + for (; i < p.row_words; i += LANES) { data_kc[dst_base + i] = data_k[src_base + i]; } } else { - for (uint i = tid; i < p.row_words; i += LANES) { + uint i = tid; + for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { + data_kc[dst_base + i] = 0u; + data_kc[dst_base + i + LANES] = 0u; + data_kc[dst_base + i + 2 * LANES] = 0u; + data_kc[dst_base + i + 3 * LANES] = 0u; + } + for (; i < p.row_words; i += LANES) { data_kc[dst_base + i] = 0u; } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp index 4669e0d1c75e..306372ba8af6 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp @@ -63,13 +63,29 @@ void main() { // type here (a zero scale zeroes the block) as well as for f16. The row is -inf in the mask // either way; zeroing only keeps a garbage dot product from reaching the softmax as a NaN. const uint dst_base = row * p.row_words; + // Stepped 4*LANES for the same reason as flash_attn_gather.comp: row_words is a push + // constant, so the copy cannot be [[unroll]]ed, and an f16 row is one iteration this way. if (src < p.n_kv) { const uint src_base = src * p.nbk1; - for (uint i = tid; i < p.row_words; i += LANES) { + uint i = tid; + for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { + data_kc[dst_base + i] = data_k[src_base + i]; + data_kc[dst_base + i + LANES] = data_k[src_base + i + LANES]; + data_kc[dst_base + i + 2 * LANES] = data_k[src_base + i + 2 * LANES]; + data_kc[dst_base + i + 3 * LANES] = data_k[src_base + i + 3 * LANES]; + } + for (; i < p.row_words; i += LANES) { data_kc[dst_base + i] = data_k[src_base + i]; } } else { - for (uint i = tid; i < p.row_words; i += LANES) { + uint i = tid; + for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { + data_kc[dst_base + i] = 0u; + data_kc[dst_base + i + LANES] = 0u; + data_kc[dst_base + i + 2 * LANES] = 0u; + data_kc[dst_base + i + 3 * LANES] = 0u; + } + for (; i < p.row_words; i += LANES) { data_kc[dst_base + i] = 0u; } } From 14f7c554964db33db63cc3c6bbfe9f2b44a6030e Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 20 Aug 2026 00:38:13 +0000 Subject: [PATCH 102/109] vulkan: decode quantised K/V inside the DeepSeek V4 per-token gather The union gather decodes q8_0/q4_0 on the way in, but it exists only for n_batch > 1: a single token has nothing to deduplicate, so plain autoregressive decode takes the per-token gather, which relocated rows verbatim and left flash attention paying the same inline decode the union work measured at ~0.15 us per KV row attended. flash_attn_gather_dq.comp closes that: same row mapping as flash_attn_gather.comp, same per-type element mapping as flash_attn_gather_union_dq.comp (q4_0 materialised in element order, not packed order). The compact scratch becomes f16 and flash attention takes its f16 path, keyed off the compact state exactly as the union path already is. test-backend-ops perf, kv=11008 n_kv_raw=2304 n_top_k=512 ov=60, nb=1: type verbatim gather decoded gather q8_0 226.9 us 102.9 us q4_0 258.4 us 101.8 us f16 114.5 us 105.5 us (f16 never decodes; its two runs bound the harness spread) Both quantised types land on the f16 op time instead of trailing it 2.0-2.3x: after the gather the attention work is identical f16, and the gather's scattered read side moves fewer bytes. q8_0 and q4_0 only, as in the union: each type's element mapping must be stated and tested, an #error guards anything else reaching the shader, and every other type keeps the verbatim gather. nb=1 rows are added to the top-k perf grids so this regime stays measured. Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 33 ++++-- .../vulkan-shaders/flash_attn_gather_dq.comp | 106 ++++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 2 + tests/test-backend-ops.cpp | 6 +- 4 files changed, 138 insertions(+), 9 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_dq.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f81a23556ba3..b2df0cb6ffcb 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1094,6 +1094,7 @@ struct vk_device_struct { vk_pipeline pipeline_flash_attn_top_k_f16; vk_pipeline pipeline_flash_attn_top_k_cm_f16; vk_pipeline pipeline_flash_attn_gather_f16; + vk_pipeline pipeline_flash_attn_gather_dq[GGML_TYPE_COUNT]; vk_pipeline pipeline_flash_attn_union_f16; vk_pipeline pipeline_flash_attn_gather_union_f16; vk_pipeline pipeline_flash_attn_gather_union_dq[GGML_TYPE_COUNT]; @@ -6129,6 +6130,15 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { "flash_attn_gather_f16", flash_attn_gather_f16_len, flash_attn_gather_f16_data, "main", 5, sizeof(vk_op_flash_attn_gather_push_constants), {1, 1, 1}, {}, 1, true, true, device->subgroup_size); +#define CREATE_FA_GATHER_TOK_DQ(TYPE, NAMED) \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_dq[TYPE], \ + "flash_attn_gather_dq_" #NAMED, flash_attn_gather_dq_ ## NAMED ## _len, \ + flash_attn_gather_dq_ ## NAMED ## _data, "main", 5, \ + sizeof(vk_op_flash_attn_gather_push_constants), {1, 1, 1}, {}, 1, true, true, \ + device->subgroup_size); + CREATE_FA_GATHER_TOK_DQ(GGML_TYPE_Q4_0, q4_0) + CREATE_FA_GATHER_TOK_DQ(GGML_TYPE_Q8_0, q8_0) +#undef CREATE_FA_GATHER_TOK_DQ if (device->subgroup_arithmetic) { ggml_vk_create_pipeline(device, device->pipeline_flash_attn_union_f16, "flash_attn_union_f16", flash_attn_union_f16_len, flash_attn_union_f16_data, "main", 3, @@ -11705,8 +11715,14 @@ union_unavailable:; return false; } + // Decode on the way in, as the union path does: the gather touches each row once, while + // flash attention decodes once per query block that reads it. The union only covers + // n_batch > 1, so single-token decode lands here. + const bool tok_dq = ggml_is_quantized(k->type) && ctx->device->pipeline_flash_attn_gather_dq[k->type]; + const uint32_t row_by = tok_dq ? (uint32_t) (k->ne[0] * sizeof(ggml_fp16_t)) : k_row_bytes; + const uint32_t ns = (uint32_t) q->ne[3]; - const size_t kc_sz = (size_t) ns * kv_c * k_row_bytes; + const size_t kc_sz = (size_t) ns * kv_c * row_by; const size_t mc_sz = (size_t) ns * n_batch * kv_c * sizeof(ggml_fp16_t); if (ctx->prealloc_size_y < kc_sz + mc_sz) { @@ -11717,19 +11733,21 @@ union_unavailable:; ggml_vk_sync_buffers(ctx, subctx); } - vk_pipeline pipeline = ctx->device->pipeline_flash_attn_gather_f16; + vk_pipeline pipeline = tok_dq ? ctx->device->pipeline_flash_attn_gather_dq[k->type] + : ctx->device->pipeline_flash_attn_gather_f16; ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + // the fused decoder steps K in BLOCKS and writes elements; the verbatim one does words const vk_op_flash_attn_gather_push_constants pc = { (uint32_t) k->ne[1], (uint32_t) n_kv_raw, (uint32_t) top_k->ne[0], kv_c, - (uint32_t) (k->nb[1] / 4), - (uint32_t) (k->nb[3] / 4), + tok_dq ? (uint32_t) (k->nb[1] / ggml_type_size(k->type)) : (uint32_t) (k->nb[1] / 4), + tok_dq ? (uint32_t) (k->nb[3] / ggml_type_size(k->type)) : (uint32_t) (k->nb[3] / 4), (uint32_t) (top_k->nb[1] / sizeof(int32_t)), (uint32_t) (top_k->nb[3] / sizeof(int32_t)), (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (mask->nb[3] / sizeof(ggml_fp16_t)), (uint32_t) mask->ne[3], - n_batch, k_row_words, + n_batch, tok_dq ? (uint32_t) k->ne[0] : k_row_words, }; st.kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); @@ -11791,8 +11809,9 @@ union_unavailable:; st.active = true; st.kv_c = kv_c; st.n_batch = n_batch; - st.row_bytes = k_row_bytes; - st.row_elems = (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); + st.row_bytes = row_by; + st.row_elems = tok_dq ? (uint32_t) k->ne[0] : (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); + st.dequantized = tok_dq; return true; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_dq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_dq.comp new file mode 100644 index 000000000000..f2711539641a --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_dq.comp @@ -0,0 +1,106 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : 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 + +#define FLOAT_TYPE float +#include "types.glsl" + +// Gather + DEQUANTISE for the per-token compact form: same row mapping as the verbatim +// flash_attn_gather.comp, same decode as flash_attn_gather_union_dq.comp. One variant per K +// type, via DATA_A_*. +// +// The union carries the decode only for n_batch > 1, because it needs more than one token to +// deduplicate. Single-token decode is the common case and it lands here; a verbatim gather +// leaves flash attention to decode inline, measured 1.98x the f16 op at q8_0 and 2.26x at +// q4_0, kv 11008. + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer A { A_TYPE data_a[]; }; +layout(binding = 1) readonly buffer TopBuf { int data_top[]; }; +layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; +layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_kv_raw; + uint n_top_k; + uint kv_c; + uint nbk1; // K source row stride, BLOCKS + uint nbk3; // K source stream stride, BLOCKS + uint nbt1; + uint nbt3; + uint nbm1; + uint nbm3; + uint nem3; + uint n_batch; + uint row_elems; // elements per K row; 64 lanes x 8 covers the 512 the gate pins +} p; + +void main() { + const uint row = gl_WorkGroupID.x; + const uint stream = gl_WorkGroupID.z; + const uint tid = gl_LocalInvocationIndex; + + const uint ALL_TOKENS = 0xffffffffu; + uint src = p.n_kv; + uint owner = ALL_TOKENS; + if (row < p.n_kv_raw) { + src = row; + } else { + const uint off = row - p.n_kv_raw; + const uint tok = off / p.n_top_k; + const uint slot = off - tok * p.n_top_k; + if (tok < p.n_batch) { + owner = tok; + const int idx = data_top[stream * p.nbt3 + tok * p.nbt1 + slot]; + if (idx >= 0 && uint(idx) < p.n_kv - p.n_kv_raw) { + src = p.n_kv_raw + uint(idx); + } + } + } + + const uint dst_base = (stream * p.kv_c + row) * p.row_elems; + const uint e0 = tid * 8; + if (src < p.n_kv) { + const uint ib = stream * p.nbk3 + src * p.nbk1 + e0 / QUANT_K; + const uint iqs = e0 % QUANT_K; + const float d = float(data_a[ib].d); +#if defined(DATA_A_Q8_0) + // element e is qs[e] + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_kc[dst_base + e0 + l] = float16_t(d * float(data_a[ib].qs[iqs + l])); + } +#elif defined(DATA_A_Q4_0) + // byte j carries element j in its low nibble and element j+16 in its high one + const uint shift = (iqs >> 4) * 4; + const uint byte0 = iqs & 0xF; + [[unroll]] for (uint l = 0; l < 8; ++l) { + const float q = float((data_a[ib].qs[byte0 + l] >> shift) & 0xF) - 8.0f; + data_kc[dst_base + e0 + l] = float16_t(d * q); + } +#else +#error "no element mapping written for this K type" +#endif + } else { + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_kc[dst_base + e0 + l] = float16_t(0.0); + } + } + + const float NEG_INF = uintBitsToFloat(0xff800000); + if (tid < p.n_batch) { + const uint mc_idx = (stream * p.n_batch + tid) * p.kv_c + row; + float mv = NEG_INF; + if (src < p.n_kv && (owner == ALL_TOKENS || owner == tid)) { + mv = float(data_m[(stream % p.nem3) * p.nbm3 + tid * p.nbm1 + src]); + } + data_mc[mc_idx] = float16_t(mv); + } +} 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 4ef16f6165a2..38e67264bb8e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -821,6 +821,8 @@ void process_shaders() { for (const auto& tname : {"q4_0", "q8_0"}) { string_to_spv("flash_attn_gather_union_dq_" + std::string(tname), "flash_attn_gather_union_dq.comp", {{"DATA_A_" + to_uppercase(tname), "1"}}); + string_to_spv("flash_attn_gather_dq_" + std::string(tname), "flash_attn_gather_dq.comp", + {{"DATA_A_" + to_uppercase(tname), "1"}}); } string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 6cf435fc0c3c..1919182cf6e3 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10787,7 +10787,7 @@ static std::vector> make_test_cases_perf() { // worst-case compact set 2304 + nb*512 crosses kv/2 at nb=6, so those cells measure // whether the gate can be opened by the union rather than by the worst case. for (int kv : { 11008, 35584, 133888 }) { - for (int nb : { 2, 4, 8, 16 }) { + for (int nb : { 1, 2, 4, 8, 16 }) { test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, nb, 2304, 512, false, 1, 60)); } } @@ -10801,8 +10801,10 @@ static std::vector> make_test_cases_perf() { } // q8_0 K/V at the same shapes: DSv4 with -ctk q8_0 took the dense fallback before the // gather became type-agnostic, so this is the cell that says whether it now pays there. + // nb=1 is plain autoregressive decode: the union needs nb > 1 to dedup, so this width + // takes the per-token gather. Its f16 row is in the grid above. for (ggml_type tk : { GGML_TYPE_Q8_0, GGML_TYPE_Q4_0 }) { - for (int nb : { 2, 4, 8, 16 }) { + for (int nb : { 1, 2, 4, 8, 16 }) { test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, nb, 2304, 512, false, 1, 60, tk)); } } From a2b7675c0e9e29ba0c7f3ad2dc2272add3c8081d Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Thu, 20 Aug 2026 00:25:10 +0000 Subject: [PATCH 103/109] vulkan: dequantise the cache for the DeepSeek V4 sparse prefill ggml_vk_flash_attn_top_k declined any K/V but f16, so a DSv4 run with a quantised cache took dense FA for every prefill batch - O(kv) against the sparse path's O(n_kv_raw + n_top_k), a penalty that grows with depth. A 128GB community run measured it end to end: prefill -5.4/-22.2/-39.4/-52.9% against f16 at source depths 17k/33k/67k/134k with -ctk q8_0. The sparse shaders stay f16-only. A quantised cache is dequantised once per op into the prealloc_x f16 scratch by the dense path's own fused dequant+transpose pipeline, and the shaders read the scratch. K and V are the same tensor in this path, so one pass covers both at half the dense path's scratch footprint. Types without a dequant-transpose pipeline, oversized caches, and GGML_VK_FA_DEQUANT=0 keep the dense fallback. test-backend-ops perf at nb=1024, kv 5504/11008/19200/35584 (the community depths in compressed-K rows): q8_0 was 60.9 -> 427.6 ms dense (1.30x -> 8.27x of f16, growing with kv) and lands within +0.8% of f16's flat 46.9-51.7 ms with the scratch; the dequant pass itself costs under 1%. 36/36 FLASH_ATTN_EXT top-k cases pass, including 6 new quantised prefill cases (nb=64 ns=1/2, nb=128 split form). Co-Authored-By: Claude Fable 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 56 ++++++++++++++++++++++++---- tests/test-backend-ops.cpp | 17 +++++++++ 2 files changed, 65 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index b2df0cb6ffcb..ac1468c2cdb6 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11166,9 +11166,22 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { const ggml_tensor * top_k = dst->src[5]; static const char * top_k_env = getenv("GGML_VK_FA_TOPK"); + // The sparse shaders read f16 only. Serve a quantised cache by dequantising it once into + // the f16 scratch, the same pass the dense path uses. K and V are the same tensor here, so + // one pass covers both. Without this a quantised cache declines to dense FA, which costs + // O(kv) where the sparse path costs O(n_kv_raw + n_top_k). + static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); + const uint64_t kv_f16_sz = (uint64_t) ggml_nelements(k) * sizeof(ggml_fp16_t); + const bool dequant_kv = top_k && k->type != GGML_TYPE_F16 && + !(fa_dequant_env && fa_dequant_env[0] == '0') && + ctx->device->pipeline_dequant_transpose[k->type] != nullptr && + k->nb[0] == ggml_type_size(k->type) && + ggml_is_contiguously_allocated(k) && + kv_f16_sz <= ctx->device->properties.limits.maxStorageBufferRange && + ggml_vk_fa_dequant_scratch_fits(ctx, kv_f16_sz); if ((top_k_env && top_k_env[0] == '0') || !top_k || (!ctx->device->pipeline_flash_attn_top_k_f16 && !ctx->device->pipeline_flash_attn_top_k_cm_f16) || - q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16 || + q->type != GGML_TYPE_F32 || (k->type != GGML_TYPE_F16 && !dequant_kv) || v->type != k->type || !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || q->ne[0] != 512 || q->ne[1] < 64 || k->ne[0] != 512 || v->ne[0] != 512 || q->ne[2] != 64 || k->ne[2] != 1 || v->ne[2] != 1 || @@ -11254,14 +11267,43 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & } } + vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, k); + if (dequant_kv) { + if (ctx->prealloc_size_x < kv_f16_sz) { + ctx->prealloc_size_x = kv_f16_sz; + ggml_vk_preallocate_buffers(ctx, subctx); + } + vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type]; + ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1); + if (ctx->prealloc_x_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + const vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0, kv_f16_sz }; + const uint32_t k_nel = (uint32_t) ggml_nelements(k); + const std::vector tr_pc = { (uint32_t) k->ne[0], (uint32_t) k->ne[2], (uint32_t) k->ne[1], 0, k_nel }; + ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, tr_pc, { k_nel, 1, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + ctx->prealloc_x_need_sync = true; + k_buf = k_dst; + ggml_vk_perf_mark_subop(ctx, subctx, "FA_KV_DEQUANT (sub-op)"); + } + // strides of what the shaders actually read: the source cache, or the contiguous + // [HS, KV, n_head_kv, ns] f16 scratch the dequant just wrote + const uint32_t k_stride = dequant_kv ? (uint32_t) k->ne[0] : (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)); + const uint32_t k_nb3_el = dequant_kv ? (uint32_t) ((uint64_t) k->ne[0] * k->ne[1] * k->ne[2]) + : (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)); + const uint32_t k_nb2_byte = dequant_kv ? (uint32_t) ((uint64_t) k->ne[0] * k->ne[1] * sizeof(ggml_fp16_t)) + : (uint32_t) k->nb[2]; + const uint32_t k_nb3_byte = dequant_kv ? (uint32_t) (k_nb3_el * sizeof(ggml_fp16_t)) : (uint32_t) k->nb[3]; + vk_op_flash_attn_top_k_push_constants pc = { (uint32_t) q->ne[1], (uint32_t) k->ne[1], (uint32_t) n_kv_raw, (uint32_t) top_k->ne[0], (uint32_t) q->ne[2], (uint32_t) (q->nb[1] / sizeof(float)), (uint32_t) (q->nb[2] / sizeof(float)), (uint32_t) (q->nb[3] / sizeof(float)), - (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)), - (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + k_stride, + k_nb3_el, (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), (uint32_t) (mask->nb[3] / sizeof(ggml_fp16_t)), (uint32_t) (top_k->nb[1] / sizeof(int32_t)), @@ -11297,7 +11339,6 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & vk_fa_tuning_params tuning = get_fa_tuning_params(ctx->device, D, D, N, raw_kv, GGML_TYPE_F16, GGML_TYPE_F16, f32acc); const uint32_t q_stride = (uint32_t) (q->nb[1] / sizeof(float)); - const uint32_t k_stride = (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)); const bool aligned = raw_kv % tuning.block_cols == 0 && (q_stride & 7) == 0 && (k_stride & 7) == 0; const vk_fa_pipeline_state raw_state = get_fa_pipeline_state(ctx->device, tuning, D, D, aligned, f32acc, true, false, false, GGML_TYPE_F16, GGML_TYPE_F16); @@ -11339,7 +11380,6 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & const uint32_t packed_gqa = mask_stride_in_split_kv | 1u; const uint32_t packed_partitions = (partitions << 16) | 1; const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); - const vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, k); const vk_subbuffer mask_buf = ggml_vk_tensor_subbuffer(ctx, mask); const vk_subbuffer top_buf = ggml_vk_tensor_subbuffer(ctx, top_k); const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); @@ -11365,8 +11405,8 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & 1, NS, (uint32_t) mask->ne[1], (uint32_t) mask->ne[2], (uint32_t) mask->ne[3], q_stride, (uint32_t) q->nb[2], (uint32_t) q->nb[3], - k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], - k_stride, (uint32_t) k->nb[2], (uint32_t) k->nb[3], + k_stride, k_nb2_byte, k_nb3_byte, + k_stride, k_nb2_byte, k_nb3_byte, scale, 0.0f, 0.0f, n_head_log2, 1.0f, 1.0f, packed_gqa, mask_stride, packed_partitions, @@ -11399,7 +11439,7 @@ static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {q_buf, ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf, + {q_buf, k_buf, ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf, ggml_vk_tensor_subbuffer(ctx, top_k), ggml_vk_tensor_subbuffer(ctx, dst)}, pc, {(uint32_t) q->ne[1], (uint32_t) CEIL_DIV(q->ne[2], use_cm ? 32 : 8), (uint32_t) q->ne[3]}); ggml_vk_perf_mark_subop(ctx, subctx, use_cm ? "FA_TOP_K_CM (sub-op)" : "FA_TOP_K_SPARSE (sub-op)"); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 1919182cf6e3..610ef824b62a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10329,6 +10329,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 8, 2304, 512, false, 1, 60, tk)); test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, 16, 2304, 512, false, 1, 86, tk)); test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 1, 1024, 512, false, 1, 0, tk)); + // prefill widths (nb >= 64), where the sparse shaders run on a dequantised f16 scratch + // instead of the cache. ns=2 covers the scratch's stream stride, and the 4096 case is + // wide enough for the raw/selected split form. + test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false, 1, 0, tk)); + test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false, 2, 0, tk)); + test_cases.emplace_back(new test_flash_attn_ext_top_k(4096, 128, 256, 512, false, 1, 0, tk)); } test_cases.emplace_back(new test_flash_attn_ext_top_k(8192, 4, 1024, 512, false, 2)); test_cases.emplace_back(new test_flash_attn_ext_top_k( 768, 64, 64, 128, false, 2)); @@ -10808,6 +10814,17 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext_top_k(11008, nb, 2304, 512, false, 1, 60, tk)); } } + // PREFILL widths with quantised K/V, which the sparse path now serves through a one-shot + // dequant into the f16 scratch; quantised should sit within ~1% of f16 at every kv here. + // nb=1024 is the reporting user's --ubatch-size; the kv list is their four source depths + // (17k/33k/67k/134k) in compressed-K rows. GGML_VK_FA_DEQUANT=0 reproduces the old dense + // fallback, whose gap GROWS with kv (dense is O(kv), sparse O(n_kv_raw + n_top_k)); f16 + // under GGML_VK_FA_TOPK=0 is the falsification arm for attributing that gap to the gate. + for (ggml_type tk : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0 }) { + for (int kv : { 5504, 11008, 19200, 35584 }) { + test_cases.emplace_back(new test_flash_attn_ext_top_k(kv, 1024, 2304, 512, false, 1, 0, tk)); + } + } return test_cases; } From ca507634bfe1e492134cebf2d8435674939239b7 Mon Sep 17 00:00:00 2001 From: pepuscz Date: Mon, 24 Aug 2026 22:17:58 +0200 Subject: [PATCH 104/109] vulkan: use small Lightning Indexer CM for batches 4-15 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index ac1468c2cdb6..fc35bc5c3838 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -12857,6 +12857,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const if (ctx->device->pipeline_lightning_indexer_decode_cm_f16 && src0->ne[2] == 1) { return ctx->device->pipeline_lightning_indexer_decode_cm_f16; } + static const char * small_cm_env = getenv("GGML_VK_LIGHTNING_INDEXER_SMALL_CM"); + if (ctx->device->pipeline_lightning_indexer_cm_small_f16 && + src0->ne[2] >= 4 && src0->ne[2] < 16 && small_cm_env && small_cm_env[0] == '1') { + return ctx->device->pipeline_lightning_indexer_cm_small_f16; + } vk_pipeline cm = ctx->device->pipeline_lightning_indexer_cm_f16 ? ctx->device->pipeline_lightning_indexer_cm_f16 : ctx->device->pipeline_lightning_indexer_cm_small_f16; return cm && src0->ne[2] >= 16 ? cm : ctx->device->pipeline_lightning_indexer_f16; From 1ca0f9ac373e83f4a0d31efd2e0063f439654de9 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Wed, 26 Aug 2026 05:01:45 +0000 Subject: [PATCH 105/109] vulkan: hoist the Lightning Indexer K fragments out of the head loop The K tile is invariant across the head loop but coopMatLoad sits inside it, so the shader re-read the same 8 MatrixA fragments from shared memory once per head tile - 4x at N_HEAD=64. Each wave64 fragment load moves 1 KiB of LDS traffic because the 16x16 f16 fragment is replicated 4x across the subgroup. Same loads of the same data feeding the same coopMatMulAdd sequence, so the output is bit-identical; the fragments just live in registers (8 of them, 16 f16 per lane) instead of being re-fetched. gfx1151, test-backend-ops perf, kv 8704/33280/131584 x batch 1-15: 16-20% faster at every shape, no spills, no regression. Batch 1 is ordinary DSv4 decode and gets 12-18% of that. Co-Authored-By: Claude Opus 5 --- .../lightning_indexer_decode_cm.comp | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp index fd555c76806d..24d1b98d571c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp @@ -63,6 +63,17 @@ void main() { } barrier(); + // The K tile is invariant across the head loop, but coopMatLoad is inside it, so the shipped + // shader re-reads the same 8 A fragments from shared memory once per head tile (4x at + // N_HEAD=64). Each wave64 fragment load moves 1 KiB of LDS traffic because the 16x16 f16 + // fragment is replicated 4x across the subgroup, so those re-reads are the largest single + // item of per-tile cost after the multiplies themselves. Hoisting costs 8 A fragments of + // register state (16 f16 per lane each). + coopmat kmat[HEAD_SIZE / TILE]; + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(kmat[d / TILE], k_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + } + float total = 0.0; for (uint head_base = 0; head_base < N_HEAD; head_base += TILE) { for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { @@ -77,13 +88,11 @@ void main() { coopmat scores = coopmat(0.0); - coopmat kmat; coopmat qmat; [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { - coopMatLoad(kmat, k_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); coopMatLoad(qmat, q_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); - scores = coopMatMulAdd(kmat, qmat, scores); + scores = coopMatMulAdd(kmat[d / TILE], qmat, scores); } coopMatStore(scores, score_sh, 0, SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); From 06b89c3cf29aab32f777d7204d8abf4212f4f943 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Wed, 26 Aug 2026 05:01:45 +0000 Subject: [PATCH 106/109] vulkan: route the whole small-batch Lightning Indexer window to the decode CM shader Builds directly on pepuscz's PR #6 and issue #10. Their per-kernel table - 43.8 us scalar against 27.9 us small CM per 1k scanned rows at batch 5, and the same ratio at every depth - is what showed the indexer's cost is a per-tile constant rather than anything to do with tokens or bytes. Once the cost is per tile, the thing to minimise is the tile count, and that is what this changes. The finding is downstream of their work; only the choice of shader differs. The decode CM shader puts 16 HEADS in the coopmat N dimension and dispatches one workgroup per token, so it issues 4*n_batch tiles per 16 KV rows. That is the arithmetic minimum - (64 heads x n_batch tokens) / 16 columns - because 64 heads fill its 16 columns exactly, with no remainder at any batch. The small CM shader puts 16 TOKENS in N and pays a flat 64 tiles however small the batch is, so at batch 5 eleven of its sixteen columns are padding. Measured cost tracks tile count: at kv=131584 batch 1 costs 1.95 us per 1k scanned rows and batch 5 costs 26.28, i.e. 13.5x for 5x the work, and both shaders sit at 6.6-7.8 ns per tile. The decode CM shader body has no n_batch == 1 assumption - token is gl_WorkGroupID.y and it indexes q/w/mask/dst by it - so the old gate was an artefact of where it was written. gfx1151, kv=131584 (526k source tokens), us/run, shipped vs this: batch 2 2335.6 -> 419.0 (5.6x, was on the scalar path) batch 3 3492.3 -> 631.8 (5.5x, was on the scalar path) batch 4 3267.9 -> 840.8 (3.9x) batch 5 3465.6 -> 1119.4 (3.1x) <- DSpark n-max 4 verify shape batch 8 3704.9 -> 1764.1 (2.1x) batch 15 4467.9 -> 3231.6 (1.4x) Batch 16 and 32 are unchanged, which confirms the arms are isolated. PR #6's route is kept as the opt-out arm rather than deleted, and is promoted from opt-in to default-on so that one variable is enough to reach it: default decode CM for the whole 2-15 window ..._DECODE_CM_BATCH=0 small CM for 4-15, scalar for 2-3 (PR #6) ..._DECODE_CM_BATCH=0 SMALL_CM=0 scalar for 2-15 (pre-PR #6 baseline) Verified all three route as documented and pass 29/29. Keeping it is not just courtesy: the tile-count argument is hardware independent, but decode CM re-reads the K tile once per token and that part is bandwidth dependent, so the crossover need not sit in the same place on other devices, and these numbers are from one gfx1151 box. A (head, token) packing shader was also tried and is strictly worse - it reaches 4b tiles only when the batch divides 16, and it divides the workgroup count by the batch. Kept out of tree. Numerics: decode CM sums the 64 heads in tiles of 16 rather than one at a time, so f32 accumulation order differs from the small CM path. Cleared by KLD A/B on trunc10 (-c 8192, -b/-ub small so every indexer dispatch routes through the window, 8 chunks wikitext-2, same binary, only the env var differing): -ub 8 mean KLD 0.000000, max 6.0e-5, RMS dp 0.000%, same-top 99.994%, PPL 10507558.3156 identical in both arms -ub 5 mean KLD 0.000000, max 8.8e-4 but 99.9% 4.9e-5 (one tail event, no argmax flip), RMS dp 0.000%, same-top 100.000% Same class as PR #6's own measured numbers (max 5.5e-5, same-top 100.000%) and 45x tighter than the FA_WAVE32 change already shipped. Adds the eval coverage the 2-15 window never had (batches 4/8/15 at kv=256) and the perf grid these numbers came from. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 33 +++++++++++++++++++++++++--- tests/test-backend-ops.cpp | 10 ++++++++- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index fc35bc5c3838..429fabe9b30a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -12854,12 +12854,39 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32 && src0->ne[0] == 128 && src0->ne[1] == 64 && src1->ne[1] == 1 && ctx->device->pipeline_lightning_indexer_f16) { - if (ctx->device->pipeline_lightning_indexer_decode_cm_f16 && src0->ne[2] == 1) { + // Small-batch routing. Three arms, one env var each, most specific first: + // + // default decode CM for the whole 2-15 window + // ..._DECODE_CM_BATCH=0 small CM for 4-15, scalar for 2-3 (PR #6) + // ..._DECODE_CM_BATCH=0 SMALL_CM=0 scalar for 2-15 (pre-PR #6 baseline) + // + // The decode CM shader puts 16 HEADS in the coopmat N dimension and dispatches one + // workgroup per token, so it issues 4*n_batch tiles per 16 KV rows - the arithmetic + // minimum, since 64 heads fill its 16 columns exactly. The small CM shader puts 16 + // TOKENS in N and pays a flat 64 tiles no matter how small the batch. Routing the + // whole window to decode CM is 2.7x at batch 5 and 4.8x at batch 2-3 (526k source + // tokens, gfx1151); the shader body has no n_batch == 1 assumption, token is + // gl_WorkGroupID.y, so the old ne[2] == 1 gate was an artefact. + // + // That measurement only exists because of pepuscz's PR #6 and issue #10: their + // per-kernel table (43.8 us scalar vs 27.9 us small CM per 1k scanned rows at batch + // 5, the same ratio at every depth) is what showed the cost is a per-tile constant, + // which is what makes the tile count the thing to minimise. Their small CM route is + // kept as the opt-out arm rather than deleted: the tile-count argument is hardware + // independent, but decode CM re-reads the K tile once per token, and that part is + // bandwidth dependent, so the crossover need not sit here on other devices. + static const char * decode_cm_batch_env = getenv("GGML_VK_LIGHTNING_INDEXER_DECODE_CM_BATCH"); + static const bool decode_cm_batch_on = !decode_cm_batch_env || decode_cm_batch_env[0] != '0'; + const int64_t decode_cm_max = decode_cm_batch_on ? 15 : 1; + if (ctx->device->pipeline_lightning_indexer_decode_cm_f16 && src0->ne[2] <= decode_cm_max) { return ctx->device->pipeline_lightning_indexer_decode_cm_f16; } + // PR #6 as its author proposed it, now default-on so that the kill switch above is + // enough on its own to reach it. static const char * small_cm_env = getenv("GGML_VK_LIGHTNING_INDEXER_SMALL_CM"); - if (ctx->device->pipeline_lightning_indexer_cm_small_f16 && - src0->ne[2] >= 4 && src0->ne[2] < 16 && small_cm_env && small_cm_env[0] == '1') { + static const bool small_cm_on = !small_cm_env || small_cm_env[0] != '0'; + if (ctx->device->pipeline_lightning_indexer_cm_small_f16 && small_cm_on && + src0->ne[2] >= 4 && src0->ne[2] < 16) { return ctx->device->pipeline_lightning_indexer_cm_small_f16; } vk_pipeline cm = ctx->device->pipeline_lightning_indexer_cm_f16 ? diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 610ef824b62a..fdf5a1620f5f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10267,7 +10267,7 @@ static std::vector> make_test_cases_eval() { // lightning_indexer for (int kv : { 256 }) { - for (int bs : { 1, 512 }) { + for (int bs : { 1, 4, 8, 15, 512 }) { for (int nh : { 32, 64 }) { for (auto [ns, nm] : { std::pair{1, 1}, std::pair{4, 4}, std::pair{4, 1} }) { for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) { @@ -10765,6 +10765,14 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_lightning_indexer(128, 64, kv, 2048, 1, 1, GGML_TYPE_F16)); } + // DSpark verify-step indexer rows: batch 1-32 across the 4-15 small-CM routing window, + // at the compressed key counts a 128k/491k source context produces (source/4). + for (int kv : { 8704, 33280, 131584 }) { + for (int bs : { 1, 2, 3, 4, 5, 6, 8, 12, 15, 16, 32 }) { + test_cases.emplace_back(new test_lightning_indexer(128, 64, kv, bs, 1, 1, GGML_TYPE_F16)); + } + } + // sparse top-k FA at V4 decode/prefill shapes — the A/B instrument for the // gather-to-compact work (n_active = n_kv_raw + n_top_k stays fixed as kv grows). // nb 1/8 currently takes the DENSE path (the sparse shader gates on nb >= 64): From 849aca6f5a318747fff8e66bef6910fbabc145f8 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Wed, 26 Aug 2026 07:15:56 +0000 Subject: [PATCH 107/109] vulkan: hoist the Lightning Indexer K fragments out of the CM head loop Same defect as the decode CM shader, larger blast radius. `k_sh[wave]` is written once above the head loop and never touched again, but `coopMatLoad` sits inside both the head_base and head_local loops, so each wave re-reads the same 8 MatrixA fragments N_HEAD times: 512 loads where 8 would do. Each wave64 fragment load moves 1 KiB of shared-memory traffic because the 16x16 f16 fragment is replicated 4x across the subgroup. Same loads of the same data feeding the same coopMatMulAdd order, so the output is bit-identical; the fragments live in registers (8 of them, 16 f16 per lane) instead of being re-fetched. This source compiles to both the prefill pipeline and the small-batch one, so both gain. gfx1151, test-backend-ops perf, us/run: prefill (lightning_indexer_cm_f16) kv 8704 nb 2048 27492.5 -> 18207.0 1.51x kv 33280 nb 2048 104101.9 -> 68947.8 1.51x kv 131584 nb 2048 421239.8 -> 281540.5 1.50x kv 131584 nb 16 3456.2 -> 2329.2 1.48x kv 131584 nb 32 6858.6 -> 4623.6 1.48x small batch (lightning_indexer_cm_small_f16) kv 131584 nb 5 3433.1 -> 2783.3 1.23x kv 131584 nb 8 3682.3 -> 3105.1 1.19x kv 131584 nb 15 4474.4 -> 4209.5 1.06x 1.48-1.54x on prefill at every depth and batch measured, flat. The small-batch variant gains less, which says the 1-wave configuration is bound by its Q staging and epilogue rather than by fragment traffic. No regression at any shape, so the 8 held fragments (64 VGPRs per lane) neither spill nor cost occupancy even in the prefill pipeline, which runs 512 threads at 62.7 KiB of shared memory and was already pinned to one workgroup per CU. LIGHTNING_INDEXER 29/29. Co-Authored-By: Claude Opus 5 --- .../vulkan-shaders/lightning_indexer_cm.comp | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp index d4379e8899d0..6747feba505e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp @@ -88,6 +88,15 @@ void main() { } barrier(); + // k_sh[wave] is written once above and never touched again, but coopMatLoad sits inside both + // the head_base and head_local loops, so the same 8 A fragments are re-read from shared + // memory N_HEAD times per wave (512 loads where 8 would do). Each wave64 fragment load moves + // 1 KiB of LDS traffic because the 16x16 f16 fragment is replicated 4x across the subgroup. + coopmat kmat_h[HEAD_SIZE / TILE]; + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(kmat_h[d / TILE], k_sh[wave], d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + } + for (uint head_base = 0; head_base < N_HEAD; head_base += HEADS_PER_TILE) { for (uint idx = tid; idx < HEADS_PER_TILE * TILE * VEC_PER_HEAD; idx += gl_WorkGroupSize.x) { const uint head_local = idx / (TILE * VEC_PER_HEAD); @@ -113,13 +122,11 @@ void main() { [[unroll]] for (uint head_local = 0; head_local < HEADS_PER_TILE; ++head_local) { coopmat scores = coopmat(0.0); - coopmat kmat; coopmat qmat; [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { - coopMatLoad(kmat, k_sh[wave], d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); coopMatLoad(qmat, q_sh[head_local], d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); - scores = coopMatMulAdd(kmat, qmat, scores); + scores = coopMatMulAdd(kmat_h[d / TILE], qmat, scores); } coopMatStore(scores, score_sh[wave], 0, SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); From e85148e47748f9c3613d091a472e0f9a6e8f7ca1 Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Wed, 26 Aug 2026 13:29:50 +0000 Subject: [PATCH 108/109] vulkan: support arbitrary Lightning Indexer head counts via specialization constant The indexer pipelines hardcoded nh=64 (DeepSeek-V4's geometry) in supports_op and as a shader constant, so any other head count fell to the CPU backend. Qwen3.8-Flash-Next (qwen4_exp, released today) carries the same indexer at nh=4. N_HEAD becomes specialization constant 1, one pipeline per supported head count (LI_NH_VALUES = {4, 32, 64}), selected by q->ne[1] at dispatch. It must be a specialization constant, not a push constant: the head loop's trip count has to stay visible to the compiler - as a push constant the loop cannot unroll and decode costs ~55% (measured, gfx1151, kv=131584, and why this commit is not the simpler design). nh=4 does not fill a 16-wide head tile, so the tail tile stages q as zeros (relu(0)*w = 0, padded heads drop out of the sum) and the weight read clamps its index. Both guards are gated on N_HEAD % TILE != 0, a specialization-constant expression, so at nh=64 they fold away at pipeline compile and the codegen is equivalent to the previous shader - measured, because relying on the compiler to range-prove head < N_HEAD instead still cost ~45%: kv=131584, us/run before after decode b1 226.8 219.6 decode b5 1058.8 1007.4 decode b15 3174.0 2966.8 prefill b16 2329.2 2230.2 prefill b2048 281540.5 281169.1 No regression at nh=64; the small decode improvement is within a cautious reading of run-to-run variance and is not claimed. test-backend-ops LIGHTNING_INDEXER: 44/44, up from 29/29 - the suite's existing nh=32 cases had been silently reporting "not supported" under the old gate and now run against the CPU reference. Head size stays pinned at 128: the coopmat tiling and f16vec4 staging both assume it, and every known indexer model (DSv4, GLM-DSA, Qwen4-exp) uses 128. Co-Authored-By: Claude Opus 5 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 96 ++++++++++++------- .../vulkan-shaders/lightning_indexer_cm.comp | 10 +- .../lightning_indexer_decode_cm.comp | 28 +++++- 3 files changed, 94 insertions(+), 40 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 429fabe9b30a..b6fe70eb8532 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -784,6 +784,22 @@ static bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type) { return std::find(lightning_indexer_k_types.begin(), lightning_indexer_k_types.end(), type) != lightning_indexer_k_types.end(); } +// Indexer head counts we build kernels for: 4 = Qwen4-exp (qwen4_exp), 64 = DeepSeek-V4 and +// GLM-DSA, 32 kept because test-backend-ops exercises it. N_HEAD is a specialization constant, +// so each entry is a separately compiled kernel with the head loop fully unrolled - as a push +// constant the loop cannot unroll and decode costs ~55% more (measured, gfx1151, kv=131584). +static constexpr uint32_t LI_NH_VALUES[] = { 4, 32, 64 }; +#define LI_NH_COUNT (sizeof(LI_NH_VALUES) / sizeof(LI_NH_VALUES[0])) + +static int ggml_vk_li_nh_index(int64_t nh) { + for (size_t i = 0; i < LI_NH_COUNT; ++i) { + if ((int64_t) LI_NH_VALUES[i] == nh) { + return (int) i; + } + } + return -1; +} + struct vk_device_struct { std::recursive_mutex mutex; mutable std::shared_mutex pinned_memory_mutex; @@ -1087,10 +1103,11 @@ struct vk_device_struct { vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT]; // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128 vk_pipeline pipeline_gated_delta_net[4][2]; - vk_pipeline pipeline_lightning_indexer_f16; - vk_pipeline pipeline_lightning_indexer_cm_f16; - vk_pipeline pipeline_lightning_indexer_cm_small_f16; - vk_pipeline pipeline_lightning_indexer_decode_cm_f16; + // One pipeline per supported indexer head count; see LI_NH_VALUES. + vk_pipeline pipeline_lightning_indexer_f16[LI_NH_COUNT]; + vk_pipeline pipeline_lightning_indexer_cm_f16[LI_NH_COUNT]; + vk_pipeline pipeline_lightning_indexer_cm_small_f16[LI_NH_COUNT]; + vk_pipeline pipeline_lightning_indexer_decode_cm_f16[LI_NH_COUNT]; vk_pipeline pipeline_flash_attn_top_k_f16; vk_pipeline pipeline_flash_attn_top_k_cm_f16; vk_pipeline pipeline_flash_attn_gather_f16; @@ -6094,28 +6111,36 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } if (device->subgroup_arithmetic && device->subgroup_size == 64) { - ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f16, - "lightning_indexer_f16", lightning_indexer_f16_len, lightning_indexer_f16_data, "main", 5, - sizeof(vk_op_lightning_indexer_cm_push_constants), {8, 1, 1}, {device->subgroup_size}, 1, true, true, - device->subgroup_size); + for (size_t nhi = 0; nhi < LI_NH_COUNT; ++nhi) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f16[nhi], + "lightning_indexer_f16", lightning_indexer_f16_len, lightning_indexer_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {8, 1, 1}, {device->subgroup_size, LI_NH_VALUES[nhi]}, 1, true, true, + device->subgroup_size); + } #if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) if (device->coopmat_support && device->coopmat_support_16x16x16_f32acc && device->subgroup_size_control) { - ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_small_f16, - "lightning_indexer_cm_small_f16", lightning_indexer_cm_small_f16_len, lightning_indexer_cm_small_f16_data, "main", 5, - sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 16, 1}, {device->subgroup_size}, 1, true, true, - device->subgroup_size); + for (size_t nhi = 0; nhi < LI_NH_COUNT; ++nhi) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_small_f16[nhi], + "lightning_indexer_cm_small_f16", lightning_indexer_cm_small_f16_len, lightning_indexer_cm_small_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 16, 1}, {device->subgroup_size, LI_NH_VALUES[nhi]}, 1, true, true, + device->subgroup_size); + } if (device->properties.limits.maxComputeWorkGroupInvocations >= 512 && device->properties.limits.maxComputeWorkGroupSize[0] >= 512 && device->properties.limits.maxComputeSharedMemorySize >= 64 * 1024) { - ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_f16, - "lightning_indexer_cm_f16", lightning_indexer_cm_f16_len, lightning_indexer_cm_f16_data, "main", 5, - sizeof(vk_op_lightning_indexer_cm_push_constants), {128, 16, 1}, {device->subgroup_size}, 1, true, true, + for (size_t nhi = 0; nhi < LI_NH_COUNT; ++nhi) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_f16[nhi], + "lightning_indexer_cm_f16", lightning_indexer_cm_f16_len, lightning_indexer_cm_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {128, 16, 1}, {device->subgroup_size, LI_NH_VALUES[nhi]}, 1, true, true, + device->subgroup_size); + } + } + for (size_t nhi = 0; nhi < LI_NH_COUNT; ++nhi) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_decode_cm_f16[nhi], + "lightning_indexer_decode_cm_f16", lightning_indexer_decode_cm_f16_len, lightning_indexer_decode_cm_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 1, 1}, {device->subgroup_size, LI_NH_VALUES[nhi]}, 1, true, true, device->subgroup_size); } - ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_decode_cm_f16, - "lightning_indexer_decode_cm_f16", lightning_indexer_decode_cm_f16_len, lightning_indexer_decode_cm_f16_data, "main", 5, - sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 1, 1}, {device->subgroup_size}, 1, true, true, - device->subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_flash_attn_top_k_cm_f16, "flash_attn_top_k_cm_f16", flash_attn_top_k_cm_f16_len, flash_attn_top_k_cm_f16_data, "main", 6, sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true, @@ -12850,10 +12875,12 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const return nullptr; case GGML_OP_LIGHTNING_INDEXER: // fork fast path: f16 K on wave64 subgroup-arithmetic devices routes to the tuned - // scalar-64/CM kernels; anything else falls through to the generic pipeline table + // scalar-64/CM kernels (head counts in LI_NH_VALUES); anything else falls through to + // the generic pipeline table if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32 && - src0->ne[0] == 128 && src0->ne[1] == 64 && src1->ne[1] == 1 && - ctx->device->pipeline_lightning_indexer_f16) { + src0->ne[0] == 128 && src1->ne[1] == 1 && + ggml_vk_li_nh_index(src0->ne[1]) >= 0 && ctx->device->pipeline_lightning_indexer_f16[0]) { + const int nhi = ggml_vk_li_nh_index(src0->ne[1]); // Small-batch routing. Three arms, one env var each, most specific first: // // default decode CM for the whole 2-15 window @@ -12878,20 +12905,20 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const static const char * decode_cm_batch_env = getenv("GGML_VK_LIGHTNING_INDEXER_DECODE_CM_BATCH"); static const bool decode_cm_batch_on = !decode_cm_batch_env || decode_cm_batch_env[0] != '0'; const int64_t decode_cm_max = decode_cm_batch_on ? 15 : 1; - if (ctx->device->pipeline_lightning_indexer_decode_cm_f16 && src0->ne[2] <= decode_cm_max) { - return ctx->device->pipeline_lightning_indexer_decode_cm_f16; + if (ctx->device->pipeline_lightning_indexer_decode_cm_f16[nhi] && src0->ne[2] <= decode_cm_max) { + return ctx->device->pipeline_lightning_indexer_decode_cm_f16[nhi]; } // PR #6 as its author proposed it, now default-on so that the kill switch above is // enough on its own to reach it. static const char * small_cm_env = getenv("GGML_VK_LIGHTNING_INDEXER_SMALL_CM"); static const bool small_cm_on = !small_cm_env || small_cm_env[0] != '0'; - if (ctx->device->pipeline_lightning_indexer_cm_small_f16 && small_cm_on && + if (ctx->device->pipeline_lightning_indexer_cm_small_f16[nhi] && small_cm_on && src0->ne[2] >= 4 && src0->ne[2] < 16) { - return ctx->device->pipeline_lightning_indexer_cm_small_f16; + return ctx->device->pipeline_lightning_indexer_cm_small_f16[nhi]; } - vk_pipeline cm = ctx->device->pipeline_lightning_indexer_cm_f16 ? - ctx->device->pipeline_lightning_indexer_cm_f16 : ctx->device->pipeline_lightning_indexer_cm_small_f16; - return cm && src0->ne[2] >= 16 ? cm : ctx->device->pipeline_lightning_indexer_f16; + vk_pipeline cm = ctx->device->pipeline_lightning_indexer_cm_f16[nhi] ? + ctx->device->pipeline_lightning_indexer_cm_f16[nhi] : ctx->device->pipeline_lightning_indexer_cm_small_f16[nhi]; + return cm && src0->ne[2] >= 16 ? cm : ctx->device->pipeline_lightning_indexer_f16[nhi]; } // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer() if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) { @@ -13985,9 +14012,14 @@ static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& GGML_ASSERT(pipeline != nullptr); // the fork's wave64 f16 kernels take their own push-constant layout - if (pipeline == ctx->device->pipeline_lightning_indexer_f16 || - pipeline == ctx->device->pipeline_lightning_indexer_cm_f16 || - pipeline == ctx->device->pipeline_lightning_indexer_decode_cm_f16) { + bool fork_li = false; + for (size_t nhi = 0; nhi < LI_NH_COUNT && !fork_li; ++nhi) { + fork_li = pipeline == ctx->device->pipeline_lightning_indexer_f16[nhi] || + pipeline == ctx->device->pipeline_lightning_indexer_cm_f16[nhi] || + pipeline == ctx->device->pipeline_lightning_indexer_cm_small_f16[nhi] || + pipeline == ctx->device->pipeline_lightning_indexer_decode_cm_f16[nhi]; + } + if (fork_li) { ggml_vk_lightning_indexer_cm(ctx, subctx, dst, pipeline); return; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp index 6747feba505e..a2004bd93d7f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp @@ -8,6 +8,10 @@ #extension GL_KHR_shader_subgroup_basic : require layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +// Head count is a SPECIALIZATION constant, not a push constant: the head loop's trip +// count must stay visible to the compiler. As a push constant it cannot unroll and +// decode costs ~55% more (measured, gfx1151, kv=131584). +layout(constant_id = 1) const uint N_HEAD = 64; #if N_WAVES == 1 layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; #else @@ -40,7 +44,6 @@ layout(push_constant) uniform Parameters { const uint TILE = 16; const uint HEAD_SIZE = 128; -const uint N_HEAD = 64; const uint VEC_PER_HEAD = HEAD_SIZE / 4; const uint TILE_STRIDE = VEC_PER_HEAD + 2; const uint SCORE_STRIDE = TILE / 4 + 1; @@ -105,7 +108,7 @@ void main() { const uint d4 = head_idx % VEC_PER_HEAD; const uint token = token_base + token_local; f16vec4 value = f16vec4(0.0); - if (token < p.n_batch) { + if (token < p.n_batch && (N_HEAD % HEADS_PER_TILE == 0 || head_base + head_local < N_HEAD)) { const uint offset = stream * p.nbq3 + token * p.nbq2 + (head_base + head_local) * p.nbq1 + d4 * 4; value = f16vec4(data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); } @@ -115,7 +118,8 @@ void main() { const uint head_local = tid / TILE; const uint token_local = tid % TILE; const uint token = token_base + token_local; - weight_sh[head_local][token_local] = token < p.n_batch ? data_w[stream * p.nbw3 + token * p.nbw1 + head_base + head_local] : 0.0; + weight_sh[head_local][token_local] = (token < p.n_batch && (N_HEAD % HEADS_PER_TILE == 0 || head_base + head_local < N_HEAD)) + ? data_w[stream * p.nbw3 + token * p.nbw1 + head_base + head_local] : 0.0; } barrier(); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp index 24d1b98d571c..5d32e51361b3 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp @@ -7,6 +7,10 @@ #extension GL_KHR_memory_scope_semantics : require layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +// Head count is a SPECIALIZATION constant, not a push constant: the head loop's trip +// count must stay visible to the compiler. As a push constant it cannot unroll and +// decode costs ~55% more (measured, gfx1151, kv=131584). +layout(constant_id = 1) const uint N_HEAD = 64; layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; layout(binding = 0) readonly buffer QBuf { float data_q[]; }; @@ -35,7 +39,6 @@ layout(push_constant) uniform Parameters { const uint TILE = 16; const uint HEAD_SIZE = 128; -const uint N_HEAD = 64; const uint VEC_PER_HEAD = HEAD_SIZE / 4; const uint TILE_STRIDE = VEC_PER_HEAD + 2; const uint SCORE_STRIDE = TILE / 4 + 1; @@ -80,9 +83,19 @@ void main() { const uint head_local = idx / VEC_PER_HEAD; const uint d4 = idx % VEC_PER_HEAD; const uint head = head_base + head_local; - const uint offset = stream * p.nbq3 + token * p.nbq2 + head * p.nbq1 + d4 * 4; - q_sh[head_local * TILE_STRIDE + d4] = - f16vec4(data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); + // n_head need not be a multiple of TILE (Qwen4-exp indexers are nh=4), so the tail + // tile stages zeros rather than reading past the end of Q. Zero q gives score 0, + // and relu(0) * w = 0, so the padded heads drop out of the sum on their own. + // The guard exists only when N_HEAD is not tile-aligned (nh=4). The condition is a + // specialization-constant expression, so at nh=64 it folds to true and this compiles + // to the original unguarded load - measured: relying on the compiler to range-prove + // head < N_HEAD instead costs ~45% at nh=64. + f16vec4 value = f16vec4(0.0); + if (N_HEAD % TILE == 0 || head < N_HEAD) { + const uint offset = stream * p.nbq3 + token * p.nbq2 + head * p.nbq1 + d4 * 4; + value = f16vec4(data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); + } + q_sh[head_local * TILE_STRIDE + d4] = value; } barrier(); @@ -100,8 +113,13 @@ void main() { if (tid < TILE && kv_base + tid < p.n_kv) { [[unroll]] for (uint head_local = 0; head_local < TILE; ++head_local) { + // Padded heads staged q as zero, so their score is 0 and the term dies in the + // multiply. Only the weight READ needs bounding, and clamping the index does that + // without a branch - a dynamic break here forces the loop rolled and costs ~60%. + const uint head = N_HEAD % TILE == 0 ? head_base + head_local + : min(head_base + head_local, N_HEAD - 1); const float score = score_sh[tid * SCORE_STRIDE + head_local / 4][head_local % 4]; - const float weight = data_w[stream * p.nbw3 + token * p.nbw1 + head_base + head_local]; + const float weight = data_w[stream * p.nbw3 + token * p.nbw1 + head]; total += max(score, 0.0) * weight; } } From 46b7ed7cb631a3be1c58f7bff4196d965959213e Mon Sep 17 00:00:00 2001 From: Nathan Wilson Date: Sun, 30 Aug 2026 09:02:36 +0000 Subject: [PATCH 109/109] vulkan: N_HEAD spec constant for the scalar-64 lightning indexer The scalar64 shader kept its hardcoded 64-head loop when the pipelines went per-head-count; nh 4 and 32 would have run with the wrong trip count. Assisted-by: Claude Fable 5 --- .../vulkan-shaders/lightning_indexer_scalar64.comp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp index 693bb3ece8e8..83fd923fe738 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp @@ -7,6 +7,10 @@ #extension GL_KHR_shader_subgroup_basic : require layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +// Head count is a SPECIALIZATION constant, not a push constant: the head loop's trip +// count must stay visible to the compiler. As a push constant it cannot unroll and +// decode costs ~55% more (measured, gfx1151, kv=131584). +layout(constant_id = 1) const uint N_HEAD = 64; layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; layout(binding = 0) readonly buffer QBuf { float data_q[]; }; @@ -34,7 +38,6 @@ layout(push_constant) uniform Parameters { } p; const uint K_PER_GROUP = 8; -const uint N_HEAD = 64; void main() { const uint lane = gl_SubgroupInvocationID;