From 83d385b4294a27fd5b676b56531f446480bfea2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= <1629204+CISC@users.noreply.github.com> Date: Sat, 27 Jun 2026 09:30:19 +0200 Subject: [PATCH 01/17] tests : fix test-chat-template --no-common option (#25075) --- tests/test-chat-template.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test-chat-template.cpp b/tests/test-chat-template.cpp index c388dee1c..d971b2374 100644 --- a/tests/test-chat-template.cpp +++ b/tests/test-chat-template.cpp @@ -135,7 +135,7 @@ int main(int argc, char ** argv) { output_path = args[i + 1]; i++; } else if (args[i] == "--no-common") { - use_common = true; + use_common = false; } else if (tmpl_path.empty()) { tmpl_path = args[i]; } else { From 0275c0f8007df6284bf5f8156169102612e95462 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= <1629204+CISC@users.noreply.github.com> Date: Sat, 27 Jun 2026 09:30:56 +0200 Subject: [PATCH 02/17] ci : add windows-openvino to check-release (#25022) --- .github/workflows/release.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index eb7e1f20d..616fca3da 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -552,6 +552,9 @@ jobs: name: llama-bin-ubuntu-openvino-${{ env.OPENVINO_VERSION_MAJOR }}-x64.tar.gz windows-openvino: + needs: [check-release] + if: ${{ needs.check-release.outputs.should_release == 'true' }} + runs-on: windows-2022 outputs: From c299a92c38b6de6a1139617652b66081828648db Mon Sep 17 00:00:00 2001 From: Christian Kastner Date: Sat, 27 Jun 2026 09:31:29 +0200 Subject: [PATCH 03/17] binaries : Improve rpc-server and export-graph-ops names. (#25045) Tests are generally prefixed with -test, so rename export-graph-ops accordingly. rpc-server is probably too generic a name for /usr/bin. Because it should work with any ggml application, it is renamed to ggml-rpc-server. --- SECURITY.md | 2 +- tests/CMakeLists.txt | 8 ++--- tests/test-backend-ops.cpp | 2 +- ...raph-ops.cpp => test-export-graph-ops.cpp} | 2 +- tools/rpc/CMakeLists.txt | 2 +- tools/rpc/README.md | 34 +++++++++---------- 6 files changed, 25 insertions(+), 25 deletions(-) rename tests/{export-graph-ops.cpp => test-export-graph-ops.cpp} (98%) diff --git a/SECURITY.md b/SECURITY.md index a98b8e70b..0e704e328 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -80,7 +80,7 @@ To protect sensitive data from potential leaks or unauthorized access, it is cru ### Untrusted environments or networks If you can't run your models in a secure and isolated environment or if it must be exposed to an untrusted network, make sure to take the following security precautions: -* Do not use the RPC backend, [rpc-server](https://github.com/ggml-org/llama.cpp/tree/master/tools/rpc) and [llama-server](https://github.com/ggml-org/llama.cpp/tree/master/tools/server) functionality (see https://github.com/ggml-org/llama.cpp/pull/13061). +* Do not use the RPC backend, [ggml-rpc-server](https://github.com/ggml-org/llama.cpp/tree/master/tools/rpc) and [llama-server](https://github.com/ggml-org/llama.cpp/tree/master/tools/server) functionality (see https://github.com/ggml-org/llama.cpp/pull/13061). * Confirm the hash of any downloaded artifact (e.g. pre-trained model weights) matches a known-good value. * Encrypt your data if sending it over the network. diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 0dd1d7b16..24592a279 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -302,9 +302,9 @@ target_link_libraries(${TEST_TARGET} PRIVATE llama) llama_build_and_test(test-alloc.cpp) target_include_directories(test-alloc PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src) -llama_build(export-graph-ops.cpp) -target_include_directories(export-graph-ops PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src) +llama_build(test-export-graph-ops.cpp) +target_include_directories(test-export-graph-ops PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src) if (TARGET gguf-model-data) - target_link_libraries(export-graph-ops PRIVATE gguf-model-data) - target_compile_definitions(export-graph-ops PRIVATE LLAMA_HF_FETCH) + target_link_libraries(test-export-graph-ops PRIVATE gguf-model-data) + target_compile_definitions(test-export-graph-ops PRIVATE LLAMA_HF_FETCH) endif() diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 0830dbf57..09ac62a75 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9943,7 +9943,7 @@ static void usage(char ** argv) { printf(" --output specifies output format (default: console, options: console, sql, csv)\n"); printf(" --list-ops lists all available GGML operations\n"); printf(" --show-coverage shows test coverage\n"); - printf(" --test-file reads test operators from a test file generated by llama-export-graph-ops\n"); + printf(" --test-file reads test operators from a test file generated by test-export-graph-ops\n"); printf(" -j runs tests using parallel worker threads (default: 1, test mode only)\n"); } diff --git a/tests/export-graph-ops.cpp b/tests/test-export-graph-ops.cpp similarity index 98% rename from tests/export-graph-ops.cpp rename to tests/test-export-graph-ops.cpp index 64cf6dcea..7d8118dcd 100644 --- a/tests/export-graph-ops.cpp +++ b/tests/test-export-graph-ops.cpp @@ -185,7 +185,7 @@ int main(int argc, char ** argv) { return 1; } #else - LOG_ERR("export-graph-ops compiled without HF fetch support\n"); + LOG_ERR("test-export-graph-ops compiled without HF fetch support\n"); return 1; #endif } diff --git a/tools/rpc/CMakeLists.txt b/tools/rpc/CMakeLists.txt index 20f114ad9..0eee9a922 100644 --- a/tools/rpc/CMakeLists.txt +++ b/tools/rpc/CMakeLists.txt @@ -1,4 +1,4 @@ -set(TARGET rpc-server) +set(TARGET ggml-rpc-server) add_executable(${TARGET} rpc-server.cpp) target_link_libraries(${TARGET} PRIVATE ggml) target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/tools/rpc/README.md b/tools/rpc/README.md index 05b7292c0..655b65347 100644 --- a/tools/rpc/README.md +++ b/tools/rpc/README.md @@ -4,8 +4,8 @@ > This example and the RPC backend are currently in a proof-of-concept development stage. As such, the functionality is fragile and > insecure. **Never run the RPC server on an open network or in a sensitive environment!** -The `rpc-server` allows exposing `ggml` devices on a remote host. -The RPC backend communicates with one or several instances of `rpc-server` and offloads computations to them. +The `ggml-rpc-server` allows exposing `ggml` devices on a remote host. +The RPC backend communicates with one or several instances of `ggml-rpc-server` and offloads computations to them. This can be used for distributed LLM inference with `llama.cpp` in the following way: ```mermaid @@ -14,15 +14,15 @@ flowchart TD rpcb<-->|TCP|srvb rpcb<-.->|TCP|srvn subgraph hostn[Host N] - srvn[rpc-server]<-.->dev4["CUDA0"] - srvn[rpc-server]<-.->dev5["CPU"] + srvn[ggml-rpc-server]<-.->dev4["CUDA0"] + srvn[ggml-rpc-server]<-.->dev5["CPU"] end subgraph hostb[Host B] - srvb[rpc-server]<-->dev3["Metal"] + srvb[ggml-rpc-server]<-->dev3["Metal"] end subgraph hosta[Host A] - srva[rpc-server]<-->dev["CUDA0"] - srva[rpc-server]<-->dev2["CUDA1"] + srva[ggml-rpc-server]<-->dev["CUDA0"] + srva[ggml-rpc-server]<-->dev2["CUDA1"] end subgraph host[Main Host] local["Local devices"]<-->ggml[llama-cli] @@ -33,7 +33,7 @@ flowchart TD class local,dev,dev2,dev3,dev4,dev5 devcls ``` -By default, `rpc-server` exposes all available accelerator devices on the host. +By default, `ggml-rpc-server` exposes all available accelerator devices on the host. If there are no accelerators, it exposes a single `CPU` device. ## Usage @@ -41,7 +41,7 @@ If there are no accelerators, it exposes a single `CPU` device. ### Remote hosts On each remote host, build the backends for each accelerator by adding `-DGGML_RPC=ON` to the build options. -For example, to build the `rpc-server` with support for CUDA accelerators: +For example, to build the `ggml-rpc-server` with support for CUDA accelerators: ```bash mkdir build-rpc-cuda @@ -50,10 +50,10 @@ cmake .. -DGGML_CUDA=ON -DGGML_RPC=ON cmake --build . --config Release ``` -When started, the `rpc-server` will detect and expose all available `CUDA` devices: +When started, the `ggml-rpc-server` will detect and expose all available `CUDA` devices: ```bash -$ bin/rpc-server +$ bin/ggml-rpc-server ggml_cuda_init: GGML_CUDA_FORCE_MMQ: no ggml_cuda_init: GGML_CUDA_FORCE_CUBLAS: no ggml_cuda_init: found 1 CUDA devices: @@ -67,14 +67,14 @@ Devices: You can control the set of exposed CUDA devices with the `CUDA_VISIBLE_DEVICES` environment variable or the `--device` command line option. The following two commands have the same effect: ```bash -$ CUDA_VISIBLE_DEVICES=0 bin/rpc-server -p 50052 -$ bin/rpc-server --device CUDA0 -p 50052 +$ CUDA_VISIBLE_DEVICES=0 bin/ggml-rpc-server -p 50052 +$ bin/ggml-rpc-server --device CUDA0 -p 50052 ``` ### Main host On the main host build `llama.cpp` with the backends for the local devices and add `-DGGML_RPC=ON` to the build options. -Finally, when running `llama-cli` or `llama-server`, use the `--rpc` option to specify the host and port of each `rpc-server`: +Finally, when running `llama-cli` or `llama-server`, use the `--rpc` option to specify the host and port of each `ggml-rpc-server`: ```bash $ llama-cli -hf ggml-org/gemma-3-1b-it-GGUF -ngl 99 --rpc 192.168.88.10:50052,192.168.88.11:50052 @@ -90,7 +90,7 @@ This can speed up model loading significantly, especially when using large model To enable the cache, use the `-c` option: ```bash -$ bin/rpc-server -c +$ bin/ggml-rpc-server -c ``` By default, the cache is stored in the `$HOME/.cache/llama.cpp/rpc` directory and can be controlled via the `LLAMA_CACHE` environment variable. @@ -103,8 +103,8 @@ RDMA is enabled by default when `libibverbs` is found at build time. ### Troubleshooting -Use the `GGML_RPC_DEBUG` environment variable to enable debug messages from `rpc-server`: +Use the `GGML_RPC_DEBUG` environment variable to enable debug messages from `ggml-rpc-server`: ```bash -$ GGML_RPC_DEBUG=1 bin/rpc-server +$ GGML_RPC_DEBUG=1 bin/ggml-rpc-server ``` From 0b6529d818608304b277c1db668e8799a467c32d Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Sat, 27 Jun 2026 10:57:31 +0200 Subject: [PATCH 04/17] vulkan: fix step operator for 0 input (#25036) --- ggml/src/ggml-vulkan/vulkan-shaders/unary.comp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/unary.comp b/ggml/src/ggml-vulkan/vulkan-shaders/unary.comp index c62bce825..5ee5275d2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/unary.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/unary.comp @@ -42,7 +42,7 @@ float op_leaky_relu(float x) { } float op_step(float x) { - return x >= 0.0f ? 1.0f : 0.0f; + return x > 0.0f ? 1.0f : 0.0f; } float op_tanh(float x) { From 9bebfcb4bc8b12a316e96ae03f33671eac1e72fd Mon Sep 17 00:00:00 2001 From: Neo Zhang Date: Sat, 27 Jun 2026 17:13:43 +0800 Subject: [PATCH 05/17] sycl : fix failed ut cases of norm (#25044) --- ggml/src/ggml-sycl/norm.cpp | 150 ++++++++++++++++++++++++------------ 1 file changed, 102 insertions(+), 48 deletions(-) diff --git a/ggml/src/ggml-sycl/norm.cpp b/ggml/src/ggml-sycl/norm.cpp index 09fce1280..c4472e4bd 100644 --- a/ggml/src/ggml-sycl/norm.cpp +++ b/ggml/src/ggml-sycl/norm.cpp @@ -2,8 +2,10 @@ #include "ggml-sycl/common.hpp" #include "ggml-sycl/presets.hpp" -static void norm_f32(const float* x, float* dst, const int ncols, const int64_t stride_row, const int64_t stride_channel, - const int64_t stride_sample, const float eps, const sycl::nd_item<3>& item_ct1, sycl::float2* s_sum, int block_size) { +static void norm_f32(const float* x, float* dst, const int ncols, + const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample, + const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, + const float eps, const sycl::nd_item<3>& item_ct1, sycl::float2* s_sum, int block_size) { const int nrows = item_ct1.get_group_range(2); const int nchannels = item_ct1.get_group_range(1); @@ -16,16 +18,16 @@ static void norm_f32(const float* x, float* dst, const int ncols, const int64_t const int tid = item_ct1.get_local_id(2); const int nwarps = nthreads / WARP_SIZE; - const auto strided_offset = calculate_offset<3>({stride_sample, stride_channel, stride_row}, {sample, channel, row}); - const auto packed_offset = calculate_offset<3>({nchannels * nrows * ncols, nrows * ncols, ncols}, {sample, channel, row}); + const auto src_offset = calculate_offset<3>({src_stride_sample, src_stride_channel, src_stride_row}, {sample, channel, row}); + const auto dst_offset = calculate_offset<3>({dst_stride_sample, dst_stride_channel, dst_stride_row}, {sample, channel, row}); - x += strided_offset; - dst += packed_offset; + x += src_offset; + dst += dst_offset; sycl::float2 mean_var = sycl::float2(0.f, 0.f); for (int col = tid; col < ncols; col += block_size) { - const float xi = x[col]; + const float xi = x[col * src_stride_col]; mean_var.x() += xi; mean_var.y() += xi * xi; } @@ -54,7 +56,7 @@ static void norm_f32(const float* x, float* dst, const int ncols, const int64_t const float inv_std = sycl::rsqrt(var + eps); for (int col = tid; col < ncols; col += block_size) { - dst[col] = (x[col] - mean) * inv_std; + dst[col * dst_stride_col] = (x[col * src_stride_col] - mean) * inv_std; } } @@ -145,8 +147,10 @@ static void group_norm_f32(const float* x, float* dst, const int group_size, con } } -static void rms_norm_f32(const float* x, float* dst, const int ncols, const int64_t stride_row, const int64_t stride_channel, - const int64_t stride_sample, const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, int block_size) { +static void rms_norm_f32(const float* x, float* dst, const int ncols, + const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample, + const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, + const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, int block_size) { const int nrows = item_ct1.get_group_range(2); const int nchannels = item_ct1.get_group_range(1); @@ -160,17 +164,17 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols, const int6 const int tid = item_ct1.get_local_id(2); const int nwarps = nthreads / WARP_SIZE; - const auto strided_offset = calculate_offset<3>({stride_sample, stride_channel, stride_row}, {sample, channel, row}); - const auto packed_offset = calculate_offset<3>({nchannels * nrows * ncols, nrows * ncols, ncols}, {sample, channel, row}); + const auto src_offset = calculate_offset<3>({src_stride_sample, src_stride_channel, src_stride_row}, {sample, channel, row}); + const auto dst_offset = calculate_offset<3>({dst_stride_sample, dst_stride_channel, dst_stride_row}, {sample, channel, row}); - x += strided_offset; - dst += packed_offset; + x += src_offset; + dst += dst_offset; float tmp = 0.0f; // partial sum for thread in warp for (int col = tid; col < ncols; col += block_size) { - const float xi = x[col]; + const float xi = x[col * src_stride_col]; tmp += xi * xi; } @@ -198,14 +202,15 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols, const int6 const float scale = sycl::rsqrt(mean + eps); for (int col = tid; col < ncols; col += block_size) { - dst[col] = scale * x[col]; + dst[col * dst_stride_col] = scale * x[col * src_stride_col]; } } template static void l2_norm_f32(const float * x, float * dst, const int ncols, - const int64_t stride_row, const int64_t stride_channel, - const int64_t stride_sample, const float eps, + const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, + const int64_t src_stride_sample, const int64_t dst_stride_col, const int64_t dst_stride_row, + const int64_t dst_stride_channel, const int64_t dst_stride_sample, const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, const int block_size) { const int nrows = item_ct1.get_group_range(2); const int nchannels = item_ct1.get_group_range(1); @@ -215,13 +220,13 @@ static void l2_norm_f32(const float * x, float * dst, const int ncols, const int sample = item_ct1.get_group(0); const int tid = item_ct1.get_local_id(2); - x += sample*stride_sample + channel*stride_channel + row*stride_row; - dst += ((sample*nchannels + channel)*nrows + row)*ncols; + x += sample*src_stride_sample + channel*src_stride_channel + row*src_stride_row; + dst += sample*dst_stride_sample + channel*dst_stride_channel + row*dst_stride_row; float tmp = 0.0f; // partial sum for thread in warp for (int col = tid; col < ncols; col += block_size) { - const float xi = x[col]; + const float xi = x[col * src_stride_col]; tmp += xi * xi; } @@ -229,12 +234,13 @@ static void l2_norm_f32(const float * x, float * dst, const int ncols, const float scale = sycl::rsqrt(sycl::fmax(tmp, eps * eps)); for (int col = tid; col < ncols; col += block_size) { - dst[col] = scale * x[col]; + dst[col * dst_stride_col] = scale * x[col * src_stride_col]; } } static void norm_f32_sycl(const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, - const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, + const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample, + const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, const float eps, queue_ptr stream, int device) { const sycl::range<3> global_dims(nsamples, nchannels, nrows); @@ -245,7 +251,10 @@ static void norm_f32_sycl(const float * x, float * dst, const int ncols, const i sycl::nd_range<3>(global_dims * block_dims, block_dims), [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - norm_f32(x, dst, ncols, stride_row, stride_channel, stride_sample, eps, item_ct1, nullptr, WARP_SIZE); + norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, nullptr, WARP_SIZE); }); }); } @@ -265,7 +274,10 @@ static void norm_f32_sycl(const float * x, float * dst, const int ncols, const i sycl::nd_range<3>(global_dims * block_dims, block_dims), [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - norm_f32(x, dst, ncols, stride_row, stride_channel, stride_sample, eps, item_ct1, get_pointer(s_sum_acc_ct1), work_group_size); + norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, get_pointer(s_sum_acc_ct1), work_group_size); }); }); } @@ -319,7 +331,9 @@ static void group_norm_f32_sycl(const float* x, float* dst, } static void rms_norm_f32_sycl(const float* x, float* dst, const int ncols, const int nrows, const int nchannels, const int nsamples, - const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, queue_ptr stream, int device) { + const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample, + const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, + const float eps, queue_ptr stream, int device) { // printf("%s ncols=%d, nrows=%d, WARP_SIZE=%d\n", __func__, ncols, nrows, WARP_SIZE); const sycl::range<3> global_dims(nsamples, nchannels, nrows); @@ -330,7 +344,10 @@ static void rms_norm_f32_sycl(const float* x, float* dst, const int ncols, const sycl::nd_range<3>(global_dims * block_dims, block_dims), [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - rms_norm_f32(x, dst, ncols, stride_row, stride_channel, stride_sample, eps, item_ct1, nullptr, WARP_SIZE); + rms_norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, nullptr, WARP_SIZE); }); }); } @@ -350,7 +367,10 @@ static void rms_norm_f32_sycl(const float* x, float* dst, const int ncols, const sycl::nd_range<3>(global_dims * block_dims, block_dims), [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { - rms_norm_f32(x, dst, ncols, stride_row, stride_channel, stride_sample, eps, item_ct1, get_pointer(s_sum_acc_ct1), work_group_size); + rms_norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, get_pointer(s_sum_acc_ct1), work_group_size); }); }); } @@ -363,9 +383,14 @@ static void l2_norm_f32_sycl(const float * x, const int nrows, const int nchannels, const int nsamples, - const int64_t stride_row, - const int64_t stride_channel, - const int64_t stride_sample, + const int64_t src_stride_col, + const int64_t src_stride_row, + const int64_t src_stride_channel, + const int64_t src_stride_sample, + const int64_t dst_stride_col, + const int64_t dst_stride_row, + const int64_t dst_stride_channel, + const int64_t dst_stride_sample, const float eps, queue_ptr stream, int device) { @@ -379,7 +404,10 @@ static void l2_norm_f32_sycl(const float * x, block_dims), [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(warp_size)]] { - l2_norm_f32(x, dst, ncols, stride_row, stride_channel, stride_sample, eps, item_ct1, + l2_norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, nullptr, warp_size); }); }); @@ -398,7 +426,9 @@ static void l2_norm_f32_sycl(const float * x, block_dims), [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(warp_size)]] { - l2_norm_f32(x, dst, ncols, stride_row, stride_channel, stride_sample, + l2_norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, eps, item_ct1, get_pointer(s_sum_acc_ct1), work_group_size); }); }); @@ -421,12 +451,20 @@ void ggml_sycl_op_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { memcpy(&eps, dst->op_params, sizeof(float)); GGML_ASSERT(eps >= 0.0f); const size_t ts0 = ggml_type_size(src0->type); - GGML_ASSERT(nb00 == ts0); - const int64_t s01 = nb01 / ts0; - const int64_t s02 = nb02 / ts0; - const int64_t s03 = nb03 / ts0; + const size_t tdst = ggml_type_size(dst->type); + GGML_ASSERT(nb00 % ts0 == 0 && nb01 % ts0 == 0 && nb02 % ts0 == 0 && nb03 % ts0 == 0); + GGML_ASSERT(nb0 % tdst == 0 && nb1 % tdst == 0 && nb2 % tdst == 0 && nb3 % tdst == 0); + const int64_t ss0 = nb00 / ts0; + const int64_t ss1 = nb01 / ts0; + const int64_t ss2 = nb02 / ts0; + const int64_t ss3 = nb03 / ts0; + const int64_t ds0 = nb0 / tdst; + const int64_t ds1 = nb1 / tdst; + const int64_t ds2 = nb2 / tdst; + const int64_t ds3 = nb3 / tdst; - norm_f32_sycl(src0_dd, dst_dd, ne00, ne01, ne02, ne03, s01, s02, s03, eps, main_stream, ctx.device); + norm_f32_sycl(src0_dd, dst_dd, ne00, ne01, ne02, ne03, + ss0, ss1, ss2, ss3, ds0, ds1, ds2, ds3, eps, main_stream, ctx.device); } void ggml_sycl_op_group_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { @@ -465,11 +503,19 @@ void ggml_sycl_op_rms_norm(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { GGML_TENSOR_UNARY_OP_LOCALS const size_t ts0 = ggml_type_size(src0->type); - GGML_ASSERT(nb00 == ts0); - const int64_t s01 = nb01 / ts0; - const int64_t s02 = nb02 / ts0; - const int64_t s03 = nb03 / ts0; - rms_norm_f32_sycl(src0_dd, dst_dd, ne00, ne01, ne02, ne03, s01, s02, s03, eps, main_stream, ctx.device); + const size_t tdst = ggml_type_size(dst->type); + GGML_ASSERT(nb00 % ts0 == 0 && nb01 % ts0 == 0 && nb02 % ts0 == 0 && nb03 % ts0 == 0); + GGML_ASSERT(nb0 % tdst == 0 && nb1 % tdst == 0 && nb2 % tdst == 0 && nb3 % tdst == 0); + const int64_t ss0 = nb00 / ts0; + const int64_t ss1 = nb01 / ts0; + const int64_t ss2 = nb02 / ts0; + const int64_t ss3 = nb03 / ts0; + const int64_t ds0 = nb0 / tdst; + const int64_t ds1 = nb1 / tdst; + const int64_t ds2 = nb2 / tdst; + const int64_t ds3 = nb3 / tdst; + rms_norm_f32_sycl(src0_dd, dst_dd, ne00, ne01, ne02, ne03, + ss0, ss1, ss2, ss3, ds0, ds1, ds2, ds3, eps, main_stream, ctx.device); } void ggml_sycl_op_rms_norm_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { @@ -644,13 +690,21 @@ void ggml_sycl_op_l2_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst) { GGML_ASSERT(eps >= 0.0f); const size_t ts0 = ggml_type_size(src0->type); - GGML_ASSERT(nb00 == ts0); - const int64_t s01 = nb01 / ts0; - const int64_t s02 = nb02 / ts0; - const int64_t s03 = nb03 / ts0; + const size_t tdst = ggml_type_size(dst->type); + GGML_ASSERT(nb00 % ts0 == 0 && nb01 % ts0 == 0 && nb02 % ts0 == 0 && nb03 % ts0 == 0); + GGML_ASSERT(nb0 % tdst == 0 && nb1 % tdst == 0 && nb2 % tdst == 0 && nb3 % tdst == 0); + const int64_t ss0 = nb00 / ts0; + const int64_t ss1 = nb01 / ts0; + const int64_t ss2 = nb02 / ts0; + const int64_t ss3 = nb03 / ts0; + const int64_t ds0 = nb0 / tdst; + const int64_t ds1 = nb1 / tdst; + const int64_t ds2 = nb2 / tdst; + const int64_t ds3 = nb3 / tdst; /*support both WARP_SIZE or WARP_32_SIZE in code choose by hardware for better performance */ - l2_norm_f32_sycl(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream, ctx.device); + l2_norm_f32_sycl(src0_d, dst_d, ne00, ne01, ne02, ne03, + ss0, ss1, ss2, ss3, ds0, ds1, ds2, ds3, eps, stream, ctx.device); } From 0ed235ea2c17a19fc8238668653946721ed136fd Mon Sep 17 00:00:00 2001 From: Gaurav Garg Date: Sat, 27 Jun 2026 17:46:21 +0530 Subject: [PATCH 06/17] [CUDA] Added a cudaMemcpy2DAsync fast path to ggml_cuda_cpy (#25057) * [CUDA] Added a cudaMemcpy2DAsync fast path to ggml_cuda_cpy Add a CUDA ggml_cpy fast path for same-type, same-shape strided copies that are just 2D pitched block copies. When tensors are not fully contiguous but each row is contiguous, it now uses cudaMemcpy2DAsync instead of the slow element-wise scalar copy kernel. This fixes the GDN recurrent snapshot update with -np 4, where rollback slots are separated by cache stride gaps. * Add new tests that execute the new optimized strided copy path * Return unsupported for strided copy in OpenVINO, as new tests are failing --- ggml/src/ggml-cuda/cpy.cu | 45 ++++++++++++++++++++++++ ggml/src/ggml-openvino/ggml-openvino.cpp | 4 +++ tests/test-backend-ops.cpp | 35 ++++++++++++++---- 3 files changed, 77 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-cuda/cpy.cu b/ggml/src/ggml-cuda/cpy.cu index 1e625cc1c..eb5eb0eb4 100644 --- a/ggml/src/ggml-cuda/cpy.cu +++ b/ggml/src/ggml-cuda/cpy.cu @@ -386,6 +386,46 @@ static void ggml_cpy_f32_iq4_nl_cuda( (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13); } +// check if a same-type copy reduces to a 2D strided copy (height rows of width +// contiguous bytes), so it can use cudaMemcpy2DAsync instead of the scalar kernel +static bool ggml_cuda_cpy_as_memcpy_2d(const ggml_tensor * src0, const ggml_tensor * src1, + size_t & width, size_t & height, size_t & spitch, size_t & dpitch) { + // require matching shape: a reshaped copy maps elements by flat order, which the + // prefix walk below does not handle + if (src0->type != src1->type || !ggml_are_same_shape(src0, src1)) { + return false; + } + + // grow the contiguous prefix block shared by both tensors + size_t block_nb = ggml_element_size(src0); + int d = 0; + for (; d < GGML_MAX_DIMS; ++d) { + if (src0->nb[d] != block_nb || src1->nb[d] != block_nb) { + break; + } + block_nb *= src0->ne[d]; + } + + // d == 0: nothing contiguous; d == GGML_MAX_DIMS: fully contiguous (handled by memcpy) + if (d == 0 || d == GGML_MAX_DIMS) { + return false; + } + + // dim d carries the rows; everything above it must be a single element + for (int i = d + 1; i < GGML_MAX_DIMS; ++i) { + if (src0->ne[i] != 1) { + return false; + } + } + + width = block_nb; + height = src0->ne[d]; + spitch = src0->nb[d]; + dpitch = src1->nb[d]; + + return spitch >= width && dpitch >= width; +} + void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, ggml_tensor * src1) { const int64_t ne = ggml_nelements(src0); GGML_ASSERT(ne == ggml_nelements(src1)); @@ -421,6 +461,8 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg const bool can_be_transposed = nb01 == (int64_t)ggml_element_size(src0) && src0->ne[3] == 1 && nb02 == ne00 * ne01 * (int64_t)ggml_element_size(src0); + size_t mc_width = 0, mc_height = 0, mc_spitch = 0, mc_dpitch = 0; + if (src0->type == src1->type && contiguous_srcs) { GGML_ASSERT(ggml_nbytes(src0) == ggml_nbytes(src1)); #if defined(GGML_USE_MUSA) && defined(GGML_MUSA_MUDNN_COPY) @@ -431,6 +473,9 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg { CUDA_CHECK(cudaMemcpyAsync(src1_ddc, src0_ddc, ggml_nbytes(src0), cudaMemcpyDeviceToDevice, main_stream)); } + } else if (ggml_cuda_cpy_as_memcpy_2d(src0, src1, mc_width, mc_height, mc_spitch, mc_dpitch)) { + CUDA_CHECK(cudaMemcpy2DAsync(src1_ddc, mc_dpitch, src0_ddc, mc_spitch, + mc_width, mc_height, cudaMemcpyDeviceToDevice, main_stream)); } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) { if (can_be_transposed) { ggml_cpy_scalar_cuda diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index 943aef864..659dbd4b5 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -1053,6 +1053,10 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { (op->ne[0] == 2 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2)) { return true; } + // CPY into a strided view of a larger buffer (recurrent-state snapshots) not supported + if (op->view_src && ggml_nbytes(op) != ggml_nbytes(op->view_src)) { + return true; + } break; } case GGML_OP_MUL_MAT: { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 09ac62a75..15b50209c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2890,12 +2890,17 @@ struct test_cpy : public test_case { const std::array ne_dst; const std::array permute_src; const std::array permute_dst; + const std::array dst_alloc; // if set, dst is a view into a larger buffer (strided) bool _src_use_permute; bool _dst_use_permute; bool _src_transpose; bool _use_dst_shape; + bool _use_dst_alloc; std::string vars() override { + if (_use_dst_alloc) { + return VARS_TO_STR8(type_src, type_dst, ne_src, ne_dst, permute_src, permute_dst, _src_transpose, dst_alloc); + } if (_use_dst_shape) { return VARS_TO_STR7(type_src, type_dst, ne_src, ne_dst, permute_src, permute_dst, _src_transpose); } @@ -2943,12 +2948,15 @@ struct test_cpy : public test_case { std::array ne_dst = {-1, -1, -1, -1}, std::array permute_src = {0, 0, 0, 0}, std::array permute_dst = {0, 0, 0, 0}, - bool transpose_src = false) + bool transpose_src = false, + std::array dst_alloc = {0, 0, 0, 0}) : type_src(type_src), type_dst(type_dst), ne_src(ne_src), ne_dst(ne_dst), permute_src(permute_src), permute_dst(permute_dst), + dst_alloc(dst_alloc), _src_use_permute(permute_src[0] + permute_src[1] + permute_src[2] + permute_src[3] > 0), _dst_use_permute(permute_dst[0] + permute_dst[1] + permute_dst[2] + permute_dst[3] > 0), _src_transpose(transpose_src), - _use_dst_shape(ne_dst[0] >= 0 && ne_dst[1] >= 0 && ne_dst[2] >= 0 && ne_dst[3] >= 0){} + _use_dst_shape(ne_dst[0] >= 0 && ne_dst[1] >= 0 && ne_dst[2] >= 0 && ne_dst[3] >= 0), + _use_dst_alloc(dst_alloc[0] > 0){} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * src = ggml_new_tensor(ctx, type_src, 4, ne_src.data()); @@ -2966,12 +2974,23 @@ struct test_cpy : public test_case { } std::array dst_ne = _use_dst_shape ? ne_dst : std::array{src->ne[0], src->ne[1], src->ne[2], src->ne[3]}; - ggml_tensor * dst = ggml_new_tensor(ctx, type_dst, 4, dst_ne.data()); - ggml_set_name(dst, "dst"); + ggml_tensor * dst; - if (_dst_use_permute) { - dst = ggml_permute(ctx, dst, permute_dst[0], permute_dst[1], permute_dst[2], permute_dst[3]); - ggml_set_name(dst, "dst_permuted"); + if (_use_dst_alloc) { + // view a sub-block of a larger buffer -> strided dst + ggml_tensor * dst_buf = ggml_new_tensor(ctx, type_dst, 4, dst_alloc.data()); + ggml_set_name(dst_buf, "dst_buf"); + dst = ggml_view_4d(ctx, dst_buf, dst_ne[0], dst_ne[1], dst_ne[2], dst_ne[3], + dst_buf->nb[1], dst_buf->nb[2], dst_buf->nb[3], 0); + ggml_set_name(dst, "dst_view"); + } else { + dst = ggml_new_tensor(ctx, type_dst, 4, dst_ne.data()); + ggml_set_name(dst, "dst"); + + if (_dst_use_permute) { + dst = ggml_permute(ctx, dst, permute_dst[0], permute_dst[1], permute_dst[2], permute_dst[3]); + ggml_set_name(dst, "dst_permuted"); + } } ggml_tensor * out = ggml_cpy(ctx, src, dst); @@ -8181,6 +8200,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {256, 1, 4, 1}, {-1,-1,-1,-1}, {1, 2, 0, 3}, {0, 0, 0, 0})); test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {2, 2097121, 1, 1}, {-1,-1,-1,-1}, {1, 0, 2, 3})); test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {2, 2, 524281, 1}, {-1,-1,-1,-1}, {1, 0, 2, 3})); + test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {128, 2, 3, 1}, {128, 2, 3, 1}, {0, 0, 0, 0}, {0, 0, 0, 0}, false, {128, 4, 3, 1})); // strided dst + test_cases.emplace_back(new test_cpy(GGML_TYPE_F16, GGML_TYPE_F16, {128, 2, 3, 1}, {128, 2, 3, 1}, {0, 0, 0, 0}, {0, 0, 0, 0}, false, {128, 4, 3, 1})); // strided dst // CPY - different src/dst shapes (reshaping via CPY) // Use permutations of {3, 5, 7, 32}. Total elements: 3*5*7*32 = 3360. From ebd048fc5e4b43ec4e0b4abe0b9bf66e1724dad0 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang <66336067+wanghqc@users.noreply.github.com> Date: Sat, 27 Jun 2026 15:36:06 -0700 Subject: [PATCH 07/17] opencl: flash attention improvement (#25069) * opencl: rework FA kernel for f16 and f32 * opencl: flash-attention prefill prepass kernels - flash_attn_kv_pad_f16 pads the tail KV tile to a BLOCK_N multiple - flash_attn_mask_pad_f16 pads the matching mask tile - flash_attn_blk_f16 classifies each KV tile per query block as fully masked / mixed / fully unmasked, so the main kernel can skip fully-masked tiles and the mask lookup for fully-unmasked ones * opencl: FA kernels for q4_0 and q8_0 * opencl: `set_rows` for f32 to q8_0/q4_0 * opencl: dequant kernels for q4_0 and q8_0 * opencl: add FA tile tuning table with override * opencl: wire host side for FA * opencl: q4_0 MoE tensors are also SOA'ed * opencl: cosmetic fix * opencl: refactor, also clarify some code paths in comments * opencl: fix inifity for `-cl-finite-math-only` --------- Co-authored-by: Li He --- ggml/src/ggml-opencl/CMakeLists.txt | 3 + ggml/src/ggml-opencl/fa_tune.h | 91 + ggml/src/ggml-opencl/ggml-opencl.cpp | 2039 +++++++++++++++-- ggml/src/ggml-opencl/kernels/cvt.cl | 152 ++ .../src/ggml-opencl/kernels/flash_attn_f16.cl | 115 +- .../src/ggml-opencl/kernels/flash_attn_f32.cl | 111 +- .../ggml-opencl/kernels/flash_attn_f32_f16.cl | 777 ++++++- .../kernels/flash_attn_f32_q4_0.cl | 1041 +++++++++ .../kernels/flash_attn_f32_q8_0.cl | 1049 +++++++++ .../ggml-opencl/kernels/flash_attn_pre_f16.cl | 156 ++ ggml/src/ggml-opencl/kernels/set_rows.cl | 500 ++++ 11 files changed, 5617 insertions(+), 417 deletions(-) create mode 100644 ggml/src/ggml-opencl/fa_tune.h create mode 100644 ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl create mode 100644 ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl create mode 100644 ggml/src/ggml-opencl/kernels/flash_attn_pre_f16.cl diff --git a/ggml/src/ggml-opencl/CMakeLists.txt b/ggml/src/ggml-opencl/CMakeLists.txt index 82ce61d72..09efbc566 100644 --- a/ggml/src/ggml-opencl/CMakeLists.txt +++ b/ggml/src/ggml-opencl/CMakeLists.txt @@ -192,7 +192,10 @@ set(GGML_OPENCL_KERNELS mul_mm_f16_f32_kq_kqv conv2d conv2d_f16_f32 + flash_attn_pre_f16 flash_attn_f32_f16 + flash_attn_f32_q8_0 + flash_attn_f32_q4_0 flash_attn_f16 flash_attn_f32 ) diff --git a/ggml/src/ggml-opencl/fa_tune.h b/ggml/src/ggml-opencl/fa_tune.h new file mode 100644 index 000000000..1e2c6ea7e --- /dev/null +++ b/ggml/src/ggml-opencl/fa_tune.h @@ -0,0 +1,91 @@ +#pragma once + +// Flash-attention per-(dk,dv) tile tuning for the Adreno OpenCL backend. +// Isolated from ggml-opencl.cpp so the tuning numbers are easy to find and +// edit; the FA dispatch and kernel-compile logic stay in the main file. +// This header is a file section — it is #included exactly once, at the point +// in ggml-opencl.cpp where the ggml logging macros are already in scope. + +// Per-(dk, dv) FA config; shared by dispatch and supports_op. +struct ggml_opencl_fa_dim { + int dk; int dv; int bm; int bn; int n_split; int nkv_split_threshold; +}; + +// Split variant fires when n_kv >= threshold (threshold=0 -> always split). +// Default tuning covers Adreno 7xx/8xx mobile and X1-series laptop GPUs. +static const ggml_opencl_fa_dim g_fa_dims_adreno_default[] = { + { 40, 40, 64, 32, 1, 0}, { 64, 64, 64, 32, 2, 64}, + { 80, 80, 64, 32, 2, 64}, { 96, 96, 64, 32, 2, 64}, + {112, 112, 64, 32, 2, 64}, {128, 128, 64, 32, 2, 64}, + {192, 128, 16, 16, 1, 0}, + {192, 192, 16, 16, 1, 0}, + {256, 256, 16, 16, 16, 0}, +}; + +struct ggml_opencl_fa_dim_table { + const ggml_opencl_fa_dim * data; + size_t count; + + const ggml_opencl_fa_dim * begin() const { return data; } + const ggml_opencl_fa_dim * end() const { return data + count; } +}; + +// Mutable copy of the active table; GGML_OPENCL_FA_TUNE patches entries here +// at backend init without touching the const source table. +static ggml_opencl_fa_dim g_fa_dims_runtime[ + sizeof(g_fa_dims_adreno_default) / sizeof(g_fa_dims_adreno_default[0])]; + +static ggml_opencl_fa_dim_table g_opencl_fa_dims = { + g_fa_dims_adreno_default, + sizeof(g_fa_dims_adreno_default) / sizeof(g_fa_dims_adreno_default[0]), +}; + +// GGML_OPENCL_FA_TUNE=dk:dv:bm:bn:nsplit:thr[,…] — patches matching entries +// in the active table at backend init, before the first FA kernel compiles. +// Unmatched (dk,dv) pairs are warned and ignored. +static void ggml_opencl_fa_apply_env_overrides() { + const char * e = std::getenv("GGML_OPENCL_FA_TUNE"); + if (!e || !e[0]) { + return; + } + + std::string s = e; + size_t pos = 0; + while (pos < s.size()) { + size_t comma = s.find(',', pos); + std::string entry = s.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); + int dk, dv, bm, bn, nsplit, thr; + if (std::sscanf(entry.c_str(), "%d:%d:%d:%d:%d:%d", &dk, &dv, &bm, &bn, &nsplit, &thr) == 6) { + bool patched = false; + for (size_t i = 0; i < g_opencl_fa_dims.count; ++i) { + ggml_opencl_fa_dim & d = g_fa_dims_runtime[i]; + if (d.dk == dk && d.dv == dv) { + d.bm = bm; d.bn = bn; d.n_split = nsplit; d.nkv_split_threshold = thr; + GGML_LOG_INFO("ggml_opencl: FA tune override DK=%d DV=%d -> bm=%d bn=%d n_split=%d thr=%d\n", + dk, dv, bm, bn, nsplit, thr); + patched = true; + break; + } + } + if (!patched) { + GGML_LOG_WARN("ggml_opencl: FA tune override DK=%d DV=%d ignored (no matching dim)\n", dk, dv); + } + } else { + GGML_LOG_WARN("ggml_opencl: FA tune override entry malformed: '%s'\n", entry.c_str()); + } + if (comma == std::string::npos) break; + pos = comma + 1; + } +} + +// Copy the default table into the mutable runtime buffer and apply any +// GGML_OPENCL_FA_TUNE overrides. A per-generation table can be added here +// once it has been tuned on hardware. +static void ggml_cl_init_fa_dims_table() { + const size_t count = sizeof(g_fa_dims_adreno_default) / sizeof(g_fa_dims_adreno_default[0]); + for (size_t i = 0; i < count; ++i) { + g_fa_dims_runtime[i] = g_fa_dims_adreno_default[i]; + } + g_opencl_fa_dims = { g_fa_dims_runtime, count }; + ggml_opencl_fa_apply_env_overrides(); +} diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 00f20b09b..32581901b 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -29,6 +29,8 @@ #include #include #include +#include +#include #undef MIN #undef MAX @@ -53,6 +55,9 @@ //------------------------------------------------------------------------------ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor); +static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor); +static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor); +static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); // See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1. // Precompute mp (m' in the paper) and L such that division @@ -96,6 +101,7 @@ enum ADRENO_GPU_GEN { A7X, A8X, X1E, + X2E, }; enum ADRENO_CL_COMPILER_TYPE { @@ -236,6 +242,10 @@ static ADRENO_GPU_GEN get_adreno_gpu_gen(const char *device_name) { return ADRENO_GPU_GEN::X1E; } + if (strstr(device_name, "X2")) { + return ADRENO_GPU_GEN::X2E; + } + return ADRENO_GPU_GEN::ADRENO_UNKNOWN; } @@ -368,7 +378,7 @@ struct ggml_backend_opencl_device_context { cl_device_type device_type; std::string device_version; - // Initialized by ggml_cl2_init(). + // Initialized by ggml_cl_init(). ggml_backend_opencl_context * backend_ctx = nullptr; // Initialized by ggml_backend_opencl_device_get_buffer_type() @@ -384,6 +394,55 @@ struct ggml_backend_opencl_device_context { size_t global_mem_size = 0; }; +// Lazily-compiled flash-attention kernels and their per-(dk,dv) tile metadata. +// One map per (Q/KV dtype, decode/prefill, split) combination; the int maps +// hold tile dims (bm/bn), workgroup sizes and the n_kv split thresholds. +struct ggml_opencl_fa_kernels { + // f16 Q / f16 KV + std::map, cl_kernel> f16; + std::map, cl_kernel> f16_q1; + // f32 Q / f32 KV + std::map, cl_kernel> f32; + std::map, cl_kernel> f32_q1; + // f32 Q / f16 KV (mixed) + std::map, cl_kernel> f32_f16; + std::map, cl_kernel> f32_f16_split; // N_SPLIT>1 variant + std::map, cl_kernel> f32_f16_q1; + std::map, cl_kernel> f32_f16_q1_split; // flash-decoding K-split + std::map, int> f32_f16_bm; + std::map, int> f32_f16_bn; + std::map, int> f32_f16_wg_size; + std::map, int> f32_f16_split_wg_size; + std::map, int> f32_f16_split_nkv_threshold; + // f32 Q / native q8_0 KV + std::map, cl_kernel> f32_q8_0_q1; // decode + std::map, cl_kernel> f32_q8_0_q1_split; // flash-decoding pass 1 + std::map, cl_kernel> f32_q8_0; // prefill (baseline) + std::map, cl_kernel> f32_q8_0_split; // N_SPLIT>1 variant + std::map, int> f32_q8_0_split_wg_size; // wg_size = bm*n_split + std::map, int> f32_q8_0_split_nkv_threshold; // use split when n_kv >= this + std::map, int> f32_q8_0_split_bm; // per-split BLOCK_M + // f32 Q / native q4_0 KV + std::map, cl_kernel> f32_q4_0_q1; + std::map, cl_kernel> f32_q4_0_q1_split; + std::map, cl_kernel> f32_q4_0; + std::map, cl_kernel> f32_q4_0_split; + std::map, int> f32_q4_0_split_wg_size; + std::map, int> f32_q4_0_split_nkv_threshold; + std::map, int> f32_q4_0_split_bm; + // shared: flash-decoding merge + prefill prepass (kv-pad, mask-pad, blk class) + std::map, cl_kernel> f32_merge; + std::map, cl_kernel> kv_pad_f16; + std::map, cl_kernel> mask_pad_f16; + std::map, cl_kernel> blk_f16; + // generic prefill tile dims (f16 / f32 paths) + std::map, int> bm; + std::map, int> bn; + // attempted (variant, (dk, dv)) + // all attempted FA kernels appear here, but those not registered failed compilation + std::set>> variant_attempted; +}; + // backend context struct ggml_backend_opencl_context { int ref_count; @@ -397,9 +456,6 @@ struct ggml_backend_opencl_context { // argsort is loaded in supports_op because its availability depends on how // many workgroups are allowed, which requires kernel compilation. bool kernels_loaded_argsort = false; - // flash attn is loaded in supports_op because it contains multiple variants - // and takes time to compile, so we want to only compile it when needed. - bool kernels_loaded_flash_attn = false; // rest of the kernels are currently always loaded in alloc_buffer. bool kernels_loaded = false; @@ -414,13 +470,16 @@ struct ggml_backend_opencl_context { size_t max_workgroup_size; bool fp16_support; bool has_vector_subgroup_broadcast; - bool has_qcom_subgroup_shuffle = false; // cl_qcom_subgroup_shuffle + bool has_subgroup_shuffle = false; // cl_khr_subgroup_shuffle or cl_qcom_subgroup_shuffle + bool has_qcom_subgroup_shuffle = false; // specifically cl_qcom_subgroup_shuffle bool disable_fusion; bool adreno_has_large_buffer; bool adreno_use_large_buffer; ggml_cl_compiler_version adreno_cl_compiler_version; + std::string kernel_compile_opts; // cached for lazy-compiled kernels. + int adreno_wave_size; cl_bool non_uniform_workgroups; @@ -546,16 +605,13 @@ struct ggml_backend_opencl_context { cl_kernel kernel_diag_f32; cl_kernel kernel_soft_max, kernel_soft_max_4; cl_kernel kernel_soft_max_f16, kernel_soft_max_4_f16; - std::map, cl_kernel> kernels_flash_attn_f16; - std::map, cl_kernel> kernels_flash_attn_f16_q1; - std::map, cl_kernel> kernels_flash_attn_f32; - std::map, cl_kernel> kernels_flash_attn_f32_q1; - std::map, cl_kernel> kernels_flash_attn_f32_f16; - std::map, cl_kernel> kernels_flash_attn_f32_f16_q1; - std::map, int> kernels_flash_attn_bm; - std::map, int> kernels_flash_attn_bn; + ggml_opencl_fa_kernels fa; cl_kernel kernel_get_rows_f32, kernel_get_rows_f16, kernel_get_rows_q4_0; cl_kernel kernel_set_rows_f32_i64, kernel_set_rows_f32_i32, kernel_set_rows_f16_i64, kernel_set_rows_f16_i32; + cl_kernel kernel_set_rows_q8_0_i64, kernel_set_rows_q8_0_i32; + cl_kernel kernel_set_rows_q8_0_soa_i64, kernel_set_rows_q8_0_soa_i32; + cl_kernel kernel_set_rows_q4_0_i64, kernel_set_rows_q4_0_i32; + cl_kernel kernel_set_rows_q4_0_soa_i64, kernel_set_rows_q4_0_soa_i32; cl_kernel kernel_rope_norm_f32, kernel_rope_norm_f16, kernel_rope_neox_f32, kernel_rope_neox_f16; cl_kernel kernel_rope_multi_f32, kernel_rope_multi_f16, kernel_rope_vision_f32, kernel_rope_vision_f16; cl_kernel kernel_cpy_f16_f16, kernel_cpy_f16_f32, kernel_cpy_f32_f16, kernel_cpy_f32_f32, kernel_cpy_f32_f32_pack, kernel_cpy_i32_i32; @@ -589,6 +645,10 @@ struct ggml_backend_opencl_context { cl_kernel kernel_convert_block_mxfp4, kernel_convert_block_mxfp4_trans, kernel_restore_block_mxfp4, kernel_restore_block_mxfp4_trans; cl_kernel kernel_convert_block_mxfp4_trans4_ns, kernel_restore_block_mxfp4_trans4_ns; cl_kernel kernel_convert_block_q8_0, kernel_restore_block_q8_0, kernel_restore_block_q8_0_trans; + cl_kernel kernel_dequant_q8_0_f16_view_aos; + cl_kernel kernel_dequant_q8_0_f32_view_aos; + cl_kernel kernel_dequant_q4_0_f16_view_aos; + cl_kernel kernel_dequant_q4_0_f32_view_aos; cl_kernel kernel_convert_block_q6_K_noshuffle, kernel_restore_block_q6_K_noshuffle; cl_kernel kernel_convert_bf16_to_f16, kernel_convert_f16_to_bf16; cl_kernel kernel_mul_mat_q4_0_f32_8x_flat; @@ -877,7 +937,13 @@ inline std::string read_file(const std::string &path) { return text; } -static cl_program build_program_from_source(cl_context ctx, cl_device_id dev, const char* program_buffer, const std::string &compile_opts) { +// fatal=false returns NULL on compile failure instead of aborting; used for +// optional FA variants that may exhaust the Adreno compiler at large DK. +static cl_program build_program_from_source_ex(cl_context ctx, cl_device_id dev, const char* program_buffer, const std::string &compile_opts, bool fatal, const char *tag = nullptr) { + if (tag) { + GGML_LOG_INFO("ggml_opencl: compiling %s\n", tag); + } + cl_program p; char *program_log; size_t program_size; @@ -889,7 +955,10 @@ static cl_program build_program_from_source(cl_context ctx, cl_device_id dev, co p = clCreateProgramWithSource(ctx, 1, (const char**)&program_buffer, &program_size, &err); if(err < 0) { GGML_LOG_ERROR("OpenCL error creating program"); - exit(1); + if (fatal) { + exit(1); + } + return nullptr; } err = clBuildProgram(p, 0, NULL, compile_opts.c_str(), NULL, NULL); @@ -898,14 +967,22 @@ static cl_program build_program_from_source(cl_context ctx, cl_device_id dev, co program_log = (char*) malloc(log_size + 1); program_log[log_size] = '\0'; clGetProgramBuildInfo(p, dev, CL_PROGRAM_BUILD_LOG, log_size + 1, program_log, NULL); - GGML_LOG_ERROR("ggml_opencl: kernel compile error:\n\n%s\n", program_log); + GGML_LOG_ERROR("ggml_opencl: kernel compile error (err=%d):\n\n%s\n", err, program_log); free(program_log); - exit(1); + clReleaseProgram(p); + if (fatal) { + exit(1); + } + return nullptr; } return p; } +static cl_program build_program_from_source(cl_context ctx, cl_device_id dev, const char* program_buffer, const std::string &compile_opts) { + return build_program_from_source_ex(ctx, dev, program_buffer, compile_opts, /*fatal=*/true); +} + static void load_cl_kernels_argsort(ggml_backend_opencl_context *backend_ctx) { // compiler options for general kernels auto opencl_c_std = @@ -932,84 +1009,6 @@ static void load_cl_kernels_argsort(ggml_backend_opencl_context *backend_ctx) { } } -static void load_cl_kernels_flash_attn(ggml_backend_opencl_context *backend_ctx) { - // compiler options for general kernels - auto opencl_c_std = - std::string("CL") + std::to_string(backend_ctx->opencl_c_version.major) + "." + std::to_string(backend_ctx->opencl_c_version.minor); - std::string compile_opts = std::string("-cl-std=") + opencl_c_std + - " -cl-mad-enable -cl-unsafe-math-optimizations" - " -cl-finite-math-only -cl-fast-relaxed-math"; - - // flash_attn - if (!backend_ctx->kernels_loaded_flash_attn) { - cl_int err; - - #ifdef GGML_OPENCL_EMBED_KERNELS - const std::string kernel_src_f16 { - #include "flash_attn_f16.cl.h" - }; - const std::string kernel_src_f32 { - #include "flash_attn_f32.cl.h" - }; - const std::string kernel_src_f32_f16 { - #include "flash_attn_f32_f16.cl.h" - }; - #else - const std::string kernel_src_f16 = read_file("flash_attn_f16.cl"); - const std::string kernel_src_f32 = read_file("flash_attn_f32.cl"); - const std::string kernel_src_f32_f16 = read_file("flash_attn_f32_f16.cl"); - #endif - - if (!kernel_src_f16.empty() && !kernel_src_f32.empty() && !kernel_src_f32_f16.empty()) { - const struct { int dk; int dv; int bm; int bn; } fa_dims[] = { - { 40, 40, 32, 32}, { 64, 64, 64, 64}, { 80, 80, 64, 32}, { 96, 96, 64, 32}, - {112, 112, 32, 32}, {128, 128, 32, 32}, {192, 128, 16, 16}, - {192, 192, 16, 16}, {256, 256, 16, 16}, - }; - - for (size_t i = 0; i < sizeof(fa_dims)/sizeof(fa_dims[0]); ++i) { - const int dk = fa_dims[i].dk; - const int dv = fa_dims[i].dv; - const int bm = fa_dims[i].bm; - const int bn = fa_dims[i].bn; - std::string OPTS = compile_opts + - " -D DK=" + std::to_string(dk) + - " -D DV=" + std::to_string(dv) + - " -D BLOCK_M=" + std::to_string(bm) + - " -D BLOCK_N=" + std::to_string(bn); - - cl_program prog_f16 = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src_f16.c_str(), OPTS); - cl_kernel k_f16, k_f16_q1; - CL_CHECK((k_f16 = clCreateKernel(prog_f16, "flash_attn_f16", &err), err)); - CL_CHECK((k_f16_q1 = clCreateKernel(prog_f16, "flash_attn_f16_q1", &err), err)); - backend_ctx->kernels_flash_attn_f16[{dk, dv}] = k_f16; - backend_ctx->kernels_flash_attn_f16_q1[{dk, dv}] = k_f16_q1; - CL_CHECK(clReleaseProgram(prog_f16)); - - cl_program prog_f32 = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src_f32.c_str(), OPTS); - cl_kernel k_f32, k_f32_q1; - CL_CHECK((k_f32 = clCreateKernel(prog_f32, "flash_attn_f32", &err), err)); - CL_CHECK((k_f32_q1 = clCreateKernel(prog_f32, "flash_attn_f32_q1", &err), err)); - backend_ctx->kernels_flash_attn_f32[{dk, dv}] = k_f32; - backend_ctx->kernels_flash_attn_f32_q1[{dk, dv}] = k_f32_q1; - CL_CHECK(clReleaseProgram(prog_f32)); - - cl_program prog_f32_f16 = build_program_from_source(backend_ctx->context, backend_ctx->device, kernel_src_f32_f16.c_str(), OPTS); - cl_kernel k_f32_f16, k_f32_f16_q1; - CL_CHECK((k_f32_f16 = clCreateKernel(prog_f32_f16, "flash_attn_f32_f16", &err), err)); - CL_CHECK((k_f32_f16_q1 = clCreateKernel(prog_f32_f16, "flash_attn_f32_f16_q1", &err), err)); - backend_ctx->kernels_flash_attn_f32_f16[{dk, dv}] = k_f32_f16; - backend_ctx->kernels_flash_attn_f32_f16_q1[{dk, dv}] = k_f32_f16_q1; - CL_CHECK(clReleaseProgram(prog_f32_f16)); - - backend_ctx->kernels_flash_attn_bm[{dk, dv}] = bm; - backend_ctx->kernels_flash_attn_bn[{dk, dv}] = bn; - } - backend_ctx->kernels_loaded_flash_attn = true; - } - } -} - static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { if (backend_ctx->kernels_loaded) { return; @@ -1028,6 +1027,8 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { compile_opts += " -qcom-enable-large-buffer "; } + backend_ctx->kernel_compile_opts = compile_opts; + GGML_LOG_INFO("ggml_opencl: loading OpenCL kernels"); // add @@ -1189,6 +1190,10 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { CL_CHECK((backend_ctx->kernel_convert_block_q8_0 = clCreateKernel(backend_ctx->program_cvt, "kernel_convert_block_q8_0", &err), err)); CL_CHECK((backend_ctx->kernel_restore_block_q8_0 = clCreateKernel(backend_ctx->program_cvt, "kernel_restore_block_q8_0", &err), err)); CL_CHECK((backend_ctx->kernel_restore_block_q8_0_trans = clCreateKernel(backend_ctx->program_cvt, "kernel_restore_block_q8_0_trans", &err), err)); + CL_CHECK((backend_ctx->kernel_dequant_q8_0_f16_view_aos = clCreateKernel(backend_ctx->program_cvt, "kernel_dequant_q8_0_f16_view_aos", &err), err)); + CL_CHECK((backend_ctx->kernel_dequant_q8_0_f32_view_aos = clCreateKernel(backend_ctx->program_cvt, "kernel_dequant_q8_0_f32_view_aos", &err), err)); + CL_CHECK((backend_ctx->kernel_dequant_q4_0_f16_view_aos = clCreateKernel(backend_ctx->program_cvt, "kernel_dequant_q4_0_f16_view_aos", &err), err)); + CL_CHECK((backend_ctx->kernel_dequant_q4_0_f32_view_aos = clCreateKernel(backend_ctx->program_cvt, "kernel_dequant_q4_0_f32_view_aos", &err), err)); CL_CHECK((backend_ctx->kernel_convert_block_q4_K = clCreateKernel(backend_ctx->program_cvt, "kernel_convert_block_q4_K", &err), err)); CL_CHECK((backend_ctx->kernel_restore_block_q4_K = clCreateKernel(backend_ctx->program_cvt, "kernel_restore_block_q4_K", &err), err)); CL_CHECK((backend_ctx->kernel_convert_block_q4_K_noshuffle = clCreateKernel(backend_ctx->program_cvt, "kernel_convert_block_q4_K_noshuffle", &err), err)); @@ -2680,6 +2685,14 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { CL_CHECK((backend_ctx->kernel_set_rows_f32_i32 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_f32_i32", &err), err)); CL_CHECK((backend_ctx->kernel_set_rows_f16_i64 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_f16_i64", &err), err)); CL_CHECK((backend_ctx->kernel_set_rows_f16_i32 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_f16_i32", &err), err)); + CL_CHECK((backend_ctx->kernel_set_rows_q8_0_i64 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_q8_0_i64", &err), err)); + CL_CHECK((backend_ctx->kernel_set_rows_q8_0_i32 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_q8_0_i32", &err), err)); + CL_CHECK((backend_ctx->kernel_set_rows_q8_0_soa_i64 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_q8_0_soa_i64", &err), err)); + CL_CHECK((backend_ctx->kernel_set_rows_q8_0_soa_i32 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_q8_0_soa_i32", &err), err)); + CL_CHECK((backend_ctx->kernel_set_rows_q4_0_i64 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_q4_0_i64", &err), err)); + CL_CHECK((backend_ctx->kernel_set_rows_q4_0_i32 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_q4_0_i32", &err), err)); + CL_CHECK((backend_ctx->kernel_set_rows_q4_0_soa_i64 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_q4_0_soa_i64", &err), err)); + CL_CHECK((backend_ctx->kernel_set_rows_q4_0_soa_i32 = clCreateKernel(backend_ctx->program_set_rows, "kernel_set_rows_q4_0_soa_i32", &err), err)); GGML_LOG_CONT("."); } @@ -3704,13 +3717,470 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { backend_ctx->kernels_loaded = true; } -// XXX static ggml_backend_opencl_context * ggml_cl2_init(ggml_backend_dev_t dev) { -// XXX static bool initialized = false; -// XXX static ggml_backend_opencl_context *backend_ctx = nullptr; - static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev); static bool ggml_opencl_is_device_supported(ggml_backend_dev_t dev); +// FA per-(dk,dv) tile tuning table + GGML_OPENCL_FA_TUNE override parsing. +#include "fa_tune.h" + +// FA variant key for the per-(dk,dv,variant) lazy compile cache. +// Kernel built on first dispatch to reduce kernel loading time. +// NB - a warmup run is recommended to get all necessary FA variants compiled +// before actual runs. +enum ggml_opencl_fa_variant { + FA_VARIANT_PRE = 0, // prepass kernels (kv_pad, mask_pad, blk) + FA_VARIANT_F16 = 1, + FA_VARIANT_F32 = 2, + FA_VARIANT_F32_F16 = 3, + FA_VARIANT_Q8_0 = 4, + FA_VARIANT_Q4_0 = 5, + FA_VARIANT_F32_F16_SPLIT = 6, + FA_VARIANT_Q8_0_SPLIT = 7, + FA_VARIANT_Q4_0_SPLIT = 8, +}; + +static std::string ggml_opencl_fa_kernel_src(ggml_opencl_fa_variant v) { +#ifdef GGML_OPENCL_EMBED_KERNELS + switch (v) { + case FA_VARIANT_F16: + return std::string{ + #include "flash_attn_f16.cl.h" + }; + case FA_VARIANT_F32: + return std::string{ + #include "flash_attn_f32.cl.h" + }; + case FA_VARIANT_F32_F16: + case FA_VARIANT_F32_F16_SPLIT: + return std::string{ + #include "flash_attn_f32_f16.cl.h" + }; + case FA_VARIANT_PRE: + return std::string{ + #include "flash_attn_pre_f16.cl.h" + }; + case FA_VARIANT_Q8_0: + case FA_VARIANT_Q8_0_SPLIT: + return std::string{ + #include "flash_attn_f32_q8_0.cl.h" + }; + case FA_VARIANT_Q4_0: + case FA_VARIANT_Q4_0_SPLIT: + return std::string{ + #include "flash_attn_f32_q4_0.cl.h" + }; + } + return {}; +#else + switch (v) { + case FA_VARIANT_F16: return read_file("flash_attn_f16.cl"); + case FA_VARIANT_F32: return read_file("flash_attn_f32.cl"); + case FA_VARIANT_F32_F16: + case FA_VARIANT_F32_F16_SPLIT: return read_file("flash_attn_f32_f16.cl"); + case FA_VARIANT_PRE: return read_file("flash_attn_pre_f16.cl"); + case FA_VARIANT_Q8_0: + case FA_VARIANT_Q8_0_SPLIT: return read_file("flash_attn_f32_q8_0.cl"); + case FA_VARIANT_Q4_0: + case FA_VARIANT_Q4_0_SPLIT: return read_file("flash_attn_f32_q4_0.cl"); + } + return {}; +#endif +} + +static std::string ggml_opencl_fa_compile_opts(ggml_backend_opencl_context * backend_ctx, + const ggml_opencl_fa_dim * cfg, + ggml_opencl_fa_variant variant) { + std::string opts = backend_ctx->kernel_compile_opts + + " -D DK=" + std::to_string(cfg->dk) + + " -D DV=" + std::to_string(cfg->dv) + + " -D BLOCK_M=" + std::to_string(cfg->bm) + + " -D BLOCK_N=" + std::to_string(cfg->bn); + + const bool is_split = variant == FA_VARIANT_F32_F16_SPLIT || + variant == FA_VARIANT_Q8_0_SPLIT || + variant == FA_VARIANT_Q4_0_SPLIT; + if (is_split) { + opts += " -D N_SPLIT=" + std::to_string(cfg->n_split); + if (backend_ctx->has_subgroup_shuffle) { + opts += backend_ctx->has_qcom_subgroup_shuffle + ? " -D cl_qcom_subgroup_shuffle=1" + : " -D cl_khr_subgroup_shuffle=1"; + } + } + return opts; +} + +// Log private memory for an FA kernel. Enable via `GGML_OPENCL_FA_LOG_SPILL=1`. +// On Adreno non-zero private_mem means spilling to global memory due to resource +// constraint and usually causes performance degradation. +// (per-work-item, no cache locality) — a strong signal to pick a config +// with smaller per-thread state (e.g. larger N_SPLIT). +static void ggml_opencl_log_fa_kernel_spill(ggml_backend_opencl_context * backend_ctx, + cl_kernel kernel, const char * name, int dk, int dv) { + static const bool enabled = []{ + const char * e = std::getenv("GGML_OPENCL_FA_LOG_SPILL"); + return e && e[0] && e[0] != '0'; + }(); + + if (!enabled || kernel == nullptr) { + return; + } + + cl_ulong priv_mem = 0; + if (clGetKernelWorkGroupInfo(kernel, backend_ctx->device, CL_KERNEL_PRIVATE_MEM_SIZE, + sizeof(priv_mem), &priv_mem, NULL) == CL_SUCCESS) { + const char * tag = priv_mem > 0 ? "SPILL" : "ok"; + GGML_LOG_INFO("ggml_opencl: [%s] %s DK=%d DV=%d private_mem=%llu bytes\n", + tag, name, dk, dv, (unsigned long long) priv_mem); + } +} + +static void ggml_opencl_ensure_fa_pre_kernels(ggml_backend_opencl_context * backend_ctx, int dk, int dv) { + const std::pair dk_dv = {dk, dv}; + if (backend_ctx->fa.kv_pad_f16.count(dk_dv) > 0) { + return; + } + + const ggml_opencl_fa_dim * cfg = nullptr; + for (const auto & d : g_opencl_fa_dims) { + if (d.dk == dk && d.dv == dv) { + cfg = &d; break; + } + } + + if (cfg == nullptr) { + GGML_ABORT("ggml_opencl: no flash_attn config for DK=%d DV=%d", dk, dv); + } + + GGML_LOG_INFO("ggml_opencl: lazy-compiling flash_attn prepass for DK=%d DV=%d\n", dk, dv); + + cl_int err; + const std::string src = ggml_opencl_fa_kernel_src(FA_VARIANT_PRE); + const std::string opts = ggml_opencl_fa_compile_opts(backend_ctx, cfg, FA_VARIANT_PRE); + + cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, src.c_str(), opts); + + cl_kernel k_kv_pad_f16, k_mask_pad_f16, k_blk_f16; + CL_CHECK((k_kv_pad_f16 = clCreateKernel(prog, "flash_attn_kv_pad_f16", &err), err)); + CL_CHECK((k_mask_pad_f16 = clCreateKernel(prog, "flash_attn_mask_pad_f16", &err), err)); + CL_CHECK((k_blk_f16 = clCreateKernel(prog, "flash_attn_blk_f16", &err), err)); + backend_ctx->fa.kv_pad_f16[{dk, dv}] = k_kv_pad_f16; + backend_ctx->fa.mask_pad_f16[{dk, dv}] = k_mask_pad_f16; + backend_ctx->fa.blk_f16[{dk, dv}] = k_blk_f16; + CL_CHECK(clReleaseProgram(prog)); + + backend_ctx->fa.f32_f16_bm[{dk, dv}] = cfg->bm; + backend_ctx->fa.f32_f16_bn[{dk, dv}] = cfg->bn; + backend_ctx->fa.f32_f16_wg_size[{dk, dv}] = cfg->bm; + backend_ctx->fa.bm[{dk, dv}] = cfg->bm; + backend_ctx->fa.bn[{dk, dv}] = cfg->bn; +} + +// Compile one (variant, dk, dv); memoised. false = compiler rejected. +static bool ggml_opencl_ensure_fa_variant(ggml_backend_opencl_context * backend_ctx, int dk, int dv, ggml_opencl_fa_variant variant) { + const std::pair dk_dv = {dk, dv}; + + const ggml_opencl_fa_dim * cfg = nullptr; + for (const auto & d : g_opencl_fa_dims) { + if (d.dk == dk && d.dv == dv) { + cfg = &d; break; + } + } + if (cfg == nullptr) { + return false; + } + + // if a variant has already been compiled + switch (variant) { + case FA_VARIANT_F16: { + if (backend_ctx->fa.f16.count(dk_dv)) { + return true; + } + break; + } + case FA_VARIANT_F32: { + if (backend_ctx->fa.f32.count(dk_dv)) { + return true; + } + break; + } + case FA_VARIANT_F32_F16: { + if (backend_ctx->fa.f32_f16.count(dk_dv)) { + return true; + } + break; + } + case FA_VARIANT_Q8_0: { + if (backend_ctx->fa.f32_q8_0.count(dk_dv)) { + return true; + } + break; + } + case FA_VARIANT_Q4_0: { + if (backend_ctx->fa.f32_q4_0.count(dk_dv)) { + return true; + } + break; + } + case FA_VARIANT_F32_F16_SPLIT: { + if (backend_ctx->fa.f32_f16_split.count(dk_dv)) { + return true; + } + break; + } + case FA_VARIANT_Q8_0_SPLIT: { + if (backend_ctx->fa.f32_q8_0_split.count(dk_dv)) { + return true; + } + break; + } + case FA_VARIANT_Q4_0_SPLIT: { + if (backend_ctx->fa.f32_q4_0_split.count(dk_dv)) { + return true; + } + break; + } + case FA_VARIANT_PRE: { + ggml_opencl_ensure_fa_pre_kernels(backend_ctx, dk, dv); + return true; + } + } + + // not registered but attempted - meaning these kernels failed to compile + const auto attempt_key = std::make_pair(variant, dk_dv); + if (backend_ctx->fa.variant_attempted.count(attempt_key)) { + return false; + } + backend_ctx->fa.variant_attempted.insert(attempt_key); + + const bool is_split = variant == FA_VARIANT_F32_F16_SPLIT || + variant == FA_VARIANT_Q8_0_SPLIT || + variant == FA_VARIANT_Q4_0_SPLIT; + const bool is_quant = variant == FA_VARIANT_Q8_0 || variant == FA_VARIANT_Q8_0_SPLIT || + variant == FA_VARIANT_Q4_0 || variant == FA_VARIANT_Q4_0_SPLIT; + if (is_quant && (dk % 32 != 0 || dv % 32 != 0)) { + return false; + } + if (is_split && cfg->n_split <= 1) { + return false; + } + if ((variant == FA_VARIANT_Q8_0_SPLIT || variant == FA_VARIANT_Q4_0_SPLIT) && + ((dk / 32) % cfg->n_split != 0 || (dv / 4) % cfg->n_split != 0)) { + return false; + } + + const std::string src = ggml_opencl_fa_kernel_src(variant); + + if (src.empty()) { + return false; + } + const std::string opts = ggml_opencl_fa_compile_opts(backend_ctx, cfg, variant); + + const char * tag = nullptr; + switch (variant) { + case FA_VARIANT_F16: tag = "fa f16"; break; + case FA_VARIANT_F32: tag = "fa f32"; break; + case FA_VARIANT_F32_F16: tag = "fa f32_f16"; break; + case FA_VARIANT_Q8_0: tag = "fa q8_0"; break; + case FA_VARIANT_Q4_0: tag = "fa q4_0"; break; + case FA_VARIANT_F32_F16_SPLIT: tag = "fa f32_f16 split"; break; + case FA_VARIANT_Q8_0_SPLIT: tag = "fa q8_0 split"; break; + case FA_VARIANT_Q4_0_SPLIT: tag = "fa q4_0 split"; break; + default: break; + } + cl_program prog = build_program_from_source_ex( + backend_ctx->context, backend_ctx->device, src.c_str(), opts, /*fatal=*/false, tag); + + if (!prog) { + return false; + } + + cl_int err; + switch (variant) { + case FA_VARIANT_F16: { + cl_kernel k, kq1; + CL_CHECK((k = clCreateKernel(prog, "flash_attn_f16", &err), err)); + CL_CHECK((kq1 = clCreateKernel(prog, "flash_attn_f16_q1", &err), err)); + backend_ctx->fa.f16[{dk, dv}] = k; + backend_ctx->fa.f16_q1[{dk, dv}] = kq1; + break; + } + case FA_VARIANT_F32: { + cl_kernel k, kq1; + CL_CHECK((k = clCreateKernel(prog, "flash_attn_f32", &err), err)); + CL_CHECK((kq1 = clCreateKernel(prog, "flash_attn_f32_q1", &err), err)); + backend_ctx->fa.f32[{dk, dv}] = k; + backend_ctx->fa.f32_q1[{dk, dv}] = kq1; + break; + } + case FA_VARIANT_F32_F16: { + cl_kernel k, kq1; + CL_CHECK((k = clCreateKernel(prog, "flash_attn_f32_f16", &err), err)); + CL_CHECK((kq1 = clCreateKernel(prog, "flash_attn_f32_f16_q1", &err), err)); + backend_ctx->fa.f32_f16[{dk, dv}] = k; + backend_ctx->fa.f32_f16_q1[{dk, dv}] = kq1; + ggml_opencl_log_fa_kernel_spill(backend_ctx, k, "flash_attn_f32_f16", dk, dv); + ggml_opencl_log_fa_kernel_spill(backend_ctx, kq1, "flash_attn_f32_f16_q1", dk, dv); + cl_kernel k_split = clCreateKernel(prog, "flash_attn_f32_f16_q1_split", &err); + if (err == CL_SUCCESS) { + backend_ctx->fa.f32_f16_q1_split[{dk, dv}] = k_split; + ggml_opencl_log_fa_kernel_spill(backend_ctx, k_split, "flash_attn_f32_f16_q1_split", dk, dv); + } + cl_kernel k_merge = clCreateKernel(prog, "flash_attn_f32_merge", &err); + if (err == CL_SUCCESS) { + backend_ctx->fa.f32_merge[{dk, dv}] = k_merge; + } + break; + } + case FA_VARIANT_Q8_0: + case FA_VARIANT_Q4_0: { + const bool is_q8 = variant == FA_VARIANT_Q8_0; + const std::string base = is_q8 ? "flash_attn_f32_q8_0" : "flash_attn_f32_q4_0"; + const std::string name_q1 = base + "_q1"; + const std::string name_q1_split = base + "_q1_split"; + auto & m_q1 = is_q8 ? backend_ctx->fa.f32_q8_0_q1 : backend_ctx->fa.f32_q4_0_q1; + auto & m_prefill = is_q8 ? backend_ctx->fa.f32_q8_0 : backend_ctx->fa.f32_q4_0; + auto & m_q1_split = is_q8 ? backend_ctx->fa.f32_q8_0_q1_split : backend_ctx->fa.f32_q4_0_q1_split; + + cl_kernel k, kq1; + CL_CHECK((kq1 = clCreateKernel(prog, name_q1.c_str(), &err), err)); + CL_CHECK((k = clCreateKernel(prog, base.c_str(), &err), err)); + m_q1[{dk, dv}] = kq1; + m_prefill[{dk, dv}] = k; + ggml_opencl_log_fa_kernel_spill(backend_ctx, kq1, name_q1.c_str(), dk, dv); + ggml_opencl_log_fa_kernel_spill(backend_ctx, k, base.c_str(), dk, dv); + cl_kernel k_split = clCreateKernel(prog, name_q1_split.c_str(), &err); + if (err == CL_SUCCESS) { + m_q1_split[{dk, dv}] = k_split; + ggml_opencl_log_fa_kernel_spill(backend_ctx, k_split, name_q1_split.c_str(), dk, dv); + } + if (!backend_ctx->fa.f32_merge.count({dk, dv})) { + cl_kernel k_merge = clCreateKernel(prog, "flash_attn_f32_merge", &err); + if (err == CL_SUCCESS) { + backend_ctx->fa.f32_merge[{dk, dv}] = k_merge; + } + } + break; + } + case FA_VARIANT_F32_F16_SPLIT: { + cl_kernel k; + CL_CHECK((k = clCreateKernel(prog, "flash_attn_f32_f16", &err), err)); + backend_ctx->fa.f32_f16_split[{dk, dv}] = k; + backend_ctx->fa.f32_f16_split_wg_size[{dk, dv}] = cfg->bm * cfg->n_split; + backend_ctx->fa.f32_f16_split_nkv_threshold[{dk, dv}] = cfg->nkv_split_threshold; + break; + } + case FA_VARIANT_Q8_0_SPLIT: + case FA_VARIANT_Q4_0_SPLIT: { + const bool is_q8 = variant == FA_VARIANT_Q8_0_SPLIT; + cl_kernel k; + CL_CHECK((k = clCreateKernel(prog, is_q8 ? "flash_attn_f32_q8_0" : "flash_attn_f32_q4_0", &err), err)); + auto & split = is_q8 ? backend_ctx->fa.f32_q8_0_split : backend_ctx->fa.f32_q4_0_split; + auto & split_wg = is_q8 ? backend_ctx->fa.f32_q8_0_split_wg_size : backend_ctx->fa.f32_q4_0_split_wg_size; + auto & split_bm = is_q8 ? backend_ctx->fa.f32_q8_0_split_bm : backend_ctx->fa.f32_q4_0_split_bm; + auto & split_thresh = is_q8 ? backend_ctx->fa.f32_q8_0_split_nkv_threshold : backend_ctx->fa.f32_q4_0_split_nkv_threshold; + split[{dk, dv}] = k; + split_wg[{dk, dv}] = cfg->bm * cfg->n_split; + split_bm[{dk, dv}] = cfg->bm; + split_thresh[{dk, dv}] = 0; // quant prefill: always split + break; + } + default: + break; + } + CL_CHECK(clReleaseProgram(prog)); + return true; +} + +// Compile a quant FA split kernel with a hand-picked (BLOCK_M, N_SPLIT) that +// overrides the default fa_dims tuning, for the DK values where the default +// N_SPLIT is degenerate for quant prefill: +// DK=256: default N_SPLIT=16 leaves DK/32=8 blocks -> 0 blocks/split. +// Override N_SPLIT=8 (1 block/split), BLOCK_M=16. +// DK=96 : DK/32 = 3 blocks, not divisible by the default N_SPLIT=2 -> +// override N_SPLIT=3. BLOCK_M must be 16, not 32: the N_SPLIT=3 +// QK-partial reduction uses sub_group_shuffle, so all 3 split +// threads of a query must land in one subgroup — WG_SIZE = +// BLOCK_M*N_SPLIT must be <= the 64-lane Adreno subgroup (16*3=48). +static bool ggml_opencl_ensure_fa_quant_split_override( + ggml_backend_opencl_context * backend_ctx, + int dk, int dv, int quant_bm, int quant_n_split, bool is_q8_0 +) { + const std::pair dk_dv = {dk, dv}; + if (is_q8_0 && backend_ctx->fa.f32_q8_0_split.count(dk_dv)) { + return true; + } + if (!is_q8_0 && backend_ctx->fa.f32_q4_0_split.count(dk_dv)) { + return true; + } + + const ggml_opencl_fa_variant variant = is_q8_0 ? FA_VARIANT_Q8_0_SPLIT : FA_VARIANT_Q4_0_SPLIT; + const auto attempt_key = std::make_pair(variant, dk_dv); + if (backend_ctx->fa.variant_attempted.count(attempt_key)) { + return false; + } + + backend_ctx->fa.variant_attempted.insert(attempt_key); + + std::string shuffle_opts; + if (backend_ctx->has_subgroup_shuffle) { + shuffle_opts = backend_ctx->has_qcom_subgroup_shuffle + ? " -D cl_qcom_subgroup_shuffle=1" + : " -D cl_khr_subgroup_shuffle=1"; + } + const ggml_opencl_fa_dim * cfg = nullptr; + for (const auto & d : g_opencl_fa_dims) { + if (d.dk == dk && d.dv == dv) { + cfg = &d; break; + } + } + if (cfg == nullptr) { + return false; + } + + // BLK_PREPASS_BM is the prepass-kernel BLOCK_M, needed so the quant kernel + // indexes the blk[] classification buffer correctly. + std::string opts = backend_ctx->kernel_compile_opts + shuffle_opts + + " -D DK=" + std::to_string(dk) + + " -D DV=" + std::to_string(dv) + + " -D BLOCK_M=" + std::to_string(quant_bm) + + " -D BLOCK_N=" + std::to_string(cfg->bn) + + " -D N_SPLIT=" + std::to_string(quant_n_split) + + " -D BLK_PREPASS_BM=" + std::to_string(cfg->bm); + + const std::string src = ggml_opencl_fa_kernel_src(variant); + if (src.empty()) { + return false; + } + + const std::string tag = std::string("fa ") + (is_q8_0 ? "q8_0" : "q4_0") + + " split DK=" + std::to_string(dk); + cl_program prog = build_program_from_source_ex( + backend_ctx->context, backend_ctx->device, src.c_str(), opts, /*fatal=*/false, tag.c_str()); + + if (!prog) { + return false; + } + + cl_int err; + cl_kernel k; + if (is_q8_0) { + CL_CHECK((k = clCreateKernel(prog, "flash_attn_f32_q8_0", &err), err)); + backend_ctx->fa.f32_q8_0_split[dk_dv] = k; + backend_ctx->fa.f32_q8_0_split_wg_size[dk_dv] = quant_bm * quant_n_split; + backend_ctx->fa.f32_q8_0_split_bm[dk_dv] = quant_bm; + backend_ctx->fa.f32_q8_0_split_nkv_threshold[dk_dv] = 0; + } else { + CL_CHECK((k = clCreateKernel(prog, "flash_attn_f32_q4_0", &err), err)); + backend_ctx->fa.f32_q4_0_split[dk_dv] = k; + backend_ctx->fa.f32_q4_0_split_wg_size[dk_dv] = quant_bm * quant_n_split; + backend_ctx->fa.f32_q4_0_split_bm[dk_dv] = quant_bm; + backend_ctx->fa.f32_q4_0_split_nkv_threshold[dk_dv] = 0; + } + CL_CHECK(clReleaseProgram(prog)); + return true; +} + namespace /* anonymous */ { extern struct ggml_backend_device_i ggml_backend_opencl_device_i; } @@ -3955,6 +4425,8 @@ static void ggml_opencl_print_backend_info(ggml_backend_opencl_device_context * backend_ctx->driver_version.c_str()); GGML_LOG_INFO("ggml_opencl: vector subgroup broadcast support: %s\n", backend_ctx->has_vector_subgroup_broadcast ? "true" : "false"); + GGML_LOG_INFO("ggml_opencl: subgroup shuffle support: %s\n", + backend_ctx->has_subgroup_shuffle ? "true" : "false"); GGML_LOG_INFO("ggml_opencl: device FP16 support: %s\n", backend_ctx->fp16_support ? "true" : "false"); GGML_LOG_INFO("ggml_opencl: mem base addr align: %u\n", @@ -4111,6 +4583,8 @@ static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev) { backend_ctx->gpu_family = dev_ctx->gpu_family; backend_ctx->adreno_gen = dev_ctx->adreno_gen; if (backend_ctx->gpu_family == GPU_FAMILY::ADRENO) { + ggml_cl_init_fa_dims_table(); + // Use wave size of 64 for all Adreno GPUs. backend_ctx->adreno_wave_size = 64; } @@ -4156,6 +4630,11 @@ static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev) { // check Adreno large buffer support backend_ctx->adreno_has_large_buffer = strstr(ext_buffer, "cl_qcom_large_buffer") != NULL; + // subgroup shuffle support (N_SPLIT>1 FA kernel) + backend_ctx->has_qcom_subgroup_shuffle = strstr(ext_buffer, "cl_qcom_subgroup_shuffle") != NULL; + backend_ctx->has_subgroup_shuffle = + strstr(ext_buffer, "cl_khr_subgroup_shuffle") != NULL || + backend_ctx->has_qcom_subgroup_shuffle; cl_uint base_align_in_bits; CL_CHECK(clGetDeviceInfo(device, CL_DEVICE_MEM_BASE_ADDR_ALIGN, sizeof(cl_uint), &base_align_in_bits, NULL)); @@ -5100,6 +5579,8 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te switch (op->type) { case GGML_TYPE_F16: case GGML_TYPE_F32: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q4_0: return (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32); default: return false; @@ -5175,9 +5656,10 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te case GGML_UNARY_OP_TANH: case GGML_UNARY_OP_NEG: case GGML_UNARY_OP_EXP: - return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16; + // Adreno F16 exp/expm1 overflow even post-half->float convert. + return op->src[0]->type == GGML_TYPE_F32; case GGML_UNARY_OP_EXPM1: - return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16; + return op->src[0]->type == GGML_TYPE_F32; case GGML_UNARY_OP_SOFTPLUS: return op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16; default: @@ -5250,7 +5732,10 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te return true; } else if (op->src[0]->type == GGML_TYPE_F32) { return op->src[1]->type == GGML_TYPE_F32; - } else if (op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q4_1 || + } else if (op->src[0]->type == GGML_TYPE_Q4_0) { + // Non-contig src0 routes through on-device dequant-to-f16. + return op->src[1]->type == GGML_TYPE_F32; + } else if (op->src[0]->type == GGML_TYPE_Q4_1 || op->src[0]->type == GGML_TYPE_Q5_0 || op->src[0]->type == GGML_TYPE_Q5_1 || op->src[0]->type == GGML_TYPE_MXFP4 || op->src[0]->type == GGML_TYPE_IQ4_NL || @@ -5339,43 +5824,55 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te return op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]); case GGML_OP_MEAN: return op->src[0]->type == GGML_TYPE_F32; - case GGML_OP_FLASH_ATTN_EXT: - { - load_cl_kernels_flash_attn(backend_ctx); + case GGML_OP_FLASH_ATTN_EXT: { + const ggml_tensor * q = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * v = op->src[2]; - const ggml_tensor * q = op->src[0]; - const ggml_tensor * k = op->src[1]; - const ggml_tensor * v = op->src[2]; + const int dk = q->ne[0]; + const int dv = v->ne[0]; - const int dk = q->ne[0]; - const int dv = v->ne[0]; + const struct { int dk; int dv; } supported_dims[] = { + { 40, 40}, { 64, 64}, { 80, 80}, { 96, 96}, + {112, 112}, {128, 128}, {192, 128}, + {192, 192}, {256, 256}, + }; - const struct { int dk; int dv; } supported_dims[] = { - { 40, 40}, { 64, 64}, { 80, 80}, { 96, 96}, - {112, 112}, {128, 128}, {192, 128}, - {192, 192}, {256, 256}, - }; - - bool dims_supported = false; - for (size_t i = 0; i < sizeof(supported_dims)/sizeof(supported_dims[0]); ++i) { - if (supported_dims[i].dk == dk && supported_dims[i].dv == dv) { - dims_supported = true; - break; - } + bool dims_supported = false; + for (size_t i = 0; i < sizeof(supported_dims)/sizeof(supported_dims[0]); ++i) { + if (supported_dims[i].dk == dk && supported_dims[i].dv == dv) { + dims_supported = true; + break; } - if (!dims_supported) { - return false; - } - - const bool is_f32_f32 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F32 && - v->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; - const bool is_f16_f16 = q->type == GGML_TYPE_F16 && k->type == GGML_TYPE_F16 && - v->type == GGML_TYPE_F16 && op->type == GGML_TYPE_F16; - const bool is_f32_f16 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && - v->type == GGML_TYPE_F16 && op->type == GGML_TYPE_F32; - - return is_f32_f32 || is_f16_f16 || is_f32_f16; } + if (!dims_supported) { + return false; + } + + const bool is_f32_f32 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F32 && + v->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; + const bool is_f16_f16 = q->type == GGML_TYPE_F16 && k->type == GGML_TYPE_F16 && + v->type == GGML_TYPE_F16 && op->type == GGML_TYPE_F16; + const bool is_f32_f16 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && + v->type == GGML_TYPE_F16 && op->type == GGML_TYPE_F32; + const bool is_f32_q8_0 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_Q8_0 && + v->type == GGML_TYPE_Q8_0 && op->type == GGML_TYPE_F32 && + dk % 32 == 0 && dv % 32 == 0; + const bool is_f32_q4_0 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_Q4_0 && + v->type == GGML_TYPE_Q4_0 && op->type == GGML_TYPE_F32 && + dk % 32 == 0 && dv % 32 == 0; + + // Asymmetric KV: host-dequants both sides to F32, uses f32 kernel. + auto is_kv_type_ok = [](ggml_type t) { + return t == GGML_TYPE_F16 || t == GGML_TYPE_F32 || + t == GGML_TYPE_Q4_0 || t == GGML_TYPE_Q8_0; + }; + const bool is_f32_asym = q->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && + k->type != v->type && + is_kv_type_ok(k->type) && is_kv_type_ok(v->type); + + return is_f32_f32 || is_f16_f16 || is_f32_f16 || is_f32_q8_0 || is_f32_q4_0 || is_f32_asym; + } default: return false; } @@ -5737,6 +6234,9 @@ struct ggml_backend_opencl_buffer_context { temp_tensor_extras_q6_K.push_back(e); } temp_tensor_extras_q6_K_in_use.clear(); + + q8_0_soa_tensors.clear(); + q4_0_soa_tensors.clear(); } // Pools for extras. Available extras are in `temp_tensor_extras`. Extras @@ -5767,6 +6267,17 @@ struct ggml_backend_opencl_buffer_context { std::vector temp_tensor_extras_q6_K; std::vector temp_tensor_extras_q6_K_in_use; + // q8_0 tensors with AoS->SoA layout conversion installed by set_tensor. + // Two types of tensors get SOA'ed - normal weights and MoE weights. + // In Q8_0's case, we only have normal weights. If we ever have Q8_0 as MoE + // weights, they need to be added to this set in `set_tensors`. + std::unordered_set q8_0_soa_tensors; + + // Same for q4_0. KV-cache q4_0 tensors are allocated but never pass + // through set_tensor, so they stay AoS and aren't in this set. + // In Q4_0's case, in addition to normal weights, we have MoE weights. + std::unordered_set q4_0_soa_tensors; + // The buffer_context is initially created by ggml_backend_buft_alloc_buffer // before any tensor is initialized (at the beginning of alloc_tensor_range). // Hence, there is always a buffer object in this vector. When each tensor is @@ -5848,6 +6359,10 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer, // buffers for quantized bits and scales, which are then populated by the // conversion kernel. if (tensor->type == GGML_TYPE_Q4_0) { + // Views can't SoA-ify here — parent owns the layout (see q8_0 guard). + if (tensor->view_src != nullptr || !ggml_is_contiguous(tensor)) { + return; + } // Tensors should have been preallocated, therefore they should // already have ggml_tensor_extra_cl as extra. ggml_tensor_extra_cl * extra_orig = (ggml_tensor_extra_cl *)tensor->extra; @@ -5937,6 +6452,8 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer, }; extra->q_img = clCreateImage(context, CL_MEM_READ_ONLY, &img_format_q, &img_desc_q, NULL, &err); tensor->extra = extra; + // MoE tensors are also SOA'ed + ctx->q4_0_soa_tensors.insert(tensor); return; } @@ -5965,6 +6482,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer, CL_CHECK(clReleaseMemObject(data_device)); tensor->extra = extra; + ctx->q4_0_soa_tensors.insert(tensor); // transpose the weights and scales #ifdef GGML_OPENCL_USE_ADRENO_KERNELS @@ -6516,6 +7034,11 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer, return; } if (tensor->type == GGML_TYPE_Q8_0) { + // Views share the parent's buffer; parent owns SoA conversion. + if (tensor->view_src != nullptr || !ggml_is_contiguous(tensor)) { + return; + } + ggml_tensor_extra_cl * extra_orig = (ggml_tensor_extra_cl *)tensor->extra; GGML_ASSERT(extra_orig && "Tesnors in OpenCL backend should have been allocated and initialized"); @@ -6571,6 +7094,7 @@ static void ggml_backend_opencl_buffer_set_tensor(ggml_backend_buffer_t buffer, CL_CHECK(clReleaseMemObject(data_device)); tensor->extra = extra; + ctx->q8_0_soa_tensors.insert(tensor); // Transpose the weights and scales #ifdef GGML_OPENCL_USE_ADRENO_KERNELS @@ -7226,7 +7750,18 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer, // To properly support this, we need to restore block_q4_0 struct arrays // from the flattened buffers. if (tensor->type == GGML_TYPE_Q4_0) { - ggml_tensor_extra_cl_q4_0 * extra = (ggml_tensor_extra_cl_q4_0 *)tensor->extra; + // KV-cache q4_0 stays AoS — direct readback, no SoA restore. + if (!ggml_cl_is_q4_0_soa(tensor)) { + ggml_tensor_extra_cl * extra_aos = (ggml_tensor_extra_cl *) tensor->extra; + CL_CHECK(clEnqueueReadBuffer( + queue, extra_aos->data_device, CL_TRUE, + extra_aos->offset + tensor->view_offs + offset, + size, data, 0, NULL, NULL)); + return; + } + // SoA extra lives on the parent tensor — follow view_src. + const ggml_tensor * extra_src = tensor->view_src != nullptr ? tensor->view_src : tensor; + ggml_tensor_extra_cl_q4_0 * extra = (ggml_tensor_extra_cl_q4_0 *)extra_src->extra; #ifdef GGML_OPENCL_USE_ADRENO_KERNELS if (use_adreno_moe_kernels(backend_ctx, tensor)) { @@ -7697,7 +8232,18 @@ static void ggml_backend_opencl_buffer_get_tensor(ggml_backend_buffer_t buffer, return; } if (tensor->type == GGML_TYPE_Q8_0) { - ggml_tensor_extra_cl_q8_0 * extra = (ggml_tensor_extra_cl_q8_0 *)tensor->extra; + // KV-cache q8_0 stays AoS (see Q4_0 branch). + if (!ggml_cl_is_q8_0_soa(tensor)) { + ggml_tensor_extra_cl * extra_aos = (ggml_tensor_extra_cl *) tensor->extra; + CL_CHECK(clEnqueueReadBuffer( + queue, extra_aos->data_device, CL_TRUE, + extra_aos->offset + tensor->view_offs + offset, + size, data, 0, NULL, NULL)); + return; + } + // SoA extra lives on the parent — follow view_src. + const ggml_tensor * extra_src = tensor->view_src != nullptr ? tensor->view_src : tensor; + ggml_tensor_extra_cl_q8_0 * extra = (ggml_tensor_extra_cl_q8_0 *)extra_src->extra; cl_int err; cl_mem data_device = clCreateBuffer(context, CL_MEM_READ_WRITE, @@ -8821,6 +9367,34 @@ static void ggml_cl_get_rows(ggml_backend_t backend, const ggml_tensor * src0, c backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst); } +// check if a Q8_0 tensor has been SOA'ed in set_tensor +// we store SOA'ed tensors in a map in set_tensor, check against that map +static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor) { + if (tensor == nullptr || tensor->type != GGML_TYPE_Q8_0 || tensor->buffer == nullptr) { + return false; + } + auto * ctx = (ggml_backend_opencl_buffer_context *) tensor->buffer->context; + if (ctx == nullptr) { + return false; + } + const ggml_tensor * key = tensor->view_src != nullptr ? tensor->view_src : tensor; + return ctx->q8_0_soa_tensors.count(key) > 0; +} + +// check if a Q4_0 tensor has been SOA'ed in set_tensor +// we store SOA'ed tensors in a map in set_tensor, check against that map +static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor) { + if (tensor == nullptr || tensor->type != GGML_TYPE_Q4_0 || tensor->buffer == nullptr) { + return false; + } + auto * ctx = (ggml_backend_opencl_buffer_context *) tensor->buffer->context; + if (ctx == nullptr) { + return false; + } + const ggml_tensor * key = tensor->view_src != nullptr ? tensor->view_src : tensor; + return ctx->q4_0_soa_tensors.count(key) > 0; +} + static void ggml_cl_set_rows(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { GGML_ASSERT(src0); GGML_ASSERT(src0->extra); @@ -8834,26 +9408,14 @@ static void ggml_cl_set_rows(ggml_backend_t backend, const ggml_tensor * src0, c // ne2 = ne02 // ne3 = ne03 - const int ne01 = src0->ne[1]; - const int ne02 = src0->ne[2]; - const int ne03 = src0->ne[3]; + GGML_TENSOR_LOCALS(int, ne0, src0, ne); + GGML_TENSOR_LOCALS(cl_ulong, nb0, src0, nb); - const cl_ulong nb01 = src0->nb[1]; - const cl_ulong nb02 = src0->nb[2]; - const cl_ulong nb03 = src0->nb[3]; + GGML_TENSOR_LOCALS(int, ne1, src1, ne); + GGML_TENSOR_LOCALS(cl_ulong, nb1, src1, nb); - const int ne11 = src1->ne[1]; - const int ne12 = src1->ne[2]; - - const cl_ulong nb10 = src1->nb[0]; - const cl_ulong nb11 = src1->nb[1]; - const cl_ulong nb12 = src1->nb[2]; - - const int ne0 = dst->ne[0]; - - const cl_ulong nb1 = dst->nb[1]; - const cl_ulong nb2 = dst->nb[2]; - const cl_ulong nb3 = dst->nb[3]; + GGML_TENSOR_LOCALS(int, ne, dst, ne); + GGML_TENSOR_LOCALS(cl_ulong, nb, dst, nb); const int nblk0 = ne0/ggml_blck_size(dst->type); @@ -8861,31 +9423,49 @@ static void ggml_cl_set_rows(ggml_backend_t backend, const ggml_tensor * src0, c ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *)src0->extra; ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *)src1->extra; - ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra; cl_ulong offset0 = extra0->offset + src0->view_offs; cl_ulong offset1 = extra1->offset + src1->view_offs; - cl_ulong offsetd = extrad->offset + dst->view_offs; + + const bool q8_0_soa = dst->type == GGML_TYPE_Q8_0 && ggml_cl_is_q8_0_soa(dst); + const bool q4_0_soa = dst->type == GGML_TYPE_Q4_0 && ggml_cl_is_q4_0_soa(dst); + const bool is_soa = q8_0_soa || q4_0_soa; cl_kernel kernel; - switch (dst->type) { - case GGML_TYPE_F32: - if (src1->type == GGML_TYPE_I64) { - kernel = backend_ctx->kernel_set_rows_f32_i64; - } else { - kernel = backend_ctx->kernel_set_rows_f32_i32; - } - break; - case GGML_TYPE_F16: - if (src1->type == GGML_TYPE_I64) { - kernel = backend_ctx->kernel_set_rows_f16_i64; - } else { - kernel = backend_ctx->kernel_set_rows_f16_i32; - } - break; - default: - GGML_ABORT("not implemented"); + if (q8_0_soa) { + kernel = (src1->type == GGML_TYPE_I64) + ? backend_ctx->kernel_set_rows_q8_0_soa_i64 + : backend_ctx->kernel_set_rows_q8_0_soa_i32; + } else if (q4_0_soa) { + kernel = (src1->type == GGML_TYPE_I64) + ? backend_ctx->kernel_set_rows_q4_0_soa_i64 + : backend_ctx->kernel_set_rows_q4_0_soa_i32; + } else { + switch (dst->type) { + case GGML_TYPE_F32: + kernel = (src1->type == GGML_TYPE_I64) + ? backend_ctx->kernel_set_rows_f32_i64 + : backend_ctx->kernel_set_rows_f32_i32; + break; + case GGML_TYPE_F16: + kernel = (src1->type == GGML_TYPE_I64) + ? backend_ctx->kernel_set_rows_f16_i64 + : backend_ctx->kernel_set_rows_f16_i32; + break; + case GGML_TYPE_Q8_0: + kernel = (src1->type == GGML_TYPE_I64) + ? backend_ctx->kernel_set_rows_q8_0_i64 + : backend_ctx->kernel_set_rows_q8_0_i32; + break; + case GGML_TYPE_Q4_0: + kernel = (src1->type == GGML_TYPE_I64) + ? backend_ctx->kernel_set_rows_q4_0_i64 + : backend_ctx->kernel_set_rows_q4_0_i32; + break; + default: + GGML_ABORT("not implemented"); + } } fastdiv_vals ne11_ = init_fastdiv_values(ne11); @@ -8895,21 +9475,65 @@ static void ggml_cl_set_rows(ggml_backend_t backend, const ggml_tensor * src0, c CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0)); CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extra1->data_device)); CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset1)); - CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &extrad->data_device)); - CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offsetd)); - CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &ne01)); - CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &nb01)); - CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_ulong), &nb02)); - CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &nb03)); - CL_CHECK(clSetKernelArg(kernel, 10, sizeof(fastdiv_vals), &ne11_)); - CL_CHECK(clSetKernelArg(kernel, 11, sizeof(fastdiv_vals), &ne12_)); - CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_ulong), &nb10)); - CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &nb11)); - CL_CHECK(clSetKernelArg(kernel, 14, sizeof(cl_ulong), &nb12)); - CL_CHECK(clSetKernelArg(kernel, 15, sizeof(int), &nblk0)); - CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &nb1)); - CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &nb2)); - CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &nb3)); + + if (is_soa) { + // The q/d subbuffers in q8_0/q4_0 extras are interchangeable here. + // For views (e.g. ggml_set_rows' `out`), follow view_src for the SoA extra. + const ggml_tensor * soa_src = dst->view_src != nullptr ? dst->view_src : dst; + cl_mem q_mem = nullptr; + cl_mem d_mem = nullptr; + if (q8_0_soa) { + ggml_tensor_extra_cl_q8_0 * e = (ggml_tensor_extra_cl_q8_0 *)soa_src->extra; + q_mem = e->q; + d_mem = e->d; + } else { + ggml_tensor_extra_cl_q4_0 * e = (ggml_tensor_extra_cl_q4_0 *)soa_src->extra; + q_mem = e->q; + d_mem = e->d; + } + cl_ulong offset_q = 0; + cl_ulong offset_d = 0; + const int ne1_dst = dst->ne[1]; + const int ne2_dst = dst->ne[2]; + const int ne3_dst = dst->ne[3]; + + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &q_mem)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offset_q)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_mem), &d_mem)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &offset_d)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(int), &ne01)); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &nb01)); + CL_CHECK(clSetKernelArg(kernel, 10, sizeof(cl_ulong), &nb02)); + CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &nb03)); + CL_CHECK(clSetKernelArg(kernel, 12, sizeof(fastdiv_vals), &ne11_)); + CL_CHECK(clSetKernelArg(kernel, 13, sizeof(fastdiv_vals), &ne12_)); + CL_CHECK(clSetKernelArg(kernel, 14, sizeof(cl_ulong), &nb10)); + CL_CHECK(clSetKernelArg(kernel, 15, sizeof(cl_ulong), &nb11)); + CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &nb12)); + CL_CHECK(clSetKernelArg(kernel, 17, sizeof(int), &nblk0)); + CL_CHECK(clSetKernelArg(kernel, 18, sizeof(int), &ne1_dst)); + CL_CHECK(clSetKernelArg(kernel, 19, sizeof(int), &ne2_dst)); + CL_CHECK(clSetKernelArg(kernel, 20, sizeof(int), &ne3_dst)); + } else { + ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra; + cl_ulong offsetd = extrad->offset + dst->view_offs; + + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &extrad->data_device)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offsetd)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(int), &ne01)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &nb01)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_ulong), &nb02)); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &nb03)); + CL_CHECK(clSetKernelArg(kernel, 10, sizeof(fastdiv_vals), &ne11_)); + CL_CHECK(clSetKernelArg(kernel, 11, sizeof(fastdiv_vals), &ne12_)); + CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_ulong), &nb10)); + CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &nb11)); + CL_CHECK(clSetKernelArg(kernel, 14, sizeof(cl_ulong), &nb12)); + CL_CHECK(clSetKernelArg(kernel, 15, sizeof(int), &nblk0)); + CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &nb1)); + CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &nb2)); + CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &nb3)); + } int nth0 = 64; if (backend_ctx->gpu_family == INTEL) { @@ -11483,14 +12107,370 @@ static void ggml_cl_timestep_embedding(ggml_backend_t backend, const ggml_tensor backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, NULL, dst); } +struct ggml_cl_flash_attn_temp_buffer { + cl_mem data = nullptr; + + ~ggml_cl_flash_attn_temp_buffer() { + if (data != nullptr) { + CL_CHECK(clReleaseMemObject(data)); + data = nullptr; + } + } +}; + +// Resolve the source buffer + strides for an FA KV tensor: keep the +// caller-supplied AoS buffer if non-NULL, else fall back to tensor->extra. +static void ggml_cl_flash_attn_resolve_src( + const ggml_tensor * tensor, + cl_mem & buf, + cl_ulong & offset, + cl_ulong & nb1, + cl_ulong & nb2, + cl_ulong & nb3) { + if (buf != NULL) { + return; + } + ggml_tensor_extra_cl * extra = (ggml_tensor_extra_cl *) tensor->extra; + GGML_ASSERT(extra && extra->data_device); + buf = extra->data_device; + offset = extra->offset + tensor->view_offs; + nb1 = tensor->nb[1]; + nb2 = tensor->nb[2]; + nb3 = tensor->nb[3]; +} + +// Read a (possibly strided-view) tensor from device into a tight host buffer. +// dim 0 is always tight; a strided view is gathered row-by-row. +static void ggml_cl_flash_attn_read_tensor_host( + ggml_backend_opencl_context * backend_ctx, + const ggml_tensor * tensor, + cl_mem src_buffer, cl_ulong src_offset, + cl_ulong src_nb1, cl_ulong src_nb2, cl_ulong src_nb3, + size_t row_bytes, void * dst, size_t total_bytes +) { + const bool contiguous_layout = + src_nb1 == row_bytes && + src_nb2 == row_bytes * (cl_ulong) tensor->ne[1] && + src_nb3 == src_nb2 * (cl_ulong) tensor->ne[2]; + + if (contiguous_layout) { + CL_CHECK(clEnqueueReadBuffer(backend_ctx->queue, src_buffer, CL_TRUE, + src_offset, total_bytes, dst, 0, NULL, NULL)); + return; + } + + size_t dst_off = 0; + for (int64_t i3 = 0; i3 < tensor->ne[3]; ++i3) { + for (int64_t i2 = 0; i2 < tensor->ne[2]; ++i2) { + for (int64_t i1 = 0; i1 < tensor->ne[1]; ++i1) { + const cl_ulong row_src_off = src_offset + + (cl_ulong) i3 * src_nb3 + + (cl_ulong) i2 * src_nb2 + + (cl_ulong) i1 * src_nb1; + CL_CHECK(clEnqueueReadBuffer(backend_ctx->queue, src_buffer, CL_TRUE, + row_src_off, row_bytes, + (uint8_t *) dst + dst_off, 0, NULL, NULL)); + dst_off += row_bytes; + } + } + } + GGML_ASSERT(dst_off == total_bytes); +} + +// Rebuild AoS q8_0/q4_0 bytes from a SoA tensor into a temp buffer. +// Returns false if the tensor is not SoA-quantised (already AoS). +static bool ggml_cl_flash_attn_reconstruct_aos( + ggml_backend_opencl_context * backend_ctx, + const ggml_tensor * tensor, + ggml_cl_flash_attn_temp_buffer & temp, + cl_mem & out_buf, + cl_ulong & out_offset, + cl_ulong & out_nb1, + cl_ulong & out_nb2, + cl_ulong & out_nb3 +) { + if (tensor == nullptr) { + return false; + } + const bool is_q8_0 = tensor->type == GGML_TYPE_Q8_0 && ggml_cl_is_q8_0_soa(tensor); + const bool is_q4_0 = tensor->type == GGML_TYPE_Q4_0 && ggml_cl_is_q4_0_soa(tensor); + if (!is_q8_0 && !is_q4_0) { + return false; + } + + // For views, SoA extra is on view_src (view->extra is pre-SoA). + // Noshuffle layout only applies to 2D weights, as determined by `use_adreno_kernels`, + // where ne2 == 1 and ne3 == 1 -- these are never FA inputs. + // Therefore, we use `restore_block_qk_0` kernels, not `restore_block_qk_0_noshuffle`. + const ggml_tensor * soa_src = tensor->view_src ? tensor->view_src : tensor; + cl_mem extra_q = NULL; + cl_mem extra_d = NULL; + if (is_q8_0) { + auto * e = (ggml_tensor_extra_cl_q8_0 *) soa_src->extra; + GGML_ASSERT(e && e->q && e->d); + extra_q = e->q; + extra_d = e->d; + } else { + auto * e = (ggml_tensor_extra_cl_q4_0 *) soa_src->extra; + GGML_ASSERT(e && e->q && e->d); + extra_q = e->q; + extra_d = e->d; + } + + // Reconstruct the whole parent; view offsets then work naturally. + const size_t parent_nbytes = ggml_nbytes(soa_src); + cl_int err; + temp.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, parent_nbytes, NULL, &err); + CL_CHECK(err); + + cl_kernel kernel = is_q8_0 ? backend_ctx->kernel_restore_block_q8_0 + : backend_ctx->kernel_restore_block_q4_0; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra_q)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra_d)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &temp.data)); + + const size_t n_blocks = (size_t) ggml_nelements(soa_src) / ggml_blck_size(soa_src->type); + size_t global_work_size[] = { n_blocks, 1, 1 }; + size_t local_work_size[] = { 1, 1, 1 }; + CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL, + global_work_size, local_work_size, 0, NULL, NULL)); + + out_buf = temp.data; + out_offset = tensor->view_offs; + out_nb1 = tensor->nb[1]; + out_nb2 = tensor->nb[2]; + out_nb3 = tensor->nb[3]; + return true; +} + +// GPU dequant of a contiguous q4_0/q8_0 KV tensor to f16/f32. Caller supplies +// src_buf when reconstructing from SoA. Returns false for non-contig layouts +// (the kernel indexes blocks tightly within ne[0]) so the caller can fall back +// to the host path. +static bool ggml_cl_flash_attn_dequant_kv_gpu( + ggml_backend_opencl_context * backend_ctx, + const ggml_tensor * tensor, + ggml_type target_type, + cl_mem in_src_buf, + cl_ulong in_src_offset, + cl_ulong in_src_nb1, + cl_ulong in_src_nb2, + cl_ulong in_src_nb3, + ggml_cl_flash_attn_temp_buffer & temp, + cl_mem & out_buf, + cl_ulong & out_offset, + cl_ulong & out_nb1, + cl_ulong & out_nb2, + cl_ulong & out_nb3 +) { + GGML_ASSERT(tensor->type == GGML_TYPE_Q8_0 || tensor->type == GGML_TYPE_Q4_0); + GGML_ASSERT(target_type == GGML_TYPE_F16 || target_type == GGML_TYPE_F32); + + const bool is_q8_0 = tensor->type == GGML_TYPE_Q8_0; + + cl_mem src_buf = in_src_buf; + cl_ulong src_offset = in_src_offset; + cl_ulong src_nb1 = in_src_nb1; + cl_ulong src_nb2 = in_src_nb2; + cl_ulong src_nb3 = in_src_nb3; + ggml_cl_flash_attn_resolve_src(tensor, src_buf, src_offset, src_nb1, src_nb2, src_nb3); + + if (tensor->nb[0] != (cl_ulong) ggml_type_size(tensor->type)) { + return false; + } + + const size_t n_blocks = (size_t) ggml_nelements(tensor) / 32; // block size is 32 + const size_t elem_size = ggml_type_size(target_type); + const size_t out_bytes = n_blocks * 32 * elem_size; + const cl_int nblk0_arg = (cl_int) (tensor->ne[0] / 32); + const cl_int ne1_arg = (cl_int) tensor->ne[1]; + const cl_int ne2_arg = (cl_int) tensor->ne[2]; + const cl_int ne3_arg = (cl_int) tensor->ne[3]; + + cl_int err; + temp.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, out_bytes, NULL, &err); + CL_CHECK(err); + + cl_kernel kernel; + if (target_type == GGML_TYPE_F16) { + kernel = is_q8_0 ? backend_ctx->kernel_dequant_q8_0_f16_view_aos + : backend_ctx->kernel_dequant_q4_0_f16_view_aos; + } else { + kernel = is_q8_0 ? backend_ctx->kernel_dequant_q8_0_f32_view_aos + : backend_ctx->kernel_dequant_q4_0_f32_view_aos; + } + + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &src_buf)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &src_offset)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_ulong), &src_nb1)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &src_nb2)); + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_ulong), &src_nb3)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_int), &nblk0_arg)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_int), &ne1_arg)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_int), &ne2_arg)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_int), &ne3_arg)); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_mem), &temp.data)); + + size_t global_ws[3] = { (size_t) nblk0_arg, (size_t) ne1_arg, (size_t) ne2_arg * (size_t) ne3_arg }; + CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL, + global_ws, NULL, 0, NULL, NULL)); + + out_buf = temp.data; + out_offset = 0; + out_nb1 = (cl_ulong) tensor->ne[0] * elem_size; + out_nb2 = out_nb1 * (cl_ulong) tensor->ne[1]; + out_nb3 = out_nb2 * (cl_ulong) tensor->ne[2]; + return true; +} + +static bool ggml_cl_flash_attn_prepare_quantized_tensor( + ggml_backend_opencl_context * backend_ctx, + const ggml_tensor * tensor, + ggml_type target_type, + ggml_cl_flash_attn_temp_buffer & temp, + cl_mem & data_device, + cl_ulong & offset, + cl_ulong & nb1, + cl_ulong & nb2, + cl_ulong & nb3 +) { + if (!ggml_is_quantized(tensor->type)) { + return false; + } + + // Caller-supplied AoS buffer wins over tensor->extra when present. + cl_mem src_buffer = data_device; + cl_ulong src_offset = offset; + cl_ulong src_nb1 = nb1; + cl_ulong src_nb2 = nb2; + cl_ulong src_nb3 = nb3; + ggml_cl_flash_attn_resolve_src(tensor, src_buffer, src_offset, src_nb1, src_nb2, src_nb3); + + const int64_t n = ggml_nelements(tensor); + const size_t row_bytes = (size_t) (tensor->ne[0] / ggml_blck_size(tensor->type)) * ggml_type_size(tensor->type); + // tight-packed byte count (ggml_nbytes includes stride gaps). + const size_t total_bytes = (size_t) (n / ggml_blck_size(tensor->type)) * ggml_type_size(tensor->type); + std::vector host_quant(total_bytes); + + sync_with_other_backends(backend_ctx); + ggml_cl_flash_attn_read_tensor_host(backend_ctx, tensor, src_buffer, src_offset, + src_nb1, src_nb2, src_nb3, + row_bytes, host_quant.data(), total_bytes); + + std::vector host_f32(n); + ggml_get_type_traits(tensor->type)->to_float(host_quant.data(), host_f32.data(), n); + + const size_t bytes_per_elem = ggml_type_size(target_type); + const size_t buffer_size = (size_t) n * bytes_per_elem; + + std::vector host_linear(buffer_size); + if (target_type == GGML_TYPE_F32) { + memcpy(host_linear.data(), host_f32.data(), buffer_size); + } else { + GGML_ASSERT(target_type == GGML_TYPE_F16); + ggml_fp32_to_fp16_row(host_f32.data(), (ggml_fp16_t *) host_linear.data(), n); + } + + cl_int err; + temp.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, buffer_size, NULL, &err); + CL_CHECK(err); + CL_CHECK(clEnqueueWriteBuffer(backend_ctx->queue, temp.data, CL_TRUE, 0, buffer_size, host_linear.data(), 0, NULL, NULL)); + + data_device = temp.data; + offset = 0; + nb1 = (cl_ulong) (tensor->ne[0] * bytes_per_elem); + nb2 = (cl_ulong) (tensor->ne[1] * nb1); + nb3 = (cl_ulong) (tensor->ne[2] * nb2); + + static bool warned = false; + if (!warned) { + GGML_LOG_WARN("ggml_opencl: OpenCL flash attention dequantizes GPU-resident quantized KV cache into temporary linear buffers; performance may be poor\n"); + warned = true; + } + + return true; +} + +// Host-side F16 -> F32 for the asymmetric-KV F32 fallback path. +static bool ggml_cl_flash_attn_convert_f16_to_f32( + ggml_backend_opencl_context * backend_ctx, + const ggml_tensor * tensor, + ggml_cl_flash_attn_temp_buffer & temp, + cl_mem & data_device, + cl_ulong & offset, + cl_ulong & nb1, + cl_ulong & nb2, + cl_ulong & nb3 +) { + if (tensor->type != GGML_TYPE_F16) { + return false; + } + + cl_mem src_buffer = data_device; + cl_ulong src_offset = offset; + cl_ulong src_nb1 = nb1; + cl_ulong src_nb2 = nb2; + cl_ulong src_nb3 = nb3; + ggml_cl_flash_attn_resolve_src(tensor, src_buffer, src_offset, src_nb1, src_nb2, src_nb3); + + const int64_t n = ggml_nelements(tensor); + const size_t row_bytes = (size_t) tensor->ne[0] * sizeof(ggml_fp16_t); + const size_t total_bytes = (size_t) n * sizeof(ggml_fp16_t); + std::vector host_f16(total_bytes); + + sync_with_other_backends(backend_ctx); + ggml_cl_flash_attn_read_tensor_host(backend_ctx, tensor, src_buffer, src_offset, + src_nb1, src_nb2, src_nb3, + row_bytes, host_f16.data(), total_bytes); + + std::vector host_f32(n); + ggml_fp16_to_fp32_row((const ggml_fp16_t *) host_f16.data(), host_f32.data(), n); + + const size_t f32_bytes = (size_t) n * sizeof(float); + cl_int err; + temp.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, f32_bytes, NULL, &err); + CL_CHECK(err); + CL_CHECK(clEnqueueWriteBuffer(backend_ctx->queue, temp.data, CL_TRUE, 0, + f32_bytes, host_f32.data(), 0, NULL, NULL)); + + data_device = temp.data; + offset = 0; + nb1 = (cl_ulong) (tensor->ne[0] * sizeof(float)); + nb2 = (cl_ulong) (tensor->ne[1] * nb1); + nb3 = (cl_ulong) (tensor->ne[2] * nb2); + + static bool warned = false; + if (!warned) { + GGML_LOG_WARN("ggml_opencl: OpenCL flash attention asymmetric KV converts an F16 cache to F32 host-side; performance may be poor\n"); + warned = true; + } + + return true; +} + +// Flash-Decoding (K-split) dispatch thresholds. FD fires for non-causal +// attention with n_kv >= FD_MIN_N_KV and d_head <= FD_MAX_DK; the KV range is +// split into ~n_kv/FD_KV_PER_SPLIT partials, clamped to [FD_MIN_SPLITS, +// FD_MAX_SPLITS]. Multi-query FD is restricted to small heads +// (d_head <= FD_MAX_DK_MULTI) and capped at FD_MAX_N_Q_MULTI queries. +static constexpr int FD_MIN_N_KV = 2048; +static constexpr int FD_KV_PER_SPLIT = 2048; +static constexpr int FD_MIN_SPLITS = 2; +static constexpr int FD_MAX_SPLITS = 16; +static constexpr int FD_MAX_DK = 128; +static constexpr int FD_MAX_DK_MULTI = 64; +static constexpr int FD_MAX_N_Q_MULTI = 8; + static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, const ggml_tensor * k, ggml_tensor * dst) { const ggml_tensor * v = dst->src[2]; const ggml_tensor * mask = dst->src[3]; const ggml_tensor * sinks = dst->src[4]; + GGML_ASSERT(q->extra); GGML_ASSERT(k->extra); GGML_ASSERT(v->extra); GGML_ASSERT(dst->extra); + if (mask) { GGML_ASSERT(mask->extra); } @@ -11508,87 +12488,463 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co const int n_head_kv = k->ne[2]; const int n_batch = q->ne[3]; + // Per-variant lazy compile for this (dk, dv). + ggml_opencl_ensure_fa_pre_kernels(backend_ctx, d_head_q, d_head_v); + cl_kernel kernel = NULL; const bool is_f16 = q->type == GGML_TYPE_F16; - const bool is_mixed = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16; - const std::pair dk_dv = {d_head_q, d_head_v}; + const bool is_mixed = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16; + const bool is_q8_0 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_Q8_0 && v->type == GGML_TYPE_Q8_0; + const bool is_q4_0 = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_Q4_0 && v->type == GGML_TYPE_Q4_0; - if (n_q == 1) { - if (is_mixed) { - kernel = backend_ctx->kernels_flash_attn_f32_f16_q1.at(dk_dv); - } else if (is_f16) { - kernel = backend_ctx->kernels_flash_attn_f16_q1.at(dk_dv); + if (is_f16) { + ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_F16); + } else if (is_mixed) { + ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_F32_F16); + ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_F32_F16_SPLIT); + } else if (is_q8_0) { + ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_Q8_0); + if (d_head_q == 96 && d_head_v == 96) { + ggml_opencl_ensure_fa_quant_split_override(backend_ctx, 96, 96, /*quant_bm=*/16, /*quant_n_split=*/3, /*is_q8_0=*/true); + } else if (d_head_q == 256 && d_head_v == 256) { + ggml_opencl_ensure_fa_quant_split_override(backend_ctx, 256, 256, /*quant_bm=*/16, /*quant_n_split=*/8, /*is_q8_0=*/true); } else { - kernel = backend_ctx->kernels_flash_attn_f32_q1.at(dk_dv); + ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_Q8_0_SPLIT); + } + } else if (is_q4_0) { + ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_Q4_0); + if (d_head_q == 96 && d_head_v == 96) { + ggml_opencl_ensure_fa_quant_split_override(backend_ctx, 96, 96, /*quant_bm=*/16, /*quant_n_split=*/3, /*is_q8_0=*/false); + } else if (d_head_q == 256 && d_head_v == 256) { + ggml_opencl_ensure_fa_quant_split_override(backend_ctx, 256, 256, /*quant_bm=*/16, /*quant_n_split=*/8, /*is_q8_0=*/false); + } else { + ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_Q4_0_SPLIT); } } else { - if (is_mixed) { - kernel = backend_ctx->kernels_flash_attn_f32_f16.at(dk_dv); - } else if (is_f16) { - kernel = backend_ctx->kernels_flash_attn_f16.at(dk_dv); - } else { - kernel = backend_ctx->kernels_flash_attn_f32.at(dk_dv); - } + ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_F32); } - GGML_ASSERT(kernel != NULL); + + const std::pair dk_dv = {d_head_q, d_head_v}; + const bool use_native_q8_0_q1 = is_q8_0 && n_q == 1 && + backend_ctx->fa.f32_q8_0_q1.count(dk_dv) > 0; + // Native q8_0 prefill — reads q8_0 directly, wg_size = cfg->bm. + const bool use_native_q8_0 = is_q8_0 && n_q > 1 && + backend_ctx->fa.f32_q8_0.count(dk_dv) > 0; + const bool use_native_q4_0_q1 = is_q4_0 && n_q == 1 && + backend_ctx->fa.f32_q4_0_q1.count(dk_dv) > 0; + const bool use_native_q4_0 = is_q4_0 && n_q > 1 && + backend_ctx->fa.f32_q4_0.count(dk_dv) > 0; + const int block_m = n_q > 1 + ? (is_mixed ? backend_ctx->fa.f32_f16_bm.at(dk_dv) : backend_ctx->fa.bm.at(dk_dv)) + : 0; + const int block_n = is_mixed + ? backend_ctx->fa.f32_f16_bn.at(dk_dv) + : backend_ctx->fa.bn.at(dk_dv); + // Pick split variant only when n_kv crosses the per-(dk,dv) threshold. + const bool use_split_kernel = (n_q > 1 && is_mixed && + backend_ctx->fa.f32_f16_split.count(dk_dv) > 0 && + n_kv >= backend_ctx->fa.f32_f16_split_nkv_threshold.at(dk_dv)); + const bool use_split_q8_0 = (use_native_q8_0 && + backend_ctx->fa.f32_q8_0_split.count(dk_dv) > 0 && + n_kv >= backend_ctx->fa.f32_q8_0_split_nkv_threshold.at(dk_dv)); + const bool use_split_q4_0 = (use_native_q4_0 && + backend_ctx->fa.f32_q4_0_split.count(dk_dv) > 0 && + n_kv >= backend_ctx->fa.f32_q4_0_split_nkv_threshold.at(dk_dv)); + const int wg_size_fa = (n_q > 1 && is_mixed) + ? (use_split_kernel + ? backend_ctx->fa.f32_f16_split_wg_size.at(dk_dv) + : backend_ctx->fa.f32_f16_wg_size.at(dk_dv)) + : block_m; ggml_tensor_extra_cl * extra_q = (ggml_tensor_extra_cl *)q->extra; - ggml_tensor_extra_cl * extra_k = (ggml_tensor_extra_cl *)k->extra; - ggml_tensor_extra_cl * extra_v = (ggml_tensor_extra_cl *)v->extra; ggml_tensor_extra_cl * extra_o = (ggml_tensor_extra_cl *)dst->extra; ggml_tensor_extra_cl * extra_mask = mask ? (ggml_tensor_extra_cl *)mask->extra : NULL; ggml_tensor_extra_cl * extra_sinks = sinks ? (ggml_tensor_extra_cl *)sinks->extra : NULL; + // SoA q8_0/q4_0 K/V: data_device aliases the `q` subbuffer; reconstruct + // AoS into a temp buffer below. AoS tensors use extra_k/v->data_device. + const bool k_soa = ggml_cl_is_q8_0_soa(k) || ggml_cl_is_q4_0_soa(k); + const bool v_soa = ggml_cl_is_q8_0_soa(v) || ggml_cl_is_q4_0_soa(v); + ggml_tensor_extra_cl * extra_k = k_soa ? nullptr : (ggml_tensor_extra_cl *)k->extra; + ggml_tensor_extra_cl * extra_v = v_soa ? nullptr : (ggml_tensor_extra_cl *)v->extra; + cl_ulong offset_q = extra_q->offset + q->view_offs; - cl_ulong offset_k = extra_k->offset + k->view_offs; - cl_ulong offset_v = extra_v->offset + v->view_offs; + cl_ulong offset_k = k_soa ? 0 : extra_k->offset + k->view_offs; + cl_ulong offset_v = v_soa ? 0 : extra_v->offset + v->view_offs; cl_ulong offset_o = extra_o->offset + dst->view_offs; cl_mem mask_buffer = extra_mask ? extra_mask->data_device : NULL; cl_ulong offset_mask = extra_mask ? extra_mask->offset + mask->view_offs : 0; cl_mem sinks_buffer = extra_sinks ? extra_sinks->data_device : NULL; cl_ulong offset_sinks = extra_sinks ? extra_sinks->offset + sinks->view_offs : 0; - const cl_ulong q_nb1 = q->nb[1], q_nb2 = q->nb[2], q_nb3 = q->nb[3]; - const cl_ulong k_nb1 = k->nb[1], k_nb2 = k->nb[2], k_nb3 = k->nb[3]; - const cl_ulong v_nb1 = v->nb[1], v_nb2 = v->nb[2], v_nb3 = v->nb[3]; - const cl_ulong o_nb1 = dst->nb[1], o_nb2 = dst->nb[2], o_nb3 = dst->nb[3]; + const cl_ulong q_nb1 = q->nb[1]; + const cl_ulong q_nb2 = q->nb[2]; + const cl_ulong q_nb3 = q->nb[3]; + + cl_ulong k_nb1 = k->nb[1]; + cl_ulong k_nb2 = k->nb[2]; + cl_ulong k_nb3 = k->nb[3]; + + cl_ulong v_nb1 = v->nb[1]; + cl_ulong v_nb2 = v->nb[2]; + cl_ulong v_nb3 = v->nb[3]; + + const cl_ulong o_nb1 = dst->nb[1]; + const cl_ulong o_nb2 = dst->nb[2]; + const cl_ulong o_nb3 = dst->nb[3]; + const cl_ulong mask_nb1 = mask ? mask->nb[1] : 0; const cl_ulong mask_nb2 = mask ? mask->nb[2] : 0; const cl_ulong mask_nb3 = mask ? mask->nb[3] : 0; const int mask_ne2 = mask ? mask->ne[2] : 0; const int mask_ne3 = mask ? mask->ne[3] : 0; - float scale, max_bias, logit_softcap; + float scale; + float max_bias; + float logit_softcap; + const float * params = (const float *)dst->op_params; scale = params[0]; max_bias = params[1]; logit_softcap = params[2]; + if (n_q == 1) { + if (use_native_q8_0_q1) { + kernel = backend_ctx->fa.f32_q8_0_q1.at(dk_dv); + } else if (use_native_q4_0_q1) { + kernel = backend_ctx->fa.f32_q4_0_q1.at(dk_dv); + } else if (is_mixed) { + kernel = backend_ctx->fa.f32_f16_q1.at(dk_dv); + } else if (is_f16) { + kernel = backend_ctx->fa.f16_q1.at(dk_dv); + } else { + kernel = backend_ctx->fa.f32_q1.at(dk_dv); + } + } else { + if (use_native_q8_0) { + kernel = use_split_q8_0 + ? backend_ctx->fa.f32_q8_0_split.at(dk_dv) + : backend_ctx->fa.f32_q8_0.at(dk_dv); + } else if (use_native_q4_0) { + kernel = use_split_q4_0 + ? backend_ctx->fa.f32_q4_0_split.at(dk_dv) + : backend_ctx->fa.f32_q4_0.at(dk_dv); + } else if (is_mixed) { + kernel = use_split_kernel + ? backend_ctx->fa.f32_f16_split.at(dk_dv) + : backend_ctx->fa.f32_f16.at(dk_dv); + } else if (is_f16) { + kernel = backend_ctx->fa.f16.at(dk_dv); + } else { + kernel = backend_ctx->fa.f32.at(dk_dv); + } + } + GGML_ASSERT(kernel != NULL); + + ggml_cl_flash_attn_temp_buffer temp_k; + ggml_cl_flash_attn_temp_buffer temp_v; + ggml_cl_flash_attn_temp_buffer temp_k_pad; + ggml_cl_flash_attn_temp_buffer temp_v_pad; + ggml_cl_flash_attn_temp_buffer temp_mask_pad; + ggml_cl_flash_attn_temp_buffer temp_blk; + const ggml_type kv_target_type = is_f16 ? GGML_TYPE_F16 : GGML_TYPE_F32; + + cl_mem k_data_device = k_soa ? NULL : extra_k->data_device; + cl_mem v_data_device = v_soa ? NULL : extra_v->data_device; + + // SoA q8_0/q4_0 -> reconstruct AoS for downstream kernels that expect + // tight records (no-op when k/v is already AoS). + ggml_cl_flash_attn_temp_buffer temp_k_aos; + ggml_cl_flash_attn_temp_buffer temp_v_aos; + ggml_cl_flash_attn_reconstruct_aos(backend_ctx, k, temp_k_aos, + k_data_device, offset_k, k_nb1, k_nb2, k_nb3); + ggml_cl_flash_attn_reconstruct_aos(backend_ctx, v, temp_v_aos, + v_data_device, offset_v, v_nb1, v_nb2, v_nb3); + + // currently FA kernels support KV cache with f16, f32, q4_0 and q8_0. + // there two cases that these kernels cannot cover, + // 1. KV cache types are q4_0 or q8_0, but the FA kernels fail to compile + // 2. KV cache types not currently supported by an FA kernel, e.g., q4_1 + // these two cases are supported here by dequantizing to f32/f16 and this + // causes performance degradation. + // For q4_0 or q8_0 cases that fail kernel compilation, dequant happens in GPU; + // for types that do not have FA kernels, dequant happens on host. + if (!use_native_q8_0_q1 && !use_native_q8_0 && + !use_native_q4_0_q1 && !use_native_q4_0) { + // for q4_0, q8_0 FA kernels that fail to compile + bool k_done = false; + bool v_done = false; + if (k->type == GGML_TYPE_Q8_0 || k->type == GGML_TYPE_Q4_0) { + k_done = ggml_cl_flash_attn_dequant_kv_gpu( + backend_ctx, k, kv_target_type, k_data_device, offset_k, k_nb1, k_nb2, k_nb3, + temp_k, k_data_device, offset_k, k_nb1, k_nb2, k_nb3); + } + if (v->type == GGML_TYPE_Q8_0 || v->type == GGML_TYPE_Q4_0) { + v_done = ggml_cl_flash_attn_dequant_kv_gpu( + backend_ctx, v, kv_target_type, v_data_device, offset_v, v_nb1, v_nb2, v_nb3, + temp_v, v_data_device, offset_v, v_nb1, v_nb2, v_nb3); + } + if (!k_done) { + ggml_cl_flash_attn_prepare_quantized_tensor( + backend_ctx, k, kv_target_type, temp_k, k_data_device, offset_k, k_nb1, k_nb2, k_nb3); + } + if (!v_done) { + ggml_cl_flash_attn_prepare_quantized_tensor( + backend_ctx, v, kv_target_type, temp_v, v_data_device, offset_v, v_nb1, v_nb2, v_nb3); + } + // Asymmetric KV on the F32 fallback path: convert the F16 side to F32 + // too. (Symmetric F16 / mixed paths handle F16 directly.) + if (kv_target_type == GGML_TYPE_F32 && !is_mixed && !is_f16) { + ggml_cl_flash_attn_convert_f16_to_f32(backend_ctx, k, temp_k, k_data_device, offset_k, k_nb1, k_nb2, k_nb3); + ggml_cl_flash_attn_convert_f16_to_f32(backend_ctx, v, temp_v, v_data_device, offset_v, v_nb1, v_nb2, v_nb3); + } + } + + cl_mem k_pad_buffer = NULL; + cl_mem v_pad_buffer = NULL; + cl_mem mask_pad_buffer = NULL; + cl_mem blk_buffer = NULL; + cl_ulong mask_pad_nb1 = 0; + cl_ulong mask_pad_nb2 = 0; + cl_ulong mask_pad_nb3 = 0; + + // Flash-Decoding K-split decision. Resolved here, before the prefill + // prepass, because KV-pad and blk prepass are pure overhead when FD fires. const int is_causal = (mask == NULL && n_q > 1 && n_q == n_kv); + const int fd_max_n_q = (d_head_q <= FD_MAX_DK_MULTI) ? FD_MAX_N_Q_MULTI : 1; + cl_kernel fd_k_split = NULL; + if (n_q >= 1 && n_q <= fd_max_n_q && n_kv >= FD_MIN_N_KV && !is_causal && + d_head_q <= FD_MAX_DK && + backend_ctx->fa.f32_merge.count(dk_dv) > 0) { + if (is_mixed && backend_ctx->fa.f32_f16_q1_split.count(dk_dv) > 0) { + fd_k_split = backend_ctx->fa.f32_f16_q1_split.at(dk_dv); + } else if (is_q8_0 && backend_ctx->fa.f32_q8_0_q1_split.count(dk_dv) > 0) { + fd_k_split = backend_ctx->fa.f32_q8_0_q1_split.at(dk_dv); + } else if (is_q4_0 && backend_ctx->fa.f32_q4_0_q1_split.count(dk_dv) > 0) { + fd_k_split = backend_ctx->fa.f32_q4_0_q1_split.at(dk_dv); + } + } + const bool use_fd = (fd_k_split != NULL); + + const int n_q_blocks = n_q > 1 ? (n_q + block_m - 1) / block_m : 0; + const int n_kv_blocks = n_kv > 0 ? (n_kv + block_n - 1) / block_n : 0; + // KV pad + blk prepass are pure overhead when FD will fire — skip them. + const bool use_mixed_prepass = is_mixed && n_q > 1 && !use_fd; + const bool use_kv_pad = use_mixed_prepass && (n_kv % block_n != 0); + // blk prepass: per-KV-tile mask class (0=masked, 1=mixed, 2=unmasked). + // Consumed identically by f32_f16, q8_0 and q4_0 prefill kernels. + const bool use_quant_prepass = (use_native_q8_0 || use_native_q4_0) && !use_fd; + const bool use_blk_mask = (use_mixed_prepass || use_quant_prepass) && mask_buffer != NULL; + + if (use_kv_pad) { + cl_int err; + + const size_t k_pad_size = (size_t) k_nb1 * (size_t) block_n * (size_t) n_head_kv * (size_t) n_batch; + temp_k_pad.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, k_pad_size, NULL, &err); + CL_CHECK(err); + k_pad_buffer = temp_k_pad.data; + + const size_t v_pad_size = (size_t) v_nb1 * (size_t) block_n * (size_t) n_head_kv * (size_t) n_batch; + temp_v_pad.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, v_pad_size, NULL, &err); + CL_CHECK(err); + v_pad_buffer = temp_v_pad.data; + + cl_kernel kernel_kv_pad = backend_ctx->fa.kv_pad_f16.at(dk_dv); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 0, sizeof(cl_mem), &k_data_device)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 1, sizeof(cl_ulong), &offset_k)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 2, sizeof(cl_mem), &v_data_device)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 3, sizeof(cl_ulong), &offset_v)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 4, sizeof(cl_mem), &k_pad_buffer)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 5, sizeof(cl_mem), &v_pad_buffer)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 6, sizeof(int), &n_kv)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 7, sizeof(int), &n_head_kv)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 8, sizeof(int), &n_batch)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 9, sizeof(cl_ulong), &k_nb1)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 10, sizeof(cl_ulong), &k_nb2)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 11, sizeof(cl_ulong), &k_nb3)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 12, sizeof(cl_ulong), &v_nb1)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 13, sizeof(cl_ulong), &v_nb2)); + CL_CHECK(clSetKernelArg(kernel_kv_pad, 14, sizeof(cl_ulong), &v_nb3)); + + size_t global_work_size[] = { (size_t) block_n, (size_t) n_head_kv, (size_t) n_batch }; + backend_ctx->enqueue_ndrange_kernel(kernel_kv_pad, 3, global_work_size, NULL, dst); + + if (mask_buffer != NULL) { + mask_pad_nb1 = (cl_ulong) block_n * (cl_ulong) sizeof(ggml_fp16_t); + mask_pad_nb2 = (cl_ulong) n_q * mask_pad_nb1; + mask_pad_nb3 = (cl_ulong) mask_ne2 * mask_pad_nb2; + + const size_t mask_pad_size = (size_t) mask_ne3 * (size_t) mask_pad_nb3; + temp_mask_pad.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, mask_pad_size, NULL, &err); + CL_CHECK(err); + mask_pad_buffer = temp_mask_pad.data; + + cl_kernel kernel_mask_pad = backend_ctx->fa.mask_pad_f16.at(dk_dv); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 0, sizeof(cl_mem), &mask_buffer)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 1, sizeof(cl_ulong), &offset_mask)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 2, sizeof(cl_mem), &mask_pad_buffer)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 3, sizeof(int), &n_q)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 4, sizeof(int), &n_kv)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 5, sizeof(cl_ulong), &mask_nb1)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 6, sizeof(cl_ulong), &mask_nb2)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 7, sizeof(cl_ulong), &mask_nb3)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 8, sizeof(int), &mask_ne2)); + CL_CHECK(clSetKernelArg(kernel_mask_pad, 9, sizeof(int), &mask_ne3)); + + size_t global_work_size_mask[] = { (size_t) block_n, (size_t) n_q, (size_t) (mask_ne2 * mask_ne3) }; + backend_ctx->enqueue_ndrange_kernel(kernel_mask_pad, 3, global_work_size_mask, NULL, dst); + } + } + + if (use_blk_mask) { + cl_int err; + const size_t blk_size = (size_t) n_kv_blocks * (size_t) n_q_blocks * (size_t) mask_ne2 * (size_t) mask_ne3; + temp_blk.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, blk_size, NULL, &err); + if (err != CL_SUCCESS) { + // Flush before retry — reclaim deferred driver deallocations. + CL_CHECK(clFinish(backend_ctx->queue)); + temp_blk.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, blk_size, NULL, &err); + } + CL_CHECK(err); + blk_buffer = temp_blk.data; + + cl_kernel kernel_blk = backend_ctx->fa.blk_f16.at(dk_dv); + CL_CHECK(clSetKernelArg(kernel_blk, 0, sizeof(cl_mem), &mask_buffer)); + CL_CHECK(clSetKernelArg(kernel_blk, 1, sizeof(cl_ulong), &offset_mask)); + CL_CHECK(clSetKernelArg(kernel_blk, 2, sizeof(cl_mem), &blk_buffer)); + CL_CHECK(clSetKernelArg(kernel_blk, 3, sizeof(int), &n_q)); + CL_CHECK(clSetKernelArg(kernel_blk, 4, sizeof(int), &n_kv)); + CL_CHECK(clSetKernelArg(kernel_blk, 5, sizeof(cl_ulong), &mask_nb1)); + CL_CHECK(clSetKernelArg(kernel_blk, 6, sizeof(cl_ulong), &mask_nb2)); + CL_CHECK(clSetKernelArg(kernel_blk, 7, sizeof(cl_ulong), &mask_nb3)); + CL_CHECK(clSetKernelArg(kernel_blk, 8, sizeof(int), &mask_ne2)); + CL_CHECK(clSetKernelArg(kernel_blk, 9, sizeof(int), &mask_ne3)); + + size_t global_work_size_blk[] = { (size_t) n_kv_blocks, (size_t) n_q_blocks, (size_t) (mask_ne2 * mask_ne3) }; + backend_ctx->enqueue_ndrange_kernel(kernel_blk, 3, global_work_size_blk, NULL, dst); + } const int n_head_log2_val = n_head > 0 ? 1u << (int)floorf(log2f((float)n_head)) : 0; const float n_head_log2_f = n_head_log2_val > 0 ? (float)n_head_log2_val : 1.0f; const float m0 = powf(2.0f, -(max_bias) / n_head_log2_f); const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2_f); - CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra_q->data_device)); - CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset_q)); - CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extra_k->data_device)); - CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset_k)); - CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &extra_v->data_device)); - CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offset_v)); - CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_mem), &extra_o->data_device)); - CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &offset_o)); - CL_CHECK(clSetKernelArg(kernel, 8, sizeof(float), &scale)); - CL_CHECK(clSetKernelArg(kernel, 9, sizeof(int), &n_q)); - CL_CHECK(clSetKernelArg(kernel, 10, sizeof(int), &n_kv)); - CL_CHECK(clSetKernelArg(kernel, 11, sizeof(int), &is_causal)); - CL_CHECK(clSetKernelArg(kernel, 12, sizeof(int), &n_head)); - CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &q_nb1)); CL_CHECK(clSetKernelArg(kernel, 14, sizeof(cl_ulong), &q_nb2)); CL_CHECK(clSetKernelArg(kernel, 15, sizeof(cl_ulong), &q_nb3)); - CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &k_nb1)); CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &k_nb2)); CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &k_nb3)); - CL_CHECK(clSetKernelArg(kernel, 19, sizeof(cl_ulong), &v_nb1)); CL_CHECK(clSetKernelArg(kernel, 20, sizeof(cl_ulong), &v_nb2)); CL_CHECK(clSetKernelArg(kernel, 21, sizeof(cl_ulong), &v_nb3)); - CL_CHECK(clSetKernelArg(kernel, 22, sizeof(cl_ulong), &o_nb1)); CL_CHECK(clSetKernelArg(kernel, 23, sizeof(cl_ulong), &o_nb2)); CL_CHECK(clSetKernelArg(kernel, 24, sizeof(cl_ulong), &o_nb3)); + if (use_fd) { + int n_splits = (n_kv + FD_KV_PER_SPLIT - 1) / FD_KV_PER_SPLIT; + if (n_splits < FD_MIN_SPLITS) { + n_splits = FD_MIN_SPLITS; + } + if (n_splits > FD_MAX_SPLITS) { + n_splits = FD_MAX_SPLITS; + } + const int kv_per_split = (n_kv + n_splits - 1) / n_splits; + + const int fa_partial_floats = 2 + d_head_v; + const size_t partial_size_bytes = + (size_t) n_batch * n_head * n_q * n_splits * fa_partial_floats * sizeof(float); + + ggml_cl_flash_attn_temp_buffer temp_partial; + cl_int err; + temp_partial.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, + partial_size_bytes, NULL, &err); + if (err != CL_SUCCESS) { + CL_CHECK(clFinish(backend_ctx->queue)); + temp_partial.data = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, + partial_size_bytes, NULL, &err); + } + CL_CHECK(err); + + cl_kernel k_split = fd_k_split; + int argi = 0; + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &extra_q->data_device)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_q)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &k_data_device)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_k)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &v_data_device)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_v)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(float), &scale)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &n_q)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &n_kv)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &n_head)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &q_nb1)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &q_nb2)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &q_nb3)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &k_nb1)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &k_nb2)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &k_nb3)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &v_nb1)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &v_nb2)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &v_nb3)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(float), &max_bias)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(float), &m0)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(float), &m1)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &n_head_log2_val)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(float), &logit_softcap)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &n_head_kv)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &mask_buffer)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_mask)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &mask_nb1)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &mask_nb2)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &mask_nb3)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &mask_ne2)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &mask_ne3)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &temp_partial.data)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &n_splits)); + CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &kv_per_split)); + + const size_t fd_wg = 64; // matches Q1_WG_SIZE in the kernel + size_t fd_lws[3] = { fd_wg, 1, 1 }; + // gid(2) packs q_idx * n_splits + split_idx. + size_t fd_gws[3] = { fd_wg, (size_t)(n_head * n_batch), (size_t)(n_splits * n_q) }; + backend_ctx->enqueue_ndrange_kernel(k_split, 3, fd_gws, fd_lws, dst); + + cl_kernel k_merge = backend_ctx->fa.f32_merge.at(dk_dv); + argi = 0; + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(cl_mem), &temp_partial.data)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(cl_mem), &extra_o->data_device)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(cl_ulong), &offset_o)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(int), &n_head)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(int), &n_splits)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(cl_ulong), &o_nb1)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(cl_ulong), &o_nb2)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(cl_ulong), &o_nb3)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(cl_mem), &sinks_buffer)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(cl_ulong), &offset_sinks)); + CL_CHECK(clSetKernelArg(k_merge, argi++, sizeof(int), &n_q)); + + const size_t merge_wg = (size_t) (d_head_v / 4); // one lane per float4 + size_t merge_lws[3] = { merge_wg, 1, 1 }; + size_t merge_gws[3] = { merge_wg, (size_t)(n_head * n_batch), (size_t) n_q }; + backend_ctx->enqueue_ndrange_kernel(k_merge, 3, merge_gws, merge_lws, dst); + return; + } + + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra_q->data_device)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset_q)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &k_data_device)); + CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset_k)); + CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &v_data_device)); + CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offset_v)); + CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_mem), &extra_o->data_device)); + CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &offset_o)); + CL_CHECK(clSetKernelArg(kernel, 8, sizeof(float), &scale)); + CL_CHECK(clSetKernelArg(kernel, 9, sizeof(int), &n_q)); + CL_CHECK(clSetKernelArg(kernel, 10, sizeof(int), &n_kv)); + CL_CHECK(clSetKernelArg(kernel, 11, sizeof(int), &is_causal)); + CL_CHECK(clSetKernelArg(kernel, 12, sizeof(int), &n_head)); + CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &q_nb1)); + CL_CHECK(clSetKernelArg(kernel, 14, sizeof(cl_ulong), &q_nb2)); + CL_CHECK(clSetKernelArg(kernel, 15, sizeof(cl_ulong), &q_nb3)); + CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &k_nb1)); + CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &k_nb2)); + CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &k_nb3)); + CL_CHECK(clSetKernelArg(kernel, 19, sizeof(cl_ulong), &v_nb1)); + CL_CHECK(clSetKernelArg(kernel, 20, sizeof(cl_ulong), &v_nb2)); + CL_CHECK(clSetKernelArg(kernel, 21, sizeof(cl_ulong), &v_nb3)); + CL_CHECK(clSetKernelArg(kernel, 22, sizeof(cl_ulong), &o_nb1)); + CL_CHECK(clSetKernelArg(kernel, 23, sizeof(cl_ulong), &o_nb2)); + CL_CHECK(clSetKernelArg(kernel, 24, sizeof(cl_ulong), &o_nb3)); CL_CHECK(clSetKernelArg(kernel, 25, sizeof(float), &max_bias)); CL_CHECK(clSetKernelArg(kernel, 26, sizeof(float), &m0)); CL_CHECK(clSetKernelArg(kernel, 27, sizeof(float), &m1)); @@ -11604,15 +12960,45 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co CL_CHECK(clSetKernelArg(kernel, 37, sizeof(int), &mask_ne3)); CL_CHECK(clSetKernelArg(kernel, 38, sizeof(cl_mem), &sinks_buffer)); CL_CHECK(clSetKernelArg(kernel, 39, sizeof(cl_ulong), &offset_sinks)); + if (n_q > 1 && is_mixed) { + CL_CHECK(clSetKernelArg(kernel, 40, sizeof(cl_mem), &k_pad_buffer)); + CL_CHECK(clSetKernelArg(kernel, 41, sizeof(cl_mem), &v_pad_buffer)); + CL_CHECK(clSetKernelArg(kernel, 42, sizeof(cl_mem), &mask_pad_buffer)); + CL_CHECK(clSetKernelArg(kernel, 43, sizeof(cl_mem), &blk_buffer)); + CL_CHECK(clSetKernelArg(kernel, 44, sizeof(int), &n_kv_blocks)); + CL_CHECK(clSetKernelArg(kernel, 45, sizeof(cl_ulong), &mask_pad_nb1)); + CL_CHECK(clSetKernelArg(kernel, 46, sizeof(cl_ulong), &mask_pad_nb2)); + CL_CHECK(clSetKernelArg(kernel, 47, sizeof(cl_ulong), &mask_pad_nb3)); + } else if (use_native_q8_0 || use_native_q4_0) { + // arg 40 = blk classification buffer (NULL disables prepass opt). + CL_CHECK(clSetKernelArg(kernel, 40, sizeof(cl_mem), &blk_buffer)); + } if (n_q == 1) { const size_t wg_size = 64; size_t local_work_size[] = { wg_size, 1 }; size_t global_work_size[] = { wg_size, (size_t)(n_head * n_batch) }; backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size, local_work_size, dst); + } else if (use_native_q8_0 || use_native_q4_0) { + // Native quant prefill. The split variant may override BLOCK_M + // (e.g. DK=96 quant uses BM=16). + const bool use_split = use_native_q8_0 ? use_split_q8_0 : use_split_q4_0; + int bm; + size_t wg_size; + if (use_split) { + bm = use_native_q8_0 ? backend_ctx->fa.f32_q8_0_split_bm.at(dk_dv) + : backend_ctx->fa.f32_q4_0_split_bm.at(dk_dv); + wg_size = use_native_q8_0 ? backend_ctx->fa.f32_q8_0_split_wg_size.at(dk_dv) + : backend_ctx->fa.f32_q4_0_split_wg_size.at(dk_dv); + } else { + bm = backend_ctx->fa.bm.at(dk_dv); + wg_size = (size_t) bm; + } + size_t local_work_size[] = { wg_size, 1 }; + size_t global_work_size[] = { (size_t)((n_q + bm - 1) / bm) * wg_size, (size_t)(n_head * n_batch) }; + backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size, local_work_size, dst); } else { - const int block_m = backend_ctx->kernels_flash_attn_bm.at(dk_dv); - const size_t wg_size = block_m; + const size_t wg_size = (size_t) wg_size_fa; size_t local_work_size[] = { wg_size, 1 }; size_t global_work_size[] = { (size_t)((n_q + block_m - 1) / block_m) * wg_size, (size_t)(n_head * n_batch) }; backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size, local_work_size, dst); @@ -13004,7 +14390,9 @@ static void ggml_cl_mul_mat_q8_0_f32_adreno(ggml_backend_t backend, const ggml_t ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *)src1->extra; ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra; - ggml_tensor_extra_cl_q8_0 * extra0_q8_0 = (ggml_tensor_extra_cl_q8_0 *)src0->extra; + // SoA extra lives on view_src (view->extra is pre-SoA). + const ggml_tensor * soa0_src = src0->view_src != nullptr ? src0->view_src : src0; + ggml_tensor_extra_cl_q8_0 * extra0_q8_0 = (ggml_tensor_extra_cl_q8_0 *)soa0_src->extra; cl_ulong offset1 = extra1->offset + src1->view_offs; cl_ulong offsetd = extrad->offset + dst->view_offs; @@ -13756,6 +15144,122 @@ static void ggml_cl_mul_mat_q5_K_f32_adreno(ggml_backend_t backend, const ggml_t #endif } +// Dequant a possibly-strided q4_0/q8_0 tensor to tight-packed f16. Returns a +// temp cl_mem the caller must release. SoA inputs are reconstructed into a +// temp AoS buffer reported via *extra_reconstruct (also caller-released). +// this is for quantized K cache without FA. +static cl_mem ggml_cl_mul_mat_dequant_quant_to_f16( + ggml_backend_opencl_context * backend_ctx, + const ggml_tensor * tensor, + cl_mem * extra_reconstruct /* out, may be NULL */ +) { + GGML_ASSERT(tensor->type == GGML_TYPE_Q4_0 || tensor->type == GGML_TYPE_Q8_0); + + if (extra_reconstruct) { + *extra_reconstruct = NULL; + } + + cl_mem src_buf; + cl_ulong src_offset; + cl_ulong src_nb1; + cl_ulong src_nb2; + cl_ulong src_nb3; + + const bool is_soa = tensor->type == GGML_TYPE_Q4_0 + ? ggml_cl_is_q4_0_soa(tensor) + : ggml_cl_is_q8_0_soa(tensor); + + if (is_soa) { + // Reconstruct full parent AoS; view's own nb[] then index it correctly. + const ggml_tensor * parent = tensor->view_src ? tensor->view_src : tensor; + const ggml_tensor * soa_src = parent; + const size_t block_bytes = (size_t) ggml_type_size(tensor->type); + const size_t blck_size = (size_t) ggml_blck_size(tensor->type); + const size_t parent_row_blocks = (size_t) parent->ne[0] / blck_size; + const size_t parent_row_bytes = parent_row_blocks * block_bytes; + const size_t parent_nbytes = (size_t) ggml_nelements(parent) / blck_size * block_bytes; + + cl_int err; + cl_mem aos = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, parent_nbytes, NULL, &err); + CL_CHECK(err); + + cl_kernel kernel; + if (tensor->type == GGML_TYPE_Q8_0) { + auto * extra = (ggml_tensor_extra_cl_q8_0 *) soa_src->extra; + kernel = backend_ctx->kernel_restore_block_q8_0; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra->q)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra->d)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &aos)); + } else { + auto * extra = (ggml_tensor_extra_cl_q4_0 *) soa_src->extra; + kernel = backend_ctx->kernel_restore_block_q4_0; + CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra->q)); + CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra->d)); + CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &aos)); + } + + const size_t n_blocks = parent_nbytes / block_bytes; + size_t gws_rec[] = { n_blocks, 1, 1 }; + size_t lws_rec[] = { 1, 1, 1 }; + CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL, gws_rec, lws_rec, 0, NULL, NULL)); + + (void) parent_row_blocks; + (void) parent_row_bytes; + src_buf = aos; + src_offset = tensor->view_offs; + src_nb1 = tensor->nb[1]; + src_nb2 = tensor->nb[2]; + src_nb3 = tensor->nb[3]; + + if (extra_reconstruct) { + *extra_reconstruct = aos; + } else { + // OpenCL retains the memobj while queued kernels reference it. + CL_CHECK(clReleaseMemObject(aos)); + } + } else { + auto * extra = (ggml_tensor_extra_cl *) tensor->extra; + GGML_ASSERT(extra && extra->data_device); + src_buf = extra->data_device; + src_offset = extra->offset + tensor->view_offs; + src_nb1 = tensor->nb[1]; + src_nb2 = tensor->nb[2]; + src_nb3 = tensor->nb[3]; + } + + const cl_int nblk0 = (cl_int) (tensor->ne[0] / ggml_blck_size(tensor->type)); + const cl_int ne1_ = (cl_int) tensor->ne[1]; + const cl_int ne2_ = (cl_int) tensor->ne[2]; + const cl_int ne3_ = (cl_int) tensor->ne[3]; + + const size_t out_bytes = (size_t) ggml_nelements(tensor) * sizeof(ggml_fp16_t); + + cl_int err; + cl_mem out = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, out_bytes, NULL, &err); + CL_CHECK(err); + + cl_kernel dq_kernel = tensor->type == GGML_TYPE_Q8_0 + ? backend_ctx->kernel_dequant_q8_0_f16_view_aos + : backend_ctx->kernel_dequant_q4_0_f16_view_aos; + + CL_CHECK(clSetKernelArg(dq_kernel, 0, sizeof(cl_mem), &src_buf)); + CL_CHECK(clSetKernelArg(dq_kernel, 1, sizeof(cl_ulong), &src_offset)); + CL_CHECK(clSetKernelArg(dq_kernel, 2, sizeof(cl_ulong), &src_nb1)); + CL_CHECK(clSetKernelArg(dq_kernel, 3, sizeof(cl_ulong), &src_nb2)); + CL_CHECK(clSetKernelArg(dq_kernel, 4, sizeof(cl_ulong), &src_nb3)); + CL_CHECK(clSetKernelArg(dq_kernel, 5, sizeof(cl_int), &nblk0)); + CL_CHECK(clSetKernelArg(dq_kernel, 6, sizeof(cl_int), &ne1_)); + CL_CHECK(clSetKernelArg(dq_kernel, 7, sizeof(cl_int), &ne2_)); + CL_CHECK(clSetKernelArg(dq_kernel, 8, sizeof(cl_int), &ne3_)); + CL_CHECK(clSetKernelArg(dq_kernel, 9, sizeof(cl_mem), &out)); + + size_t gws[3] = { (size_t) nblk0, (size_t) ne1_, (size_t) (ne2_ * ne3_) }; + size_t lws[3] = { 1, 1, 1 }; + CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, dq_kernel, 3, NULL, gws, lws, 0, NULL, NULL)); + + return out; +} + static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { GGML_ASSERT(src0); GGML_ASSERT(src0->extra); @@ -13770,6 +15274,31 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context; + // Non-contig quant src0: on-device dequant to f16 then native f16 MUL_MAT. + if ((src0t == GGML_TYPE_Q4_0 || src0t == GGML_TYPE_Q8_0) && !ggml_is_contiguous(src0)) { + cl_mem f16_buf = ggml_cl_mul_mat_dequant_quant_to_f16(backend_ctx, src0, nullptr); + + ggml_tensor fake_src0 = *src0; + ggml_tensor_extra_cl fake_extra = {}; + fake_extra.data_device = f16_buf; + fake_extra.offset = 0; + fake_src0.type = GGML_TYPE_F16; + fake_src0.extra = &fake_extra; + fake_src0.view_src = nullptr; + fake_src0.view_offs = 0; + fake_src0.nb[0] = sizeof(ggml_fp16_t); + fake_src0.nb[1] = fake_src0.nb[0] * src0->ne[0]; + fake_src0.nb[2] = fake_src0.nb[1] * src0->ne[1]; + fake_src0.nb[3] = fake_src0.nb[2] * src0->ne[2]; + + ggml_cl_mul_mat(backend, &fake_src0, src1, dst); + + // Safe to release now: OpenCL retains the memobj while queued + // kernels that reference it are still in flight. + CL_CHECK(clReleaseMemObject(f16_buf)); + return; + } + ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *)src0->extra; ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *)src1->extra; ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra; @@ -13779,16 +15308,19 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co cl_ulong offsetd = extrad->offset + dst->view_offs; #ifdef GGML_OPENCL_SOA_Q - ggml_tensor_extra_cl_q4_0 * extra0_q4_0 = (ggml_tensor_extra_cl_q4_0 *)src0->extra; - ggml_tensor_extra_cl_q4_1 * extra0_q4_1 = (ggml_tensor_extra_cl_q4_1 *)src0->extra; - ggml_tensor_extra_cl_q5_0 * extra0_q5_0 = (ggml_tensor_extra_cl_q5_0 *)src0->extra; - ggml_tensor_extra_cl_q5_1 * extra0_q5_1 = (ggml_tensor_extra_cl_q5_1 *)src0->extra; - ggml_tensor_extra_cl_mxfp4 * extra0_mxfp4 = (ggml_tensor_extra_cl_mxfp4 *)src0->extra; - ggml_tensor_extra_cl_q8_0 * extra0_q8_0 = (ggml_tensor_extra_cl_q8_0 *)src0->extra; - ggml_tensor_extra_cl_iq4_nl * extra0_iq4_nl = (ggml_tensor_extra_cl_iq4_nl *)src0->extra; - ggml_tensor_extra_cl_q4_K * extra0_q4_K = (ggml_tensor_extra_cl_q4_K *)src0->extra; - ggml_tensor_extra_cl_q5_K * extra0_q5_K = (ggml_tensor_extra_cl_q5_K *)src0->extra; - ggml_tensor_extra_cl_q6_K * extra0_q6_K = (ggml_tensor_extra_cl_q6_K *)src0->extra; + // view->extra stays pre-SoA; cast to the SoA struct would SIGSEGV. + // Follow view_src to reach the real SoA extra. + const ggml_tensor * soa0_src = src0->view_src != nullptr ? src0->view_src : src0; + ggml_tensor_extra_cl_q4_0 * extra0_q4_0 = (ggml_tensor_extra_cl_q4_0 *)soa0_src->extra; + ggml_tensor_extra_cl_q4_1 * extra0_q4_1 = (ggml_tensor_extra_cl_q4_1 *)soa0_src->extra; + ggml_tensor_extra_cl_q5_0 * extra0_q5_0 = (ggml_tensor_extra_cl_q5_0 *)soa0_src->extra; + ggml_tensor_extra_cl_q5_1 * extra0_q5_1 = (ggml_tensor_extra_cl_q5_1 *)soa0_src->extra; + ggml_tensor_extra_cl_mxfp4 * extra0_mxfp4 = (ggml_tensor_extra_cl_mxfp4 *)soa0_src->extra; + ggml_tensor_extra_cl_q8_0 * extra0_q8_0 = (ggml_tensor_extra_cl_q8_0 *)soa0_src->extra; + ggml_tensor_extra_cl_iq4_nl * extra0_iq4_nl = (ggml_tensor_extra_cl_iq4_nl *)soa0_src->extra; + ggml_tensor_extra_cl_q4_K * extra0_q4_K = (ggml_tensor_extra_cl_q4_K *)soa0_src->extra; + ggml_tensor_extra_cl_q5_K * extra0_q5_K = (ggml_tensor_extra_cl_q5_K *)soa0_src->extra; + ggml_tensor_extra_cl_q6_K * extra0_q6_K = (ggml_tensor_extra_cl_q6_K *)soa0_src->extra; #endif GGML_TENSOR_LOCALS(int, ne0, src0, ne); @@ -15543,15 +17075,18 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, GGML_UNUSED(offset0); #ifdef GGML_OPENCL_SOA_Q - ggml_tensor_extra_cl_q4_0 * extra0_q4_0 = (ggml_tensor_extra_cl_q4_0 *)src0->extra; - ggml_tensor_extra_cl_q4_1 * extra0_q4_1 = (ggml_tensor_extra_cl_q4_1 *)src0->extra; - ggml_tensor_extra_cl_q5_0 * extra0_q5_0 = (ggml_tensor_extra_cl_q5_0 *)src0->extra; - ggml_tensor_extra_cl_q5_1 * extra0_q5_1 = (ggml_tensor_extra_cl_q5_1 *)src0->extra; - ggml_tensor_extra_cl_q4_K * extra0_q4_K = (ggml_tensor_extra_cl_q4_K *)src0->extra; - ggml_tensor_extra_cl_q5_K * extra0_q5_K = (ggml_tensor_extra_cl_q5_K *)src0->extra; - ggml_tensor_extra_cl_q6_K * extra0_q6_K = (ggml_tensor_extra_cl_q6_K *)src0->extra; - ggml_tensor_extra_cl_mxfp4 * extra0_mxfp4 = (ggml_tensor_extra_cl_mxfp4 *)src0->extra; - ggml_tensor_extra_cl_q8_0 * extra0_q8_0 = (ggml_tensor_extra_cl_q8_0 *)src0->extra; + // SoA extra lives on view_src (view->extra is pre-SoA). + const ggml_tensor * soa0_src = src0->view_src != nullptr ? src0->view_src : src0; + ggml_tensor_extra_cl_q4_0 * extra0_q4_0 = (ggml_tensor_extra_cl_q4_0 *)soa0_src->extra; + ggml_tensor_extra_cl_q4_1 * extra0_q4_1 = (ggml_tensor_extra_cl_q4_1 *)soa0_src->extra; + ggml_tensor_extra_cl_q5_0 * extra0_q5_0 = (ggml_tensor_extra_cl_q5_0 *)soa0_src->extra; + ggml_tensor_extra_cl_q5_1 * extra0_q5_1 = (ggml_tensor_extra_cl_q5_1 *)soa0_src->extra; + ggml_tensor_extra_cl_q4_K * extra0_q4_K = (ggml_tensor_extra_cl_q4_K *)soa0_src->extra; + ggml_tensor_extra_cl_q5_K * extra0_q5_K = (ggml_tensor_extra_cl_q5_K *)soa0_src->extra; + ggml_tensor_extra_cl_q6_K * extra0_q6_K = (ggml_tensor_extra_cl_q6_K *)soa0_src->extra; + ggml_tensor_extra_cl_mxfp4 * extra0_mxfp4 = (ggml_tensor_extra_cl_mxfp4 *)soa0_src->extra; + ggml_tensor_extra_cl_q8_0 * extra0_q8_0 = (ggml_tensor_extra_cl_q8_0 *)soa0_src->extra; + #endif // TODO: general MoE for the following types diff --git a/ggml/src/ggml-opencl/kernels/cvt.cl b/ggml/src/ggml-opencl/kernels/cvt.cl index 226b127ab..82a130559 100644 --- a/ggml/src/ggml-opencl/kernels/cvt.cl +++ b/ggml/src/ggml-opencl/kernels/cvt.cl @@ -1582,6 +1582,158 @@ kernel void kernel_restore_block_q8_0( } } +// View-aware AoS q8_0 -> f32 dequant (f32/f32 FA path). +kernel void kernel_dequant_q8_0_f32_view_aos( + global char * src, + ulong src_offset, + ulong src_nb1, + ulong src_nb2, + ulong src_nb3, + int nblk0, + int ne1, + int ne2, + int ne3, + global float * dst +) { + int blk_i0 = get_global_id(0); + int i1 = get_global_id(1); + int batch = get_global_id(2); + + if (blk_i0 >= nblk0) return; + if (i1 >= ne1) return; + + int i2 = batch % ne2; + int i3 = batch / ne2; + if (i3 >= ne3) return; + + global char * block = src + src_offset + (ulong)i3*src_nb3 + (ulong)i2*src_nb2 + (ulong)i1*src_nb1 + (ulong)blk_i0 * (2 + QK8_0); + float d = vload_half(0, (global half *)block); + global char * qs = block + 2; + + ulong dst_row_base = ((ulong)i3 * ne2 * ne1 + (ulong)i2 * ne1 + (ulong)i1) * nblk0; + global float * out = dst + (dst_row_base + blk_i0) * QK8_0; + + for (int i = 0; i < QK8_0; ++i) { + out[i] = d * (float)qs[i]; + } +} + +// View-aware AoS q8_0 -> f16 dequant. Rows tight, batch strides may be gapped. +kernel void kernel_dequant_q8_0_f16_view_aos( + global char * src, + ulong src_offset, + ulong src_nb1, + ulong src_nb2, + ulong src_nb3, + int nblk0, + int ne1, + int ne2, + int ne3, + global half * dst +) { + int blk_i0 = get_global_id(0); + int i1 = get_global_id(1); + int batch = get_global_id(2); + + if (blk_i0 >= nblk0) return; + if (i1 >= ne1) return; + + int i2 = batch % ne2; + int i3 = batch / ne2; + if (i3 >= ne3) return; + + global char * block = src + src_offset + (ulong)i3*src_nb3 + (ulong)i2*src_nb2 + (ulong)i1*src_nb1 + (ulong)blk_i0 * (2 + QK8_0); + float d = vload_half(0, (global half *)block); + global char * qs = block + 2; + + ulong dst_row_base = ((ulong)i3 * ne2 * ne1 + (ulong)i2 * ne1 + (ulong)i1) * nblk0; + global half * out = dst + (dst_row_base + blk_i0) * QK8_0; + + for (int i = 0; i < QK8_0; ++i) { + out[i] = (half)(d * (float)qs[i]); + } +} + +// View-aware AoS q4_0 -> f32 dequant (mirrors the q8_0 view variant). +kernel void kernel_dequant_q4_0_f32_view_aos( + global char * src, + ulong src_offset, + ulong src_nb1, + ulong src_nb2, + ulong src_nb3, + int nblk0, + int ne1, + int ne2, + int ne3, + global float * dst +) { + int blk_i0 = get_global_id(0); + int i1 = get_global_id(1); + int batch = get_global_id(2); + + if (blk_i0 >= nblk0) return; + if (i1 >= ne1) return; + + int i2 = batch % ne2; + int i3 = batch / ne2; + if (i3 >= ne3) return; + + global char * block = src + src_offset + (ulong)i3*src_nb3 + (ulong)i2*src_nb2 + (ulong)i1*src_nb1 + (ulong)blk_i0 * (2 + QK4_0/2); + float d = vload_half(0, (global half *)block); + global uchar * qs = (global uchar *)(block + 2); + + ulong dst_row_base = ((ulong)i3 * ne2 * ne1 + (ulong)i2 * ne1 + (ulong)i1) * nblk0; + global float * out = dst + (dst_row_base + blk_i0) * QK4_0; + + for (int i = 0; i < QK4_0/2; ++i) { + uchar byte = qs[i]; + int q0 = (int)(byte & 0x0F) - 8; + int q1 = (int)(byte >> 4) - 8; + out[i] = d * (float)q0; + out[i + QK4_0/2] = d * (float)q1; + } +} + +// View-aware AoS q4_0 -> f16 dequant (mirrors the q8_0 view variant). +kernel void kernel_dequant_q4_0_f16_view_aos( + global char * src, + ulong src_offset, + ulong src_nb1, + ulong src_nb2, + ulong src_nb3, + int nblk0, + int ne1, + int ne2, + int ne3, + global half * dst +) { + int blk_i0 = get_global_id(0); + int i1 = get_global_id(1); + int batch = get_global_id(2); + + if (blk_i0 >= nblk0) return; + if (i1 >= ne1) return; + + int i2 = batch % ne2; + int i3 = batch / ne2; + if (i3 >= ne3) return; + + global char * block = src + src_offset + (ulong)i3*src_nb3 + (ulong)i2*src_nb2 + (ulong)i1*src_nb1 + (ulong)blk_i0 * (2 + QK4_0/2); + float d = vload_half(0, (global half *)block); + global uchar * qs = (global uchar *)(block + 2); + + ulong dst_row_base = ((ulong)i3 * ne2 * ne1 + (ulong)i2 * ne1 + (ulong)i1) * nblk0; + global half * out = dst + (dst_row_base + blk_i0) * QK4_0; + + for (int i = 0; i < QK4_0/2; ++i) { + uchar byte = qs[i]; + int q0 = (int)(byte & 0x0F) - 8; + int q1 = (int)(byte >> 4) - 8; + out[i] = (half)(d * (float)q0); + out[i + QK4_0/2] = (half)(d * (float)q1); + } +} + kernel void kernel_restore_block_q8_0_trans( global uchar * src_q, global half * src_d, diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl index 8f43c4f27..ec941b5f1 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl @@ -4,14 +4,26 @@ #define ACC_TYPE4 float4 #define DATA_TYPE half #define DATA_TYPE4 half4 -#define CONVERT_ACC4(x) convert_float4(x) -#define CONVERT_DATA4(x) convert_half4(x) +#define CONVERT_ACC4(x) ((float4)((float)(x).s0, (float)(x).s1, (float)(x).s2, (float)(x).s3)) +#define CONVERT_DATA4(x) ((half4)((half)(x).s0, (half)(x).s1, (half)(x).s2, (half)(x).s3)) #define DK_VEC (DK/4) #define DV_VEC (DV/4) #define WG_SIZE (BLOCK_M) #define Q1_WG_SIZE 64 +// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs, +// infinite operand can cause undefined behavior and miscompilation for exp. +// Therefore, a large negative value is used instead. +#define FA_M_INIT (-3.0e38f) + +// Drop full unroll at DK>=192 — Adreno compiler host-memory budget. +#if DK >= 192 +#define FA_UNROLL +#else +#define FA_UNROLL _Pragma("unroll") +#endif + inline float get_alibi_slope( const float max_bias, const uint h, const uint n_head_log2, const float m0, const float m1 ) { @@ -81,18 +93,18 @@ __kernel void flash_attn_f16( if (my_query_row < n_q) { const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + my_query_row * q_nb1; const global DATA_TYPE4* q_ptr = (const global DATA_TYPE4*)(q_base + q_row_offset); - #pragma unroll + FA_UNROLL for (int i = 0; i < DK_VEC; ++i) { q_priv[i] = CONVERT_ACC4(q_ptr[i]); } } ACC_TYPE4 o_acc[DV_VEC]; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_acc[i] = (ACC_TYPE4)(0.0f); } - ACC_TYPE m_i = -INFINITY; + ACC_TYPE m_i = FA_M_INIT; ACC_TYPE l_i = 0.0f; float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); @@ -125,49 +137,72 @@ __kernel void flash_attn_f16( continue; } - for (int j = 0; j < BLOCK_N; j += 2) { + for (int j = 0; j < BLOCK_N; j += 4) { const int k_row0 = k_start + j; const int k_row1 = k_start + j + 1; + const int k_row2 = k_start + j + 2; + const int k_row3 = k_start + j + 3; ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); ACC_TYPE4 dot_acc1 = (ACC_TYPE4)(0.0f); - #pragma unroll + ACC_TYPE4 dot_acc2 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc3 = (ACC_TYPE4)(0.0f); + FA_UNROLL for (int k = 0; k < DK_VEC; k++) { - dot_acc0 = mad(q_priv[k], CONVERT_ACC4(l_k[j][k]), dot_acc0); - dot_acc1 = mad(q_priv[k], CONVERT_ACC4(l_k[j+1][k]), dot_acc1); + const ACC_TYPE4 qk = q_priv[k]; + dot_acc0 = mad(qk, CONVERT_ACC4(l_k[j][k]), dot_acc0); + dot_acc1 = mad(qk, CONVERT_ACC4(l_k[j+1][k]), dot_acc1); + dot_acc2 = mad(qk, CONVERT_ACC4(l_k[j+2][k]), dot_acc2); + dot_acc3 = mad(qk, CONVERT_ACC4(l_k[j+3][k]), dot_acc3); } - ACC_TYPE score0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; - ACC_TYPE score1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; + ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; + ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; + ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale; + ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale; if (is_causal) { - if (k_row0 > (n_kv - n_q + my_query_row)) score0 = -INFINITY; - if (k_row1 > (n_kv - n_q + my_query_row)) score1 = -INFINITY; + const int causal_limit = n_kv - n_q + my_query_row; + if (k_row0 > causal_limit) s0 = FA_M_INIT; + if (k_row1 > causal_limit) s1 = FA_M_INIT; + if (k_row2 > causal_limit) s2 = FA_M_INIT; + if (k_row3 > causal_limit) s3 = FA_M_INIT; } - - if (k_row0 >= n_kv) score0 = -INFINITY; - if (k_row1 >= n_kv) score1 = -INFINITY; + if (k_row0 >= n_kv) s0 = FA_M_INIT; + if (k_row1 >= n_kv) s1 = FA_M_INIT; + if (k_row2 >= n_kv) s2 = FA_M_INIT; + if (k_row3 >= n_kv) s3 = FA_M_INIT; if (mask_base != NULL) { const global DATA_TYPE* mask_ptr = (const global DATA_TYPE*)(mask_base + my_query_row * mask_nb1); - if (k_row0 < n_kv) score0 += slope * (ACC_TYPE)mask_ptr[k_row0]; - if (k_row1 < n_kv) score1 += slope * (ACC_TYPE)mask_ptr[k_row1]; + if (k_row0 < n_kv) s0 += slope * (ACC_TYPE)mask_ptr[k_row0]; + if (k_row1 < n_kv) s1 += slope * (ACC_TYPE)mask_ptr[k_row1]; + if (k_row2 < n_kv) s2 += slope * (ACC_TYPE)mask_ptr[k_row2]; + if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3]; } if (logit_softcap > 0.0f) { - score0 = logit_softcap * tanh(score0 / logit_softcap); - score1 = logit_softcap * tanh(score1 / logit_softcap); + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); } - const ACC_TYPE m_new = max(m_i, max(score0, score1)); - const ACC_TYPE p0 = exp(score0 - m_new); - const ACC_TYPE p1 = exp(score1 - m_new); - const ACC_TYPE scale_prev = exp(m_i - m_new); + const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); + const ACC_TYPE scale_prev = native_exp(m_i - m_new); + const ACC_TYPE p0 = native_exp(s0 - m_new); + const ACC_TYPE p1 = native_exp(s1 - m_new); + const ACC_TYPE p2 = native_exp(s2 - m_new); + const ACC_TYPE p3 = native_exp(s3 - m_new); - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { - o_acc[i] = o_acc[i] * scale_prev + p0 * CONVERT_ACC4(l_v[j][i]) + p1 * CONVERT_ACC4(l_v[j+1][i]); + o_acc[i] = mad(p3, CONVERT_ACC4(l_v[j+3][i]), + mad(p2, CONVERT_ACC4(l_v[j+2][i]), + mad(p1, CONVERT_ACC4(l_v[j+1][i]), + mad(p0, CONVERT_ACC4(l_v[j][i]), + o_acc[i] * scale_prev)))); } - l_i = l_i * scale_prev + p0 + p1; + l_i = l_i * scale_prev + p0 + p1 + p2 + p3; m_i = m_new; } } @@ -179,7 +214,7 @@ __kernel void flash_attn_f16( const ACC_TYPE m_final = max(m_i, m_sink); const ACC_TYPE scale_o = exp(m_i - m_final); - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_acc[i] *= scale_o; } @@ -191,12 +226,12 @@ __kernel void flash_attn_f16( global DATA_TYPE4 *o_row = (global DATA_TYPE4 *)(o_base + o_row_offset); if (l_i > 0.0f) { const ACC_TYPE l_inv = 1.0f / l_i; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_row[i] = CONVERT_DATA4(o_acc[i] * l_inv); } } else { - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_row[i] = (DATA_TYPE4)(0.0f); } @@ -258,7 +293,7 @@ __kernel void flash_attn_f16_q1( ACC_TYPE4 q_priv[DK_VEC]; const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2; const global DATA_TYPE4* q_ptr = (const global DATA_TYPE4*)(q_base + q_row_offset); - #pragma unroll + FA_UNROLL for (int i = 0; i < DK_VEC; ++i) { q_priv[i] = CONVERT_ACC4(q_ptr[i]); } @@ -270,12 +305,12 @@ __kernel void flash_attn_f16_q1( sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset); } - ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : -INFINITY; + ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : FA_M_INIT; for (int k_idx = tid; k_idx < n_kv; k_idx += Q1_WG_SIZE) { const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; const global DATA_TYPE4* k_ptr = (const global DATA_TYPE4*)(k_base + k_row_offset); ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); - #pragma unroll + FA_UNROLL for (int k = 0; k < DK_VEC; k++) { dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc); } @@ -293,7 +328,7 @@ __kernel void flash_attn_f16_q1( __local ACC_TYPE local_m[Q1_WG_SIZE]; local_m[tid] = m_i; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]); barrier(CLK_LOCAL_MEM_FENCE); @@ -301,7 +336,7 @@ __kernel void flash_attn_f16_q1( const ACC_TYPE m_final = local_m[0]; ACC_TYPE4 o_acc[DV_VEC]; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); ACC_TYPE l_i = 0.0f; @@ -311,7 +346,7 @@ __kernel void flash_attn_f16_q1( const global DATA_TYPE4* k_ptr = (const global DATA_TYPE4*)(k_base + k_row_offset); const global DATA_TYPE4* v_ptr = (const global DATA_TYPE4*)(v_base + v_row_offset); ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); - #pragma unroll + FA_UNROLL for (int k = 0; k < DK_VEC; k++) { dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc); } @@ -325,7 +360,7 @@ __kernel void flash_attn_f16_q1( } const ACC_TYPE p = exp(score - m_final); l_i += p; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; i++) { o_acc[i] = mad(p, CONVERT_ACC4(v_ptr[i]), o_acc[i]); } @@ -335,7 +370,7 @@ __kernel void flash_attn_f16_q1( __local ACC_TYPE4 local_o_comp[Q1_WG_SIZE]; local_l[tid] = l_i; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_l[tid] += local_l[tid + s]; barrier(CLK_LOCAL_MEM_FENCE); @@ -354,7 +389,7 @@ __kernel void flash_attn_f16_q1( for (int i = 0; i < DV_VEC; i++) { local_o_comp[tid] = o_acc[i]; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_o_comp[tid] += local_o_comp[tid + s]; barrier(CLK_LOCAL_MEM_FENCE); @@ -364,7 +399,7 @@ __kernel void flash_attn_f16_q1( } } } else if (tid == 0) { - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) o_row[i] = (DATA_TYPE4)(0.0f); } } diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl index a6d747903..2547731c3 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl @@ -13,6 +13,18 @@ #define WG_SIZE (BLOCK_M) #define Q1_WG_SIZE 64 +// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs, +// infinite operand can cause undefined behavior and miscompilation for exp. +// Therefore, a large negative value is used instead. +#define FA_M_INIT (-3.0e38f) + +// Drop full unroll at DK>=192 — Adreno compiler host-memory budget. +#if DK >= 192 +#define FA_UNROLL +#else +#define FA_UNROLL _Pragma("unroll") +#endif + inline float get_alibi_slope( const float max_bias, const uint h, const uint n_head_log2, const float m0, const float m1 ) { @@ -82,18 +94,18 @@ __kernel void flash_attn_f32( if (my_query_row < n_q) { const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + my_query_row * q_nb1; const global DATA_TYPE4* q_ptr = (const global DATA_TYPE4*)(q_base + q_row_offset); - #pragma unroll + FA_UNROLL for (int i = 0; i < DK_VEC; ++i) { q_priv[i] = CONVERT_ACC4(q_ptr[i]); } } ACC_TYPE4 o_acc[DV_VEC]; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_acc[i] = (ACC_TYPE4)(0.0f); } - ACC_TYPE m_i = -INFINITY; + ACC_TYPE m_i = FA_M_INIT; ACC_TYPE l_i = 0.0f; float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); @@ -126,49 +138,72 @@ __kernel void flash_attn_f32( continue; } - for (int j = 0; j < BLOCK_N; j += 2) { + for (int j = 0; j < BLOCK_N; j += 4) { const int k_row0 = k_start + j; const int k_row1 = k_start + j + 1; + const int k_row2 = k_start + j + 2; + const int k_row3 = k_start + j + 3; ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); ACC_TYPE4 dot_acc1 = (ACC_TYPE4)(0.0f); - #pragma unroll + ACC_TYPE4 dot_acc2 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc3 = (ACC_TYPE4)(0.0f); + FA_UNROLL for (int k = 0; k < DK_VEC; k++) { - dot_acc0 = mad(q_priv[k], CONVERT_ACC4(l_k[j][k]), dot_acc0); - dot_acc1 = mad(q_priv[k], CONVERT_ACC4(l_k[j+1][k]), dot_acc1); + const ACC_TYPE4 qk = q_priv[k]; + dot_acc0 = mad(qk, CONVERT_ACC4(l_k[j][k]), dot_acc0); + dot_acc1 = mad(qk, CONVERT_ACC4(l_k[j+1][k]), dot_acc1); + dot_acc2 = mad(qk, CONVERT_ACC4(l_k[j+2][k]), dot_acc2); + dot_acc3 = mad(qk, CONVERT_ACC4(l_k[j+3][k]), dot_acc3); } - ACC_TYPE score0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; - ACC_TYPE score1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; + ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; + ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; + ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale; + ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale; if (is_causal) { - if (k_row0 > (n_kv - n_q + my_query_row)) score0 = -INFINITY; - if (k_row1 > (n_kv - n_q + my_query_row)) score1 = -INFINITY; + const int causal_limit = n_kv - n_q + my_query_row; + if (k_row0 > causal_limit) s0 = FA_M_INIT; + if (k_row1 > causal_limit) s1 = FA_M_INIT; + if (k_row2 > causal_limit) s2 = FA_M_INIT; + if (k_row3 > causal_limit) s3 = FA_M_INIT; } - - if (k_row0 >= n_kv) score0 = -INFINITY; - if (k_row1 >= n_kv) score1 = -INFINITY; + if (k_row0 >= n_kv) s0 = FA_M_INIT; + if (k_row1 >= n_kv) s1 = FA_M_INIT; + if (k_row2 >= n_kv) s2 = FA_M_INIT; + if (k_row3 >= n_kv) s3 = FA_M_INIT; if (mask_base != NULL) { const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base + my_query_row * mask_nb1); - if (k_row0 < n_kv) score0 += slope * (ACC_TYPE)mask_ptr[k_row0]; - if (k_row1 < n_kv) score1 += slope * (ACC_TYPE)mask_ptr[k_row1]; + if (k_row0 < n_kv) s0 += slope * (ACC_TYPE)mask_ptr[k_row0]; + if (k_row1 < n_kv) s1 += slope * (ACC_TYPE)mask_ptr[k_row1]; + if (k_row2 < n_kv) s2 += slope * (ACC_TYPE)mask_ptr[k_row2]; + if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3]; } if (logit_softcap > 0.0f) { - score0 = logit_softcap * tanh(score0 / logit_softcap); - score1 = logit_softcap * tanh(score1 / logit_softcap); + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); } - const ACC_TYPE m_new = max(m_i, max(score0, score1)); - const ACC_TYPE p0 = exp(score0 - m_new); - const ACC_TYPE p1 = exp(score1 - m_new); - const ACC_TYPE scale_prev = exp(m_i - m_new); + const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); + const ACC_TYPE scale_prev = native_exp(m_i - m_new); + const ACC_TYPE p0 = native_exp(s0 - m_new); + const ACC_TYPE p1 = native_exp(s1 - m_new); + const ACC_TYPE p2 = native_exp(s2 - m_new); + const ACC_TYPE p3 = native_exp(s3 - m_new); - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { - o_acc[i] = o_acc[i] * scale_prev + p0 * CONVERT_ACC4(l_v[j][i]) + p1 * CONVERT_ACC4(l_v[j+1][i]); + o_acc[i] = mad(p3, CONVERT_ACC4(l_v[j+3][i]), + mad(p2, CONVERT_ACC4(l_v[j+2][i]), + mad(p1, CONVERT_ACC4(l_v[j+1][i]), + mad(p0, CONVERT_ACC4(l_v[j][i]), + o_acc[i] * scale_prev)))); } - l_i = l_i * scale_prev + p0 + p1; + l_i = l_i * scale_prev + p0 + p1 + p2 + p3; m_i = m_new; } } @@ -180,7 +215,7 @@ __kernel void flash_attn_f32( const ACC_TYPE m_final = max(m_i, m_sink); const ACC_TYPE scale_o = exp(m_i - m_final); - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_acc[i] *= scale_o; } @@ -192,12 +227,12 @@ __kernel void flash_attn_f32( global DATA_TYPE4 *o_row = (global DATA_TYPE4 *)(o_base + o_row_offset); if (l_i > 0.0f) { const ACC_TYPE l_inv = 1.0f / l_i; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_row[i] = CONVERT_DATA4(o_acc[i] * l_inv); } } else { - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_row[i] = (DATA_TYPE4)(0.0f); } @@ -259,7 +294,7 @@ __kernel void flash_attn_f32_q1( ACC_TYPE4 q_priv[DK_VEC]; const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2; const global DATA_TYPE4* q_ptr = (const global DATA_TYPE4*)(q_base + q_row_offset); - #pragma unroll + FA_UNROLL for (int i = 0; i < DK_VEC; ++i) { q_priv[i] = CONVERT_ACC4(q_ptr[i]); } @@ -271,12 +306,12 @@ __kernel void flash_attn_f32_q1( sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset); } - ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : -INFINITY; + ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : FA_M_INIT; for (int k_idx = tid; k_idx < n_kv; k_idx += Q1_WG_SIZE) { const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; const global DATA_TYPE4* k_ptr = (const global DATA_TYPE4*)(k_base + k_row_offset); ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); - #pragma unroll + FA_UNROLL for (int k = 0; k < DK_VEC; k++) { dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc); } @@ -294,7 +329,7 @@ __kernel void flash_attn_f32_q1( __local ACC_TYPE local_m[Q1_WG_SIZE]; local_m[tid] = m_i; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]); barrier(CLK_LOCAL_MEM_FENCE); @@ -302,7 +337,7 @@ __kernel void flash_attn_f32_q1( const ACC_TYPE m_final = local_m[0]; ACC_TYPE4 o_acc[DV_VEC]; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); ACC_TYPE l_i = 0.0f; @@ -312,7 +347,7 @@ __kernel void flash_attn_f32_q1( const global DATA_TYPE4* k_ptr = (const global DATA_TYPE4*)(k_base + k_row_offset); const global DATA_TYPE4* v_ptr = (const global DATA_TYPE4*)(v_base + v_row_offset); ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); - #pragma unroll + FA_UNROLL for (int k = 0; k < DK_VEC; k++) { dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc); } @@ -326,7 +361,7 @@ __kernel void flash_attn_f32_q1( } const ACC_TYPE p = exp(score - m_final); l_i += p; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; i++) { o_acc[i] = mad(p, CONVERT_ACC4(v_ptr[i]), o_acc[i]); } @@ -336,7 +371,7 @@ __kernel void flash_attn_f32_q1( __local ACC_TYPE4 local_o_comp[Q1_WG_SIZE]; local_l[tid] = l_i; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_l[tid] += local_l[tid + s]; barrier(CLK_LOCAL_MEM_FENCE); @@ -355,7 +390,7 @@ __kernel void flash_attn_f32_q1( for (int i = 0; i < DV_VEC; i++) { local_o_comp[tid] = o_acc[i]; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_o_comp[tid] += local_o_comp[tid + s]; barrier(CLK_LOCAL_MEM_FENCE); @@ -365,7 +400,7 @@ __kernel void flash_attn_f32_q1( } } } else if (tid == 0) { - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) o_row[i] = (DATA_TYPE4)(0.0f); } } diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl index ec7361b9e..a7f1de325 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl @@ -1,5 +1,13 @@ #pragma OPENCL EXTENSION cl_khr_fp16 : enable +#ifdef cl_khr_subgroup_shuffle +#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable +#define HAS_SUBGROUP_SHUFFLE 1 +#elif defined(cl_qcom_subgroup_shuffle) +#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable +#define HAS_SUBGROUP_SHUFFLE 1 +#endif + #define ACC_TYPE float #define ACC_TYPE4 float4 #define Q_DATA_TYPE4 float4 @@ -12,9 +20,34 @@ #define DK_VEC (DK/4) #define DV_VEC (DV/4) -#define WG_SIZE (BLOCK_M) #define Q1_WG_SIZE 64 +// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs, +// infinite operand can cause undefined behavior and miscompilation for exp. +// Therefore, a large negative value is used instead. +#define FA_M_INIT (-3.0e38f) + +// Drop full unroll at DK>=192 — Adreno compiler host-memory budget. +#if DK >= 192 +#define FA_UNROLL +#else +#define FA_UNROLL _Pragma("unroll") +#endif + +// N_SPLIT>1 splits DK/DV across threads to cut per-thread register use. +#ifndef N_SPLIT +#define N_SPLIT 1 +#endif + +#define SPLIT_DK_VEC (DK_VEC / N_SPLIT) +#define SPLIT_DV_VEC (DV_VEC / N_SPLIT) + +#if N_SPLIT > 1 +#define WG_SIZE (BLOCK_M * N_SPLIT) +#else +#define WG_SIZE (BLOCK_M) +#endif + inline float get_alibi_slope( const float max_bias, const uint h, const uint n_head_log2, const float m0, const float m1 ) { @@ -54,19 +87,38 @@ __kernel void flash_attn_f32_f16( const int mask_ne2, const int mask_ne3, const global void* sinks_void, - const ulong sinks_offset + const ulong sinks_offset, + const global void * k_pad_void, + const global void * v_pad_void, + const global void * mask_pad_void, + const global char * blk, + const int n_kv_blocks, + const ulong mask_pad_nb1, + const ulong mask_pad_nb2, + const ulong mask_pad_nb3 ) { const int tid = get_local_id(0); const int block_q_idx = get_group_id(0); const int head_batch_idx = get_global_id(1); - const int my_query_row = block_q_idx * BLOCK_M + tid; +#if N_SPLIT > 1 + const int q_lane = tid / N_SPLIT; + const int split_idx = tid % N_SPLIT; +#else + const int q_lane = tid; + const int split_idx = 0; +#endif + + const int my_query_row = block_q_idx * BLOCK_M + q_lane; + const int query_valid = my_query_row < n_q; const int batch_idx = head_batch_idx / n_head; const int head_idx = head_batch_idx % n_head; const int gqa_ratio = n_head / n_head_kv; const int head_kv_idx = head_idx / gqa_ratio; + const int mask_head_idx = mask_void != NULL ? head_idx % mask_ne2 : 0; + const int mask_batch_idx = mask_void != NULL ? batch_idx % mask_ne3 : 0; const global char* q_base = (const global char*)q_void + q_offset; const global char* k_base = (const global char*)k_void + k_offset; @@ -75,27 +127,41 @@ __kernel void flash_attn_f32_f16( const global char* mask_base = NULL; if (mask_void != NULL) { - const int mask_head_idx = head_idx % mask_ne2; - const int mask_batch_idx = batch_idx % mask_ne3; mask_base = (const global char*)mask_void + mask_offset + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2; } + const global char* mask_pad_base = NULL; + if (mask_pad_void != NULL) { + mask_pad_base = (const global char*)mask_pad_void + mask_batch_idx * mask_pad_nb3 + mask_head_idx * mask_pad_nb2; + } + const global char* blk_base = NULL; + if (blk != NULL) { + const int n_q_blocks = (n_q + BLOCK_M - 1) / BLOCK_M; + blk_base = blk + (((mask_batch_idx * mask_ne2) + mask_head_idx) * n_q_blocks + block_q_idx) * n_kv_blocks; + } - ACC_TYPE4 q_priv[DK_VEC]; - if (my_query_row < n_q) { + ACC_TYPE4 q_priv[SPLIT_DK_VEC]; + const int dk_off = split_idx * SPLIT_DK_VEC; + if (query_valid) { const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + my_query_row * q_nb1; const global Q_DATA_TYPE4* q_ptr = (const global Q_DATA_TYPE4*)(q_base + q_row_offset); - #pragma unroll - for (int i = 0; i < DK_VEC; ++i) { - q_priv[i] = CONVERT_Q_ACC4(q_ptr[i]); + FA_UNROLL + for (int i = 0; i < SPLIT_DK_VEC; ++i) { + q_priv[i] = CONVERT_Q_ACC4(q_ptr[dk_off + i]); + } + } else { + FA_UNROLL + for (int i = 0; i < SPLIT_DK_VEC; ++i) { + q_priv[i] = (ACC_TYPE4)(0.0f); } } - ACC_TYPE4 o_acc[DV_VEC]; - #pragma unroll - for (int i = 0; i < DV_VEC; ++i) { + ACC_TYPE4 o_acc[SPLIT_DV_VEC]; + FA_UNROLL + for (int i = 0; i < SPLIT_DV_VEC; ++i) { o_acc[i] = (ACC_TYPE4)(0.0f); } - ACC_TYPE m_i = -INFINITY; + + ACC_TYPE m_i = FA_M_INIT; ACC_TYPE l_i = 0.0f; float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); @@ -103,86 +169,369 @@ __kernel void flash_attn_f32_f16( __local KV_DATA_TYPE4 l_k[BLOCK_N][DK_VEC]; __local KV_DATA_TYPE4 l_v[BLOCK_N][DV_VEC]; +#if N_SPLIT > 1 && !defined(HAS_SUBGROUP_SHUFFLE) + __local ACC_TYPE local_partial[BLOCK_N][WG_SIZE]; + __local ACC_TYPE local_p[BLOCK_M][BLOCK_N]; + __local ACC_TYPE local_softmax_scale[BLOCK_M]; + __local ACC_TYPE local_l_inv[BLOCK_M]; +#endif + for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) { + char blk_cur = 1; + if (blk_base != NULL) { + blk_cur = blk_base[k_start / BLOCK_N]; + if (blk_cur == 0) continue; + } + + const int use_kv_pad = k_pad_void != NULL && k_start + BLOCK_N > n_kv; + const int k_tile_start = use_kv_pad ? 0 : k_start; + const ulong k_tile_nb2 = use_kv_pad ? (ulong) BLOCK_N * k_nb1 : k_nb2; + const ulong k_tile_nb3 = use_kv_pad ? (ulong) n_head_kv * k_tile_nb2 : k_nb3; + const ulong v_tile_nb2 = use_kv_pad ? (ulong) BLOCK_N * v_nb1 : v_nb2; + const ulong v_tile_nb3 = use_kv_pad ? (ulong) n_head_kv * v_tile_nb2 : v_nb3; + const global char* k_tile_base = use_kv_pad ? (const global char*) k_pad_void : k_base; + const global char* v_tile_base = use_kv_pad ? (const global char*) v_pad_void : v_base; + for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) { const int row = i / DK_VEC; const int col = i % DK_VEC; - const int k_row_idx = k_start + row; - if (k_row_idx < n_kv) { - const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_row_idx * k_nb1; - l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_base + k_row_offset))[col]; + const int k_row_idx = k_tile_start + row; + if (use_kv_pad || k_row_idx < n_kv) { + const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1; + l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; + } else { + l_k[row][col] = (KV_DATA_TYPE4)(0.0h); } } for (int i = tid; i < BLOCK_N * DV_VEC; i += WG_SIZE) { const int row = i / DV_VEC; const int col = i % DV_VEC; - const int v_row_idx = k_start + row; - if (v_row_idx < n_kv) { - const ulong v_row_offset = batch_idx * v_nb3 + head_kv_idx * v_nb2 + v_row_idx * v_nb1; - l_v[row][col] = ((__global KV_DATA_TYPE4*)(v_base + v_row_offset))[col]; + const int v_row_idx = k_tile_start + row; + if (use_kv_pad || v_row_idx < n_kv) { + const ulong v_row_offset = batch_idx * v_tile_nb3 + head_kv_idx * v_tile_nb2 + v_row_idx * v_nb1; + l_v[row][col] = ((__global KV_DATA_TYPE4*)(v_tile_base + v_row_offset))[col]; + } else { + l_v[row][col] = (KV_DATA_TYPE4)(0.0h); } } barrier(CLK_LOCAL_MEM_FENCE); - if (my_query_row >= n_q) { - continue; +#if N_SPLIT > 1 && defined(HAS_SUBGROUP_SHUFFLE) + { + const int dv_off = split_idx * SPLIT_DV_VEC; + for (int j = 0; j < BLOCK_N; j += 2) { + const int k_row0 = k_start + j; + const int k_row1 = k_start + j + 1; + + ACC_TYPE partial0 = 0.0f; + ACC_TYPE partial1 = 0.0f; + FA_UNROLL + for (int k = 0; k < SPLIT_DK_VEC; k++) { + const ACC_TYPE4 qk = q_priv[k]; + ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(l_k[j ][dk_off + k]); + ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(l_k[j+1][dk_off + k]); + partial0 += dot0.s0 + dot0.s1 + dot0.s2 + dot0.s3; + partial1 += dot1.s0 + dot1.s1 + dot1.s2 + dot1.s3; + } + + FA_UNROLL + for (int step = 1; step < N_SPLIT; step <<= 1) { + partial0 += sub_group_shuffle_xor(partial0, step); + partial1 += sub_group_shuffle_xor(partial1, step); + } + + ACC_TYPE score0 = partial0 * scale; + ACC_TYPE score1 = partial1 * scale; + + if (!query_valid) { score0 = FA_M_INIT; score1 = FA_M_INIT; } + if (is_causal) { + if (k_row0 > (n_kv - n_q + my_query_row)) score0 = FA_M_INIT; + if (k_row1 > (n_kv - n_q + my_query_row)) score1 = FA_M_INIT; + } + if (k_row0 >= n_kv) score0 = FA_M_INIT; + if (k_row1 >= n_kv) score1 = FA_M_INIT; + + if (query_valid && mask_base != NULL && blk_cur != 2) { + if (use_kv_pad && mask_pad_base != NULL) { + const global MASK_DATA_TYPE* mask_ptr = + (const global MASK_DATA_TYPE*)(mask_pad_base + my_query_row * mask_pad_nb1); + score0 += slope * (ACC_TYPE)mask_ptr[j]; + score1 += slope * (ACC_TYPE)mask_ptr[j + 1]; + } else { + const global MASK_DATA_TYPE* mask_ptr = + (const global MASK_DATA_TYPE*)(mask_base + my_query_row * mask_nb1); + if (k_row0 < n_kv) score0 += slope * (ACC_TYPE)mask_ptr[k_row0]; + if (k_row1 < n_kv) score1 += slope * (ACC_TYPE)mask_ptr[k_row1]; + } + } + + if (logit_softcap > 0.0f) { + score0 = logit_softcap * tanh(score0 / logit_softcap); + score1 = logit_softcap * tanh(score1 / logit_softcap); + } + + const ACC_TYPE m_new = max(m_i, max(score0, score1)); + // Whole tile masked (m_new == FA_M_INIT): force the exp() args + // far negative so the tile contributes 0, not exp(0)=1. + const ACC_TYPE m_exp = (m_new == FA_M_INIT) ? 0.0f : m_new; + const ACC_TYPE sp = native_exp(m_i - m_exp); + const ACC_TYPE p0 = native_exp(score0 - m_exp); + const ACC_TYPE p1 = native_exp(score1 - m_exp); + + FA_UNROLL + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + o_acc[i] = o_acc[i] * sp + + p0 * CONVERT_KV_ACC4(l_v[j ][dv_off + i]) + + p1 * CONVERT_KV_ACC4(l_v[j+1][dv_off + i]); + } + l_i = l_i * sp + p0 + p1; + m_i = m_new; + } } - - for (int j = 0; j < BLOCK_N; j += 2) { - const int k_row0 = k_start + j; - const int k_row1 = k_start + j + 1; - - ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); - ACC_TYPE4 dot_acc1 = (ACC_TYPE4)(0.0f); - #pragma unroll - for (int k = 0; k < DK_VEC; k++) { - dot_acc0 = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j][k]), dot_acc0); - dot_acc1 = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j+1][k]), dot_acc1); +#elif N_SPLIT > 1 + // N_SPLIT>1 fallback (no shuffle): 3-phase local-memory reduction. + // Phase 1 — partial dots for all BLOCK_N tokens. + for (int j = 0; j < BLOCK_N; ++j) { + ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); + FA_UNROLL + for (int k = 0; k < SPLIT_DK_VEC; k++) { + dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j][dk_off + k]), dot_acc); } - ACC_TYPE score0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; - ACC_TYPE score1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; - - if (is_causal) { - if (k_row0 > (n_kv - n_q + my_query_row)) score0 = -INFINITY; - if (k_row1 > (n_kv - n_q + my_query_row)) score1 = -INFINITY; - } - - if (k_row0 >= n_kv) score0 = -INFINITY; - if (k_row1 >= n_kv) score1 = -INFINITY; - - if (mask_base != NULL) { - const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base + my_query_row * mask_nb1); - if (k_row0 < n_kv) score0 += slope * (ACC_TYPE)mask_ptr[k_row0]; - if (k_row1 < n_kv) score1 += slope * (ACC_TYPE)mask_ptr[k_row1]; - } - - if (logit_softcap > 0.0f) { - score0 = logit_softcap * tanh(score0 / logit_softcap); - score1 = logit_softcap * tanh(score1 / logit_softcap); - } - - const ACC_TYPE m_new = max(m_i, max(score0, score1)); - const ACC_TYPE p0 = exp(score0 - m_new); - const ACC_TYPE p1 = exp(score1 - m_new); - const ACC_TYPE scale_prev = exp(m_i - m_new); - - #pragma unroll - for (int i = 0; i < DV_VEC; ++i) { - o_acc[i] = o_acc[i] * scale_prev + p0 * CONVERT_KV_ACC4(l_v[j][i]) + p1 * CONVERT_KV_ACC4(l_v[j+1][i]); - } - l_i = l_i * scale_prev + p0 + p1; - m_i = m_new; + local_partial[j][tid] = + dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3; } + barrier(CLK_LOCAL_MEM_FENCE); // 1 barrier: partial dots visible + + // Phase 2 — split_idx==0 reduces partial sums and computes block softmax. + if (split_idx == 0) { + if (query_valid) { + ACC_TYPE m_new = m_i; + for (int j = 0; j < BLOCK_N; ++j) { + const int k_row = k_start + j; + ACC_TYPE score = 0.0f; + FA_UNROLL + for (int s = 0; s < N_SPLIT; s++) { + score += local_partial[j][q_lane * N_SPLIT + s]; + } + score *= scale; + + if (is_causal && k_row > (n_kv - n_q + my_query_row)) score = FA_M_INIT; + if (k_row >= n_kv) score = FA_M_INIT; + + if (mask_base != NULL && blk_cur != 2) { + if (use_kv_pad && mask_pad_base != NULL) { + const global MASK_DATA_TYPE* mask_ptr = + (const global MASK_DATA_TYPE*)(mask_pad_base + my_query_row * mask_pad_nb1); + score += slope * (ACC_TYPE)mask_ptr[j]; + } else { + const global MASK_DATA_TYPE* mask_ptr = + (const global MASK_DATA_TYPE*)(mask_base + my_query_row * mask_nb1); + if (k_row < n_kv) score += slope * (ACC_TYPE)mask_ptr[k_row]; + } + } + + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } + + m_new = max(m_new, score); + local_p[q_lane][j] = score; + } + + const ACC_TYPE m_exp = (m_new == FA_M_INIT) ? 0.0f : m_new; + const ACC_TYPE sp = native_exp(m_i - m_exp); + ACC_TYPE l_new = l_i * sp; + for (int j = 0; j < BLOCK_N; ++j) { + const ACC_TYPE p = native_exp(local_p[q_lane][j] - m_exp); + local_p[q_lane][j] = p; + l_new += p; + } + local_softmax_scale[q_lane] = sp; + l_i = l_new; + m_i = m_new; + } else { + local_softmax_scale[q_lane] = 1.0f; + for (int j = 0; j < BLOCK_N; ++j) local_p[q_lane][j] = 0.0f; + } + } + barrier(CLK_LOCAL_MEM_FENCE); + + // Phase 3 — V accumulate using broadcast probabilities. + { + const ACC_TYPE sp_block = local_softmax_scale[q_lane]; + const int dv_off = split_idx * SPLIT_DV_VEC; + FA_UNROLL + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + o_acc[i] *= sp_block; + } + for (int j = 0; j < BLOCK_N; ++j) { + const ACC_TYPE p = local_p[q_lane][j]; + FA_UNROLL + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + o_acc[i] = mad(p, CONVERT_KV_ACC4(l_v[j][dv_off + i]), o_acc[i]); + } + } + } +#else + // N_SPLIT==1: j+=4 unroll. Requires BLOCK_N % 4 == 0. + if (query_valid) { + for (int j = 0; j < BLOCK_N; j += 4) { + const int k_row0 = k_start + j; + const int k_row1 = k_start + j + 1; + const int k_row2 = k_start + j + 2; + const int k_row3 = k_start + j + 3; + + ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc1 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc2 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc3 = (ACC_TYPE4)(0.0f); + FA_UNROLL + for (int k = 0; k < DK_VEC; k++) { + const ACC_TYPE4 qk = q_priv[k]; + dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j][k]), dot_acc0); + dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k]), dot_acc1); + dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k]), dot_acc2); + dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k]), dot_acc3); + } + ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; + ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; + ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale; + ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale; + + if (is_causal) { + const int causal_limit = n_kv - n_q + my_query_row; + if (k_row0 > causal_limit) s0 = FA_M_INIT; + if (k_row1 > causal_limit) s1 = FA_M_INIT; + if (k_row2 > causal_limit) s2 = FA_M_INIT; + if (k_row3 > causal_limit) s3 = FA_M_INIT; + } + if (k_row0 >= n_kv) s0 = FA_M_INIT; + if (k_row1 >= n_kv) s1 = FA_M_INIT; + if (k_row2 >= n_kv) s2 = FA_M_INIT; + if (k_row3 >= n_kv) s3 = FA_M_INIT; + + if (mask_base != NULL && blk_cur != 2) { + if (use_kv_pad && mask_pad_base != NULL) { + const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_pad_base + my_query_row * mask_pad_nb1); + s0 += slope * (ACC_TYPE)mask_ptr[j]; + s1 += slope * (ACC_TYPE)mask_ptr[j + 1]; + s2 += slope * (ACC_TYPE)mask_ptr[j + 2]; + s3 += slope * (ACC_TYPE)mask_ptr[j + 3]; + } else { + const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base + my_query_row * mask_nb1); + if (k_row0 < n_kv) s0 += slope * (ACC_TYPE)mask_ptr[k_row0]; + if (k_row1 < n_kv) s1 += slope * (ACC_TYPE)mask_ptr[k_row1]; + if (k_row2 < n_kv) s2 += slope * (ACC_TYPE)mask_ptr[k_row2]; + if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3]; + } + } + + if (logit_softcap > 0.0f) { + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); + } + + const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); + // Whole tile masked (m_new == FA_M_INIT): force the exp() args + // far negative so the tile contributes 0, not exp(0)=1. + const ACC_TYPE m_exp = (m_new == FA_M_INIT) ? 0.0f : m_new; + const ACC_TYPE scale_prev = native_exp(m_i - m_exp); + const ACC_TYPE p0 = native_exp(s0 - m_exp); + const ACC_TYPE p1 = native_exp(s1 - m_exp); + const ACC_TYPE p2 = native_exp(s2 - m_exp); + const ACC_TYPE p3 = native_exp(s3 - m_exp); + + FA_UNROLL + for (int i = 0; i < DV_VEC; ++i) { + o_acc[i] = mad(p3, CONVERT_KV_ACC4(l_v[j+3][i]), + mad(p2, CONVERT_KV_ACC4(l_v[j+2][i]), + mad(p1, CONVERT_KV_ACC4(l_v[j+1][i]), + mad(p0, CONVERT_KV_ACC4(l_v[j][i]), + o_acc[i] * scale_prev)))); + } + l_i = l_i * scale_prev + p0 + p1 + p2 + p3; + m_i = m_new; + } + } +#endif + // End of tile: every thread must finish reading l_k/l_v before the + // next iteration's load overwrites them (WAR hazard on local memory). + barrier(CLK_LOCAL_MEM_FENCE); } - if (my_query_row < n_q) { + // Write output. +#if N_SPLIT > 1 && defined(HAS_SUBGROUP_SHUFFLE) + if (query_valid) { + ACC_TYPE sinks_sp = 1.0f; + if (sinks_void != NULL) { + const global ACC_TYPE* sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset); + const ACC_TYPE m_sink = sinks_ptr[head_idx]; + const ACC_TYPE m_final = max(m_i, m_sink); + sinks_sp = exp(m_i - m_final); + l_i = l_i * sinks_sp + exp(m_sink - m_final); + m_i = m_final; + } + const ACC_TYPE l_inv = (l_i > 0.0f) ? (1.0f / l_i) : 0.0f; + const int dv_off = split_idx * SPLIT_DV_VEC; + const ulong o_row_offset = batch_idx * o_nb3 + my_query_row * o_nb2 + head_idx * o_nb1; + global O_DATA_TYPE4 *o_row = (global O_DATA_TYPE4 *)(o_base + o_row_offset); + if (l_inv > 0.0f) { + FA_UNROLL + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + o_row[dv_off + i] = CONVERT_O_DATA4(o_acc[i] * sinks_sp * l_inv); + } + } else { + FA_UNROLL + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + o_row[dv_off + i] = (O_DATA_TYPE4)(0.0f); + } + } + } +#elif N_SPLIT > 1 + if (split_idx == 0) { + ACC_TYPE sinks_sp = 1.0f; + if (query_valid && sinks_void != NULL) { + const global ACC_TYPE* sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset); + const ACC_TYPE m_sink = sinks_ptr[head_idx]; + const ACC_TYPE m_final = max(m_i, m_sink); + sinks_sp = exp(m_i - m_final); + l_i = l_i * sinks_sp + exp(m_sink - m_final); + m_i = m_final; + } + local_softmax_scale[q_lane] = sinks_sp; + local_l_inv[q_lane] = (query_valid && l_i > 0.0f) ? (1.0f / l_i) : 0.0f; + } + barrier(CLK_LOCAL_MEM_FENCE); + + if (query_valid) { + const ACC_TYPE sinks_sp = local_softmax_scale[q_lane]; + const ACC_TYPE l_inv = local_l_inv[q_lane]; + const int dv_off = split_idx * SPLIT_DV_VEC; + const ulong o_row_offset = batch_idx * o_nb3 + my_query_row * o_nb2 + head_idx * o_nb1; + global O_DATA_TYPE4 *o_row = (global O_DATA_TYPE4 *)(o_base + o_row_offset); + if (l_inv > 0.0f) { + FA_UNROLL + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + o_row[dv_off + i] = CONVERT_O_DATA4(o_acc[i] * sinks_sp * l_inv); + } + } else { + FA_UNROLL + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + o_row[dv_off + i] = (O_DATA_TYPE4)(0.0f); + } + } + } +#else + if (query_valid) { if (sinks_void != NULL) { const global ACC_TYPE* sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset); const ACC_TYPE m_sink = sinks_ptr[head_idx]; const ACC_TYPE m_final = max(m_i, m_sink); const ACC_TYPE scale_o = exp(m_i - m_final); - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_acc[i] *= scale_o; } @@ -194,17 +543,18 @@ __kernel void flash_attn_f32_f16( global O_DATA_TYPE4 *o_row = (global O_DATA_TYPE4 *)(o_base + o_row_offset); if (l_i > 0.0f) { const ACC_TYPE l_inv = 1.0f / l_i; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_row[i] = CONVERT_O_DATA4(o_acc[i] * l_inv); } } else { - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) { o_row[i] = (O_DATA_TYPE4)(0.0f); } } } +#endif } __kernel void flash_attn_f32_f16_q1( @@ -258,13 +608,16 @@ __kernel void flash_attn_f32_f16_q1( mask_base = (const global char*)mask_void + mask_offset + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2; } - ACC_TYPE4 q_priv[DK_VEC]; + // Q is uniform across WG threads (n_q=1). Share via local memory to + // avoid per-thread q_priv[DK_VEC] dynamic-indexed private array that + // spills to DDR on Adreno. + __local ACC_TYPE4 q_shared[DK_VEC]; const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2; const global Q_DATA_TYPE4* q_ptr = (const global Q_DATA_TYPE4*)(q_base + q_row_offset); - #pragma unroll - for (int i = 0; i < DK_VEC; ++i) { - q_priv[i] = CONVERT_Q_ACC4(q_ptr[i]); + for (int i = tid; i < DK_VEC; i += Q1_WG_SIZE) { + q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]); } + barrier(CLK_LOCAL_MEM_FENCE); float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); @@ -273,14 +626,14 @@ __kernel void flash_attn_f32_f16_q1( sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset); } - ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : -INFINITY; + ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : FA_M_INIT; for (int k_idx = tid; k_idx < n_kv; k_idx += Q1_WG_SIZE) { const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; const global KV_DATA_TYPE4* k_ptr = (const global KV_DATA_TYPE4*)(k_base + k_row_offset); ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); - #pragma unroll + FA_UNROLL for (int k = 0; k < DK_VEC; k++) { - dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); + dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; if (mask_base != NULL) { @@ -296,7 +649,7 @@ __kernel void flash_attn_f32_f16_q1( __local ACC_TYPE local_m[Q1_WG_SIZE]; local_m[tid] = m_i; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]); barrier(CLK_LOCAL_MEM_FENCE); @@ -304,7 +657,7 @@ __kernel void flash_attn_f32_f16_q1( const ACC_TYPE m_final = local_m[0]; ACC_TYPE4 o_acc[DV_VEC]; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); ACC_TYPE l_i = 0.0f; @@ -314,9 +667,9 @@ __kernel void flash_attn_f32_f16_q1( const global KV_DATA_TYPE4* k_ptr = (const global KV_DATA_TYPE4*)(k_base + k_row_offset); const global KV_DATA_TYPE4* v_ptr = (const global KV_DATA_TYPE4*)(v_base + v_row_offset); ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); - #pragma unroll + FA_UNROLL for (int k = 0; k < DK_VEC; k++) { - dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); + dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; if (mask_base != NULL) { @@ -328,7 +681,7 @@ __kernel void flash_attn_f32_f16_q1( } const ACC_TYPE p = exp(score - m_final); l_i += p; - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; i++) { o_acc[i] = mad(p, CONVERT_KV_ACC4(v_ptr[i]), o_acc[i]); } @@ -338,7 +691,7 @@ __kernel void flash_attn_f32_f16_q1( __local ACC_TYPE4 local_o_comp[Q1_WG_SIZE]; local_l[tid] = l_i; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_l[tid] += local_l[tid + s]; barrier(CLK_LOCAL_MEM_FENCE); @@ -357,7 +710,7 @@ __kernel void flash_attn_f32_f16_q1( for (int i = 0; i < DV_VEC; i++) { local_o_comp[tid] = o_acc[i]; barrier(CLK_LOCAL_MEM_FENCE); - #pragma unroll + FA_UNROLL for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { if (tid < s) local_o_comp[tid] += local_o_comp[tid + s]; barrier(CLK_LOCAL_MEM_FENCE); @@ -367,7 +720,257 @@ __kernel void flash_attn_f32_f16_q1( } } } else if (tid == 0) { - #pragma unroll + FA_UNROLL for (int i = 0; i < DV_VEC; ++i) o_row[i] = (O_DATA_TYPE4)(0.0f); } } + +// Flash-decoding split pass. gid(2) = q_idx * n_splits + split_idx. +// Partial record per split: [m, l, O[DV]]. Merge kernel applies sink + norm. +#define FA_PARTIAL_FLOATS (2 + DV) + +__kernel void flash_attn_f32_f16_q1_split( + const global void * q_void, ulong q_offset, + const global void * k_void, ulong k_offset, + const global void * v_void, ulong v_offset, + const float scale, + const int n_q, + const int n_kv, + const int n_head, + const ulong q_nb1, const ulong q_nb2, const ulong q_nb3, + const ulong k_nb1, const ulong k_nb2, const ulong k_nb3, + const ulong v_nb1, const ulong v_nb2, const ulong v_nb3, + const float max_bias, + const float m0, + const float m1, + const int n_head_log2, + const float logit_softcap, + const int n_head_kv, + const global void * mask_void, + const ulong mask_offset, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3, + global float * partial_void, + const int n_splits, + const int kv_per_split +) { + const int tid = get_local_id(0); + const int head_batch_idx = get_global_id(1); + const int split_q_idx = get_global_id(2); + const int split_idx = split_q_idx % n_splits; + const int q_idx = split_q_idx / n_splits; + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + const int gqa_ratio = n_head / n_head_kv; + const int head_kv_idx = head_idx / gqa_ratio; + + const int kv_start = split_idx * kv_per_split; + const int kv_end = min(kv_start + kv_per_split, n_kv); + + const ulong record_stride = (ulong) FA_PARTIAL_FLOATS; + const ulong record_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx) + * n_splits + split_idx); + global float * rec = partial_void + record_idx * record_stride; + global float4 * rec_o = (global float4 *) (rec + 2); + + if (kv_start >= kv_end) { + // Empty split: leave sentinel partial for merge. + if (tid == 0) { + rec[0] = FA_M_INIT; + rec[1] = 0.0f; + } + return; + } + + const global char * q_base = (const global char *) q_void + q_offset; + const global char * k_base = (const global char *) k_void + k_offset; + const global char * v_base = (const global char *) v_void + v_offset; + + const global char * mask_base = NULL; + if (mask_void != NULL) { + const int mask_head_idx = head_idx % mask_ne2; + const int mask_batch_idx = batch_idx % mask_ne3; + mask_base = (const global char *) mask_void + mask_offset + + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2 + + (ulong) q_idx * mask_nb1; + } + + // Share Q via local memory (n_q=1 per split -> uniform across WG). + __local ACC_TYPE4 q_shared[DK_VEC]; + const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1; + const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset); + for (int i = tid; i < DK_VEC; i += Q1_WG_SIZE) { + q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]); + } + barrier(CLK_LOCAL_MEM_FENCE); + + const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); + + // Pass 1a — split-local max. + ACC_TYPE m_i = FA_M_INIT; + for (int k_idx = kv_start + tid; k_idx < kv_end; k_idx += Q1_WG_SIZE) { + const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; + const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_offset); + ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); + #pragma unroll + for (int k = 0; k < DK_VEC; ++k) { + dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); + } + ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (mask_base != NULL) { + const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base); + score += slope * (ACC_TYPE) mask_ptr[k_idx]; + } + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } + m_i = max(m_i, score); + } + + __local ACC_TYPE local_m[Q1_WG_SIZE]; + local_m[tid] = m_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]); + barrier(CLK_LOCAL_MEM_FENCE); + } + const ACC_TYPE m_c = local_m[0]; + + // Pass 1b — softmax-weighted V accumulate. + ACC_TYPE4 o_acc[DV_VEC]; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); + ACC_TYPE l_i = 0.0f; + + for (int k_idx = kv_start + tid; k_idx < kv_end; k_idx += Q1_WG_SIZE) { + const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; + const ulong v_row_offset = batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1; + const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_offset); + const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_offset); + ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); + #pragma unroll + for (int k = 0; k < DK_VEC; ++k) { + dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); + } + ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (mask_base != NULL) { + const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base); + score += slope * (ACC_TYPE) mask_ptr[k_idx]; + } + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } + const ACC_TYPE p = exp(score - m_c); + l_i += p; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) { + o_acc[i] = mad(p, CONVERT_KV_ACC4(v_ptr[i]), o_acc[i]); + } + } + + __local ACC_TYPE local_l[Q1_WG_SIZE]; + __local ACC_TYPE4 local_o[Q1_WG_SIZE]; + local_l[tid] = l_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_l[tid] += local_l[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + const ACC_TYPE l_c = local_l[0]; + + if (tid == 0) { + rec[0] = (float) m_c; + rec[1] = (float) l_c; + } + for (int i = 0; i < DV_VEC; ++i) { + local_o[tid] = o_acc[i]; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_o[tid] += local_o[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + if (tid == 0) { + rec_o[i] = local_o[0]; + } + } +} + +// FD Pass 2: merge per-split partials into final O. Empty splits drop via exp(-INF)=0. +__kernel void flash_attn_f32_merge( + const global float * partial_void, + global void * o_void, + const ulong o_offset, + const int n_head, + const int n_splits, + const ulong o_nb1, const ulong o_nb2, const ulong o_nb3, + const global void * sinks_void, + const ulong sinks_offset, + const int n_q +) { + const int lane = get_local_id(0); // 0..DV_VEC-1 + const int head_batch_idx = get_global_id(1); + const int q_idx = get_global_id(2); + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + + const ulong record_stride = (ulong) FA_PARTIAL_FLOATS; + const ulong record_idx_0 = (((ulong) batch_idx * n_head + head_idx) * n_q + q_idx) * n_splits; + const global float * rec0 = partial_void + record_idx_0 * record_stride; + + __local ACC_TYPE m_final_shared; + __local ACC_TYPE l_final_shared; + if (lane == 0) { + ACC_TYPE m = FA_M_INIT; + for (int c = 0; c < n_splits; ++c) { + const ACC_TYPE m_c = rec0[c * record_stride + 0]; + m = max(m, m_c); + } + ACC_TYPE m_sink = 0.0f; + bool has_sink = false; + if (sinks_void != NULL) { + const global ACC_TYPE * sinks_ptr = + (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset); + m_sink = sinks_ptr[head_idx]; + has_sink = true; + m = max(m, m_sink); + } + ACC_TYPE l = 0.0f; + for (int c = 0; c < n_splits; ++c) { + const ACC_TYPE m_c = rec0[c * record_stride + 0]; + const ACC_TYPE l_c = rec0[c * record_stride + 1]; + if (m_c > FA_M_INIT) { + l += l_c * exp(m_c - m); + } + } + if (has_sink) { + l += exp(m_sink - m); + } + m_final_shared = m; + l_final_shared = l; + } + barrier(CLK_LOCAL_MEM_FENCE); + const ACC_TYPE m_final = m_final_shared; + const ACC_TYPE l_final = l_final_shared; + const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f; + + ACC_TYPE4 o = (ACC_TYPE4)(0.0f); + for (int c = 0; c < n_splits; ++c) { + const global float * rec_c = rec0 + c * record_stride; + const ACC_TYPE m_c = rec_c[0]; + if (m_c <= FA_M_INIT) continue; + const global float4 * rec_oc = (const global float4 *) (rec_c + 2); + const ACC_TYPE scale_c = exp(m_c - m_final); + o = mad((ACC_TYPE4)(scale_c), rec_oc[lane], o); + } + o = o * l_inv; + + const ulong o_row_offset = (ulong) batch_idx * o_nb3 + (ulong) q_idx * o_nb2 + (ulong) head_idx * o_nb1; + global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) ((global char *) o_void + o_offset + o_row_offset); + o_row[lane] = CONVERT_O_DATA4(o); +} diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl new file mode 100644 index 000000000..36167ba54 --- /dev/null +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl @@ -0,0 +1,1041 @@ +#pragma OPENCL EXTENSION cl_khr_fp16 : enable +#ifdef cl_khr_integer_dot_product +#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable +#define FA_HAVE_INT_DOT 1 +#endif + +#ifdef cl_khr_subgroup_shuffle +#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable +#define HAS_SUBGROUP_SHUFFLE 1 +#elif defined(cl_qcom_subgroup_shuffle) +#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable +#define HAS_SUBGROUP_SHUFFLE 1 +#endif + +// Flash attention: Q=f32, K=q4_0, V=q4_0. +// Block = half d + uchar qs[16]; qs[j] low/high nibble -> elem j / j+16. +// Dequant: val[i] = d * (nibble_i - 8). dp4a path runs on raw 0..15 nibbles +// and applies the -8*sum(q) correction once per block (needs Q q_sum). + +#define ACC_TYPE float +#define ACC_TYPE4 float4 +#define Q_DATA_TYPE4 float4 +#define O_DATA_TYPE4 float4 +#define MASK_DATA_TYPE half +#define CONVERT_Q_ACC4(x) (x) +#define CONVERT_O_DATA4(x) (x) + +#define DK_VEC (DK/4) +#define DV_VEC (DV/4) +#define Q1_WG_SIZE 64 + +// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs, +// infinite operand can cause undefined behavior and miscompilation for exp. +// Therefore, a large negative value is used instead. +#define FA_M_INIT (-3.0e38f) + +#define QK4_0 32 +#define Q4_0_BLOCK_SIZE 18 + +#define DK_Q4_BLOCKS (DK / QK4_0) +#define DV_Q4_BLOCKS (DV / QK4_0) + +inline float dot_q4_0_f32(const global char * block_ptr, ACC_TYPE4 * q_slice) { + float d = vload_half(0, (const global half *)block_ptr); + const global uchar * qs = (const global uchar *)(block_ptr + 2); + + float sum = 0.0f; + // Low nibbles -> elems 0..15. + #pragma unroll + for (int g = 0; g < 4; ++g) { + float4 nv = (float4)((float)(int)(qs[g*4 + 0] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 1] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 2] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 3] & 0x0F) - 8.0f); + sum += dot(q_slice[g], nv); + } + // High nibbles -> elems 16..31. + #pragma unroll + for (int g = 0; g < 4; ++g) { + float4 nv = (float4)((float)(int)(qs[g*4 + 0] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 1] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 2] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 3] >> 4) - 8.0f); + sum += dot(q_slice[4 + g], nv); + } + return sum * d; +} + +#ifdef FA_HAVE_INT_DOT +inline uint pack_i8x4(char a, char b, char c, char d) { + return ((uint)(uchar)a) | + ((uint)(uchar)b) << 8 | + ((uint)(uchar)c) << 16 | + ((uint)(uchar)d) << 24; +} + +// Returns (qd, q_sum); q_sum feeds the -8*sum(q) bias correction. +typedef struct { + float qd; + int q_sum; +} q4_q_block_info; + +inline q4_q_block_info quant_q_block_int8_packed_q4(const ACC_TYPE4 * q_block, + uint * out_packed) { + float amax = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) { + float4 av = fabs(q_block[i]); + amax = fmax(amax, fmax(fmax(av.s0, av.s1), fmax(av.s2, av.s3))); + } + float qd = amax / 127.0f; + float qid = (amax > 0.0f) ? 127.0f / amax : 0.0f; + + int q_sum = 0; + #pragma unroll + for (int i = 0; i < 8; ++i) { + float4 v = q_block[i] * qid; + char a = (char)((int)round(v.s0)); + char b = (char)((int)round(v.s1)); + char c = (char)((int)round(v.s2)); + char d = (char)((int)round(v.s3)); + out_packed[i] = pack_i8x4(a, b, c, d); + q_sum += (int)a + (int)b + (int)c + (int)d; + } + q4_q_block_info info = { qd, q_sum }; + return info; +} + +// k_packed[0..3] = low nibbles (Q elems 0..15), k_packed[4..7] = high (16..31). +inline void pack_q4_0_nibbles(const global uchar * qs, uint * k_packed) { + #pragma unroll + for (int g = 0; g < 4; ++g) { + uchar b0 = qs[g*4 + 0]; + uchar b1 = qs[g*4 + 1]; + uchar b2 = qs[g*4 + 2]; + uchar b3 = qs[g*4 + 3]; + k_packed[g] = + ((uint)(b0 & 0x0F)) | + ((uint)(b1 & 0x0F)) << 8 | + ((uint)(b2 & 0x0F)) << 16 | + ((uint)(b3 & 0x0F)) << 24; + k_packed[4 + g] = + ((uint)(b0 >> 4)) | + ((uint)(b1 >> 4)) << 8 | + ((uint)(b2 >> 4)) << 16 | + ((uint)(b3 >> 4)) << 24; + } +} + +inline float dot_q4_0_int(const global char * k_block_ptr, + const uint * q_packed, + float q_d, + int q_sum) { + float kd = vload_half(0, (const global half *)k_block_ptr); + const global uchar * k_qs = (const global uchar *)(k_block_ptr + 2); + + uint k_packed[8]; + pack_q4_0_nibbles(k_qs, k_packed); + + int sum = 0; + #pragma unroll + for (int i = 0; i < 8; ++i) { + sum = dot_acc_sat_4x8packed_ss_int(q_packed[i], k_packed[i], sum); + } + // Correct raw-nibble sum: (nibble - 8) bias -> subtract 8 * q_sum. + return (float)(sum - 8 * q_sum) * q_d * kd; +} +#endif // FA_HAVE_INT_DOT + +inline void dequant_q4_0_f32(const global char * block_ptr, ACC_TYPE4 * out) { + float d = vload_half(0, (const global half *)block_ptr); + const global uchar * qs = (const global uchar *)(block_ptr + 2); + + #pragma unroll + for (int g = 0; g < 4; ++g) { + out[g] = d * (float4)((float)(int)(qs[g*4 + 0] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 1] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 2] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 3] & 0x0F) - 8.0f); + } + #pragma unroll + for (int g = 0; g < 4; ++g) { + out[4 + g] = d * (float4)((float)(int)(qs[g*4 + 0] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 1] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 2] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 3] >> 4) - 8.0f); + } +} + +// max_bias<=0 returns 1.0 so score += 1.0 * mask[k] stays a no-op multiplier. +inline float get_alibi_slope(float max_bias, int head_idx, int n_head_log2, float m0, float m1) { + if (max_bias <= 0.0f) return 1.0f; + float base = (head_idx < n_head_log2) ? m0 : m1; + int exph = (head_idx < n_head_log2) ? (head_idx + 1) : (2*(head_idx - n_head_log2) + 1); + return pow(base, (float)exph); +} + +// q1 decode: one query row per WG, threads sweep KV positions. +__kernel void flash_attn_f32_q4_0_q1( + const global void * q_void, ulong q_offset, + const global void * k_void, ulong k_offset, + const global void * v_void, ulong v_offset, + global void * o_void, ulong o_offset, + const float scale, + const int n_q, + const int n_kv, + const int is_causal, + const int n_head, + const ulong q_nb1, const ulong q_nb2, const ulong q_nb3, + const ulong k_nb1, const ulong k_nb2, const ulong k_nb3, + const ulong v_nb1, const ulong v_nb2, const ulong v_nb3, + const ulong o_nb1, const ulong o_nb2, const ulong o_nb3, + const float max_bias, + const float m0, + const float m1, + const int n_head_log2, + const float logit_softcap, + const int n_head_kv, + const global void* mask_void, + const ulong mask_offset, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3, + const global void* sinks_void, + const ulong sinks_offset +) { + const int tid = get_local_id(0); + const int head_batch_idx = get_global_id(1); + + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + + const int gqa_ratio = n_head / n_head_kv; + const int head_kv_idx = head_idx / gqa_ratio; + + const global char* q_base = (const global char*)q_void + q_offset; + const global char* k_base = (const global char*)k_void + k_offset; + const global char* v_base = (const global char*)v_void + v_offset; + global char* o_base = (global char*)o_void + o_offset; + + const global char* mask_base = NULL; + if (mask_void != NULL) { + const int mask_head_idx = head_idx % mask_ne2; + const int mask_batch_idx = batch_idx % mask_ne3; + mask_base = (const global char*)mask_void + mask_offset + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2; + } + + ACC_TYPE4 q_priv[DK_VEC]; + const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2; + const global Q_DATA_TYPE4* q_ptr = (const global Q_DATA_TYPE4*)(q_base + q_row_offset); + #pragma unroll + for (int i = 0; i < DK_VEC; ++i) { + q_priv[i] = CONVERT_Q_ACC4(q_ptr[i]); + } + +#ifdef FA_HAVE_INT_DOT + // Quantise Q once per thread: 8 uints + qd + q_sum per block. + uint q_packed[DK_Q4_BLOCKS * 8]; + float q_d_scale[DK_Q4_BLOCKS]; + int q_sum_arr[DK_Q4_BLOCKS]; + #pragma unroll + for (int b = 0; b < DK_Q4_BLOCKS; ++b) { + q4_q_block_info info = quant_q_block_int8_packed_q4(&q_priv[b * 8], &q_packed[b * 8]); + q_d_scale[b] = info.qd; + q_sum_arr[b] = info.q_sum; + } +#endif + + float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); + + const global ACC_TYPE* sinks_ptr = NULL; + if (sinks_void != NULL) { + sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset); + } + + // One-pass online softmax (FA-2): single sweep over kv positions, + // updating per-thread (m_i, l_i, o_acc) per K. Eliminates the second + // K read of the original two-pass implementation. + ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : FA_M_INIT; + ACC_TYPE l_i = 0.0f; + ACC_TYPE4 o_acc[DV_VEC]; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); + + for (int k_idx = tid; k_idx < n_kv; k_idx += Q1_WG_SIZE) { + const global char* k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; + const global char* v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1; + + ACC_TYPE score = 0.0f; + #pragma unroll + for (int b = 0; b < DK_Q4_BLOCKS; b++) { +#ifdef FA_HAVE_INT_DOT + score += dot_q4_0_int(k_row + b * Q4_0_BLOCK_SIZE, + &q_packed[b * 8], q_d_scale[b], q_sum_arr[b]); +#else + score += dot_q4_0_f32(k_row + b * Q4_0_BLOCK_SIZE, &q_priv[b * 8]); +#endif + } + score *= scale; + + if (mask_base != NULL) { + const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base); + score += slope * (ACC_TYPE)mask_ptr[k_idx]; + } + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } + + // Online softmax step. + const ACC_TYPE m_new = max(m_i, score); + const ACC_TYPE alpha = exp(m_i - m_new); + const ACC_TYPE p = exp(score - m_new); + + l_i = alpha * l_i + p; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] *= alpha; + + #pragma unroll + for (int b = 0; b < DV_Q4_BLOCKS; b++) { + ACC_TYPE4 v_dequant[8]; + dequant_q4_0_f32(v_row + b * Q4_0_BLOCK_SIZE, v_dequant); + #pragma unroll + for (int i = 0; i < 8; i++) { + o_acc[b * 8 + i] = mad(p, v_dequant[i], o_acc[b * 8 + i]); + } + } + + m_i = m_new; + } + + // Cross-thread reduce: max(m_i) -> m_final, rescale per-thread l_i and + // o_acc by alpha = exp(m_i_thread - m_final) before sum-reduce. + __local ACC_TYPE local_m[Q1_WG_SIZE]; + local_m[tid] = m_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]); + barrier(CLK_LOCAL_MEM_FENCE); + } + const ACC_TYPE m_final = local_m[0]; + + const ACC_TYPE alpha_final = exp(m_i - m_final); + l_i *= alpha_final; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] *= alpha_final; + + __local ACC_TYPE local_l[Q1_WG_SIZE]; + __local ACC_TYPE4 local_o_comp[Q1_WG_SIZE]; + local_l[tid] = l_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_l[tid] += local_l[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + + const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1; + global O_DATA_TYPE4 *o_row = (global O_DATA_TYPE4 *)(o_base + o_row_offset); + ACC_TYPE l_final = local_l[0]; + + if (sinks_ptr != NULL) { + l_final += exp(sinks_ptr[head_idx] - m_final); + } + + if (l_final > 0.0f) { + const ACC_TYPE l_inv = 1.0f / l_final; + for (int i = 0; i < DV_VEC; i++) { + local_o_comp[tid] = o_acc[i]; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_o_comp[tid] += local_o_comp[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + if (tid == 0) { + o_row[i] = CONVERT_O_DATA4(local_o_comp[0] * l_inv); + } + } + } else if (tid == 0) { + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_row[i] = (O_DATA_TYPE4)(0.0f); + } +} + +// Flash-decoding split pass for q4_0 KV. Merge kernel is type-agnostic and +// shared with the f16/q8_0 FA kernels. +#define FA_PARTIAL_FLOATS (2 + DV) + +__kernel void flash_attn_f32_q4_0_q1_split( + const global void * q_void, ulong q_offset, + const global void * k_void, ulong k_offset, + const global void * v_void, ulong v_offset, + const float scale, + const int n_q, + const int n_kv, + const int n_head, + const ulong q_nb1, const ulong q_nb2, const ulong q_nb3, + const ulong k_nb1, const ulong k_nb2, const ulong k_nb3, + const ulong v_nb1, const ulong v_nb2, const ulong v_nb3, + const float max_bias, + const float m0, + const float m1, + const int n_head_log2, + const float logit_softcap, + const int n_head_kv, + const global void * mask_void, + const ulong mask_offset, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3, + global float * partial_void, + const int n_splits, + const int kv_per_split +) { + const int tid = get_local_id(0); + const int head_batch_idx = get_global_id(1); + const int split_q_idx = get_global_id(2); + const int split_idx = split_q_idx % n_splits; + const int q_idx = split_q_idx / n_splits; + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + const int gqa_ratio = n_head / n_head_kv; + const int head_kv_idx = head_idx / gqa_ratio; + + const int kv_start = split_idx * kv_per_split; + const int kv_end = min(kv_start + kv_per_split, n_kv); + + const ulong record_stride = (ulong) FA_PARTIAL_FLOATS; + const ulong record_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx) + * n_splits + split_idx); + global float * rec = partial_void + record_idx * record_stride; + global float4 * rec_o = (global float4 *) (rec + 2); + + if (kv_start >= kv_end) { + if (tid == 0) { + rec[0] = FA_M_INIT; + rec[1] = 0.0f; + } + return; + } + + const global char * q_base = (const global char *) q_void + q_offset; + const global char * k_base = (const global char *) k_void + k_offset; + const global char * v_base = (const global char *) v_void + v_offset; + + const global char * mask_base = NULL; + if (mask_void != NULL) { + const int mask_head_idx = head_idx % mask_ne2; + const int mask_batch_idx = batch_idx % mask_ne3; + mask_base = (const global char *) mask_void + mask_offset + + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2 + + (ulong) q_idx * mask_nb1; + } + + ACC_TYPE4 q_priv[DK_VEC]; + const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1; + const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset); + #pragma unroll + for (int i = 0; i < DK_VEC; ++i) { + q_priv[i] = CONVERT_Q_ACC4(q_ptr[i]); + } + +#ifdef FA_HAVE_INT_DOT + uint q_packed[DK_Q4_BLOCKS * 8]; + float q_d_scale[DK_Q4_BLOCKS]; + int q_sum_arr[DK_Q4_BLOCKS]; + #pragma unroll + for (int b = 0; b < DK_Q4_BLOCKS; ++b) { + q4_q_block_info info = quant_q_block_int8_packed_q4(&q_priv[b * 8], &q_packed[b * 8]); + q_d_scale[b] = info.qd; + q_sum_arr[b] = info.q_sum; + } +#endif + + const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); + + // One-pass online softmax (FA-2): single sweep over the split's K range. + ACC_TYPE m_i = FA_M_INIT; + ACC_TYPE l_i = 0.0f; + ACC_TYPE4 o_acc[DV_VEC]; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); + + for (int k_idx = kv_start + tid; k_idx < kv_end; k_idx += Q1_WG_SIZE) { + const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; + const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1; + ACC_TYPE score = 0.0f; + #pragma unroll + for (int b = 0; b < DK_Q4_BLOCKS; ++b) { +#ifdef FA_HAVE_INT_DOT + score += dot_q4_0_int(k_row + b * Q4_0_BLOCK_SIZE, + &q_packed[b * 8], q_d_scale[b], q_sum_arr[b]); +#else + score += dot_q4_0_f32(k_row + b * Q4_0_BLOCK_SIZE, &q_priv[b * 8]); +#endif + } + score *= scale; + if (mask_base != NULL) { + const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base); + score += slope * (ACC_TYPE) mask_ptr[k_idx]; + } + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } + + // Online softmax step. + const ACC_TYPE m_new = max(m_i, score); + const ACC_TYPE alpha = exp(m_i - m_new); + const ACC_TYPE p = exp(score - m_new); + + l_i = alpha * l_i + p; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] *= alpha; + + #pragma unroll + for (int b = 0; b < DV_Q4_BLOCKS; ++b) { + ACC_TYPE4 v_dequant[8]; + dequant_q4_0_f32(v_row + b * Q4_0_BLOCK_SIZE, v_dequant); + #pragma unroll + for (int i = 0; i < 8; ++i) { + o_acc[b * 8 + i] = mad(p, v_dequant[i], o_acc[b * 8 + i]); + } + } + + m_i = m_new; + } + + // Cross-thread reduce: max(m_i) -> m_c, rescale per-thread l_i and o_acc + // by alpha = exp(m_i_thread - m_c) before sum-reduce. + __local ACC_TYPE local_m[Q1_WG_SIZE]; + local_m[tid] = m_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]); + barrier(CLK_LOCAL_MEM_FENCE); + } + const ACC_TYPE m_c = local_m[0]; + + const ACC_TYPE alpha_final = exp(m_i - m_c); + l_i *= alpha_final; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] *= alpha_final; + + __local ACC_TYPE local_l[Q1_WG_SIZE]; + __local ACC_TYPE4 local_o[Q1_WG_SIZE]; + local_l[tid] = l_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_l[tid] += local_l[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + const ACC_TYPE l_c = local_l[0]; + + if (tid == 0) { + rec[0] = (float) m_c; + rec[1] = (float) l_c; + } + for (int i = 0; i < DV_VEC; ++i) { + local_o[tid] = o_acc[i]; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_o[tid] += local_o[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + if (tid == 0) { + rec_o[i] = local_o[0]; + } + } +} + +// Prefill: q4_0 K/V, n_q > 1. BLOCK_M × BLOCK_N tiling. +// K in local as packed nibbles + per-block scale; V dequant -> half in local. +// Requires DK % QK4_0 == 0 and DV % QK4_0 == 0. +#define KV_DATA_TYPE4 half4 +#define CONVERT_KV_ACC4(x) convert_float4(x) + +#define DK_Q4_BLOCKS_PREFILL (DK / QK4_0) +#define DV_Q4_BLOCKS_PREFILL (DV / QK4_0) + +// N_SPLIT>1 splits DK/DV across N_SPLIT threads per query row; needs +// sub_group_shuffle_xor and DK_Q4_BLOCKS_PREFILL % N_SPLIT == 0. +#ifndef N_SPLIT +#define N_SPLIT 1 +#endif + +#if N_SPLIT > 1 +#define SPLIT_DK_VEC (DK_VEC / N_SPLIT) +#define SPLIT_DV_VEC (DV_VEC / N_SPLIT) +#define SPLIT_DK_Q4_BLOCKS (DK_Q4_BLOCKS_PREFILL / N_SPLIT) +#define WG_SIZE (BLOCK_M * N_SPLIT) +#else +#define SPLIT_DK_VEC DK_VEC +#define SPLIT_DV_VEC DV_VEC +#define SPLIT_DK_Q4_BLOCKS DK_Q4_BLOCKS_PREFILL +#define WG_SIZE BLOCK_M +#endif + +__kernel void flash_attn_f32_q4_0( + const global void * q_void, ulong q_offset, + const global void * k_void, ulong k_offset, + const global void * v_void, ulong v_offset, + global void * o_void, ulong o_offset, + const float scale, + const int n_q, + const int n_kv, + const int is_causal, + const int n_head, + const ulong q_nb1, const ulong q_nb2, const ulong q_nb3, + const ulong k_nb1, const ulong k_nb2, const ulong k_nb3, + const ulong v_nb1, const ulong v_nb2, const ulong v_nb3, + const ulong o_nb1, const ulong o_nb2, const ulong o_nb3, + const float max_bias, + const float m0, + const float m1, + const int n_head_log2, + const float logit_softcap, + const int n_head_kv, + const global void* mask_void, + const ulong mask_offset, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3, + const global void* sinks_void, + const ulong sinks_offset, + // blk: per-(qblock,kvblock) class from flash_attn_blk_f16 + // (0=masked, 1=mixed, 2=unmasked). NULL disables the prepass opt. + const global void * blk_void +) { + const int tid = get_local_id(0); + const int block_q_idx = get_group_id(0); + const int head_batch_idx = get_global_id(1); + +#if N_SPLIT > 1 + const int q_lane = tid / N_SPLIT; + const int split_idx = tid % N_SPLIT; +#else + const int q_lane = tid; + const int split_idx = 0; +#endif + const int my_query_row = block_q_idx * BLOCK_M + q_lane; + const int query_valid = my_query_row < n_q; + + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + + const int gqa_ratio = n_head / n_head_kv; + const int head_kv_idx = head_idx / gqa_ratio; + const int mask_head_idx = mask_void != NULL ? head_idx % mask_ne2 : 0; + const int mask_batch_idx = mask_void != NULL ? batch_idx % mask_ne3 : 0; + + const global char * q_base = (const global char *) q_void + q_offset; + const global char * k_base = (const global char *) k_void + k_offset; + const global char * v_base = (const global char *) v_void + v_offset; + global char * o_base = (global char *) o_void + o_offset; + + const global char * mask_base = NULL; + if (mask_void != NULL) { + mask_base = (const global char *) mask_void + mask_offset + + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2; + } + + // BLK_PREPASS_BM may differ from this kernel's BLOCK_M; scale q-block idx. + #ifndef BLK_PREPASS_BM + #define BLK_PREPASS_BM BLOCK_M + #endif + const global char * blk_base = NULL; + int n_kv_blocks = 0; + if (blk_void != NULL) { + n_kv_blocks = (n_kv + BLOCK_N - 1) / BLOCK_N; + const int n_q_blocks_prepass = (n_q + BLK_PREPASS_BM - 1) / BLK_PREPASS_BM; + const int prepass_q_block = (block_q_idx * BLOCK_M) / BLK_PREPASS_BM; + blk_base = (const global char *) blk_void + + (((mask_batch_idx * mask_ne2) + mask_head_idx) * n_q_blocks_prepass + prepass_q_block) * n_kv_blocks; + } + + const int dk_off_vec = split_idx * SPLIT_DK_VEC; + ACC_TYPE4 q_priv[SPLIT_DK_VEC]; + if (query_valid) { + const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + my_query_row * q_nb1; + const global float4 * q_ptr = (const global float4 *) (q_base + q_row_offset); + #pragma unroll + for (int i = 0; i < SPLIT_DK_VEC; ++i) { + q_priv[i] = q_ptr[dk_off_vec + i]; + } + } else { + #pragma unroll + for (int i = 0; i < SPLIT_DK_VEC; ++i) q_priv[i] = (ACC_TYPE4)(0.0f); + } + +#ifdef FA_HAVE_INT_DOT + uint q_packed_pf[SPLIT_DK_Q4_BLOCKS * 8]; + float q_d_pf[SPLIT_DK_Q4_BLOCKS]; + int q_sum_pf[SPLIT_DK_Q4_BLOCKS]; + #pragma unroll + for (int b = 0; b < SPLIT_DK_Q4_BLOCKS; ++b) { + q4_q_block_info info = quant_q_block_int8_packed_q4(&q_priv[b * 8], &q_packed_pf[b * 8]); + q_d_pf[b] = info.qd; + q_sum_pf[b] = info.q_sum; + } +#endif + + const int dv_off_vec = split_idx * SPLIT_DV_VEC; + ACC_TYPE4 o_acc[SPLIT_DV_VEC]; + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); + + ACC_TYPE m_i = FA_M_INIT; + ACC_TYPE l_i = 0.0f; + + float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); + +#ifdef FA_HAVE_INT_DOT + __local uint l_k_packed[BLOCK_N][DK_Q4_BLOCKS_PREFILL * 8]; + __local float l_k_scale [BLOCK_N][DK_Q4_BLOCKS_PREFILL]; +#else + __local half4 l_k[BLOCK_N][DK_VEC]; +#endif + + __local half4 l_v[BLOCK_N][DV_VEC]; + + for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) { + // Skip fully-masked KV tiles (uniform branch across WG). + char blk_cur = 1; + if (blk_base != NULL) { + blk_cur = blk_base[k_start / BLOCK_N]; + if (blk_cur == 0) continue; + } + + { +#ifdef FA_HAVE_INT_DOT + const int k_blocks_per_row = DK_Q4_BLOCKS_PREFILL; + const int n_blocks_total = BLOCK_N * k_blocks_per_row; + for (int i = tid; i < n_blocks_total; i += WG_SIZE) { + const int row = i / k_blocks_per_row; + const int blk = i % k_blocks_per_row; + const int k_row_idx = k_start + row; + if (k_row_idx < n_kv) { + const ulong k_row_off = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_row_idx * k_nb1; + const global char * blk_ptr = k_base + k_row_off + blk * Q4_0_BLOCK_SIZE; + const float df = (float) vload_half(0, (const global half *) blk_ptr); + const global uchar * qs = (const global uchar *)(blk_ptr + 2); + l_k_scale[row][blk] = df; + uint k_packed[8]; + pack_q4_0_nibbles(qs, k_packed); + #pragma unroll + for (int j = 0; j < 8; ++j) { + l_k_packed[row][blk * 8 + j] = k_packed[j]; + } + } else { + l_k_scale[row][blk] = 0.0f; + #pragma unroll + for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u; + } + } +#else + // Fallback: dequant q4_0 -> half in local memory. + const int k_blocks_per_row = DK_Q4_BLOCKS_PREFILL; + const int n_blocks_total = BLOCK_N * k_blocks_per_row; + for (int i = tid; i < n_blocks_total; i += WG_SIZE) { + const int row = i / k_blocks_per_row; + const int blk = i % k_blocks_per_row; + const int k_row_idx = k_start + row; + if (k_row_idx < n_kv) { + const ulong k_row_off = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_row_idx * k_nb1; + const global char * blk_ptr = k_base + k_row_off + blk * Q4_0_BLOCK_SIZE; + const float df = (float) vload_half(0, (const global half *) blk_ptr); + const global uchar * qs = (const global uchar *)(blk_ptr + 2); + #pragma unroll + for (int g = 0; g < 4; ++g) { + float4 vlo = df * (float4)((float)(int)(qs[g*4 + 0] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 1] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 2] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 3] & 0x0F) - 8.0f); + float4 vhi = df * (float4)((float)(int)(qs[g*4 + 0] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 1] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 2] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 3] >> 4) - 8.0f); + l_k[row][blk * 8 + g ] = (half4)((half)vlo.s0, (half)vlo.s1, (half)vlo.s2, (half)vlo.s3); + l_k[row][blk * 8 + 4 + g] = (half4)((half)vhi.s0, (half)vhi.s1, (half)vhi.s2, (half)vhi.s3); + } + } else { + #pragma unroll + for (int j = 0; j < 8; ++j) l_k[row][blk * 8 + j] = (half4)(0.0h); + } + } +#endif + } + // V tile load — dequant V -> half in local memory. + { + const int v_blocks_per_row = DV_Q4_BLOCKS_PREFILL; + const int n_blocks_total = BLOCK_N * v_blocks_per_row; + for (int i = tid; i < n_blocks_total; i += WG_SIZE) { + const int row = i / v_blocks_per_row; + const int blk = i % v_blocks_per_row; + const int v_row_idx = k_start + row; + if (v_row_idx < n_kv) { + const ulong v_row_off = batch_idx * v_nb3 + head_kv_idx * v_nb2 + v_row_idx * v_nb1; + const global char * blk_ptr = v_base + v_row_off + blk * Q4_0_BLOCK_SIZE; + const float df = (float) vload_half(0, (const global half *) blk_ptr); + const global uchar * qs = (const global uchar *)(blk_ptr + 2); + #pragma unroll + for (int g = 0; g < 4; ++g) { + float4 vlo = df * (float4)((float)(int)(qs[g*4 + 0] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 1] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 2] & 0x0F) - 8.0f, + (float)(int)(qs[g*4 + 3] & 0x0F) - 8.0f); + float4 vhi = df * (float4)((float)(int)(qs[g*4 + 0] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 1] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 2] >> 4) - 8.0f, + (float)(int)(qs[g*4 + 3] >> 4) - 8.0f); + l_v[row][blk * 8 + g ] = (half4)((half)vlo.s0, (half)vlo.s1, (half)vlo.s2, (half)vlo.s3); + l_v[row][blk * 8 + 4 + g] = (half4)((half)vhi.s0, (half)vhi.s1, (half)vhi.s2, (half)vhi.s3); + } + } else { + #pragma unroll + for (int j = 0; j < 8; ++j) l_v[row][blk * 8 + j] = (half4)(0.0h); + } + } + } + barrier(CLK_LOCAL_MEM_FENCE); + + // QK dot + online softmax. N_SPLIT>1 reduces per-thread partials via shuffle_xor. +#if N_SPLIT > 1 + { +#else + if (query_valid) { +#endif + const int k_blk_base = split_idx * SPLIT_DK_Q4_BLOCKS; + for (int j = 0; j < BLOCK_N; j += 4) { + const int k_row0 = k_start + j; + const int k_row1 = k_start + j + 1; + const int k_row2 = k_start + j + 2; + const int k_row3 = k_start + j + 3; + + ACC_TYPE s0, s1, s2, s3; +#ifdef FA_HAVE_INT_DOT + s0 = 0.0f; s1 = 0.0f; s2 = 0.0f; s3 = 0.0f; + #pragma unroll + for (int b_local = 0; b_local < SPLIT_DK_Q4_BLOCKS; ++b_local) { + const int b = k_blk_base + b_local; + int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0; + #pragma unroll + for (int g = 0; g < 8; ++g) { + const uint qp = q_packed_pf[b_local * 8 + g]; + sum0 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j ][b * 8 + g], sum0); + sum1 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+1][b * 8 + g], sum1); + sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2); + sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3); + } + const float qd = q_d_pf[b_local]; + const int q_sum = q_sum_pf[b_local]; + s0 += (float)(sum0 - 8 * q_sum) * qd * l_k_scale[j ][b]; + s1 += (float)(sum1 - 8 * q_sum) * qd * l_k_scale[j+1][b]; + s2 += (float)(sum2 - 8 * q_sum) * qd * l_k_scale[j+2][b]; + s3 += (float)(sum3 - 8 * q_sum) * qd * l_k_scale[j+3][b]; + } +#else + ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc1 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc2 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc3 = (ACC_TYPE4)(0.0f); + #pragma unroll + for (int k = 0; k < SPLIT_DK_VEC; ++k) { + const ACC_TYPE4 qk = q_priv[k]; + const int k_abs = dk_off_vec + k; + dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j ][k_abs]), dot_acc0); + dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k_abs]), dot_acc1); + dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k_abs]), dot_acc2); + dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k_abs]), dot_acc3); + } + s0 = dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3; + s1 = dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3; + s2 = dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3; + s3 = dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3; +#endif + +#if N_SPLIT > 1 + // Power-of-2 N_SPLIT: shuffle_xor butterfly. N_SPLIT=3 (DK=96): + // explicit 3-lane shuffle. + #if (N_SPLIT & (N_SPLIT - 1)) == 0 + #pragma unroll + for (int step = 1; step < N_SPLIT; step <<= 1) { + s0 += sub_group_shuffle_xor(s0, step); + s1 += sub_group_shuffle_xor(s1, step); + s2 += sub_group_shuffle_xor(s2, step); + s3 += sub_group_shuffle_xor(s3, step); + } + #else + const uint tri_base = (get_sub_group_local_id() / N_SPLIT) * N_SPLIT; + s0 = sub_group_shuffle(s0, tri_base + 0) + sub_group_shuffle(s0, tri_base + 1) + sub_group_shuffle(s0, tri_base + 2); + s1 = sub_group_shuffle(s1, tri_base + 0) + sub_group_shuffle(s1, tri_base + 1) + sub_group_shuffle(s1, tri_base + 2); + s2 = sub_group_shuffle(s2, tri_base + 0) + sub_group_shuffle(s2, tri_base + 1) + sub_group_shuffle(s2, tri_base + 2); + s3 = sub_group_shuffle(s3, tri_base + 0) + sub_group_shuffle(s3, tri_base + 1) + sub_group_shuffle(s3, tri_base + 2); + #endif + if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; } +#endif + s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale; + + if (is_causal) { + const int causal_limit = n_kv - n_q + my_query_row; + if (k_row0 > causal_limit) s0 = FA_M_INIT; + if (k_row1 > causal_limit) s1 = FA_M_INIT; + if (k_row2 > causal_limit) s2 = FA_M_INIT; + if (k_row3 > causal_limit) s3 = FA_M_INIT; + } + if (k_row0 >= n_kv) s0 = FA_M_INIT; + if (k_row1 >= n_kv) s1 = FA_M_INIT; + if (k_row2 >= n_kv) s2 = FA_M_INIT; + if (k_row3 >= n_kv) s3 = FA_M_INIT; + + if (query_valid && mask_base != NULL && blk_cur != 2) { + const global MASK_DATA_TYPE * mask_ptr = + (const global MASK_DATA_TYPE *) (mask_base + my_query_row * mask_nb1); + if (k_row0 < n_kv) s0 += slope * (ACC_TYPE) mask_ptr[k_row0]; + if (k_row1 < n_kv) s1 += slope * (ACC_TYPE) mask_ptr[k_row1]; + if (k_row2 < n_kv) s2 += slope * (ACC_TYPE) mask_ptr[k_row2]; + if (k_row3 < n_kv) s3 += slope * (ACC_TYPE) mask_ptr[k_row3]; + } + if (logit_softcap > 0.0f) { + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); + } + + const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); + // Whole tile masked (m_new == FA_M_INIT): force the exp() args + // far negative so the tile contributes 0, not exp(0)=1. + const ACC_TYPE m_exp = (m_new == FA_M_INIT) ? 0.0f : m_new; + const ACC_TYPE scale_prev = native_exp(m_i - m_exp); + const ACC_TYPE p0 = native_exp(s0 - m_exp); + const ACC_TYPE p1 = native_exp(s1 - m_exp); + const ACC_TYPE p2 = native_exp(s2 - m_exp); + const ACC_TYPE p3 = native_exp(s3 - m_exp); + + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + const int i_abs = dv_off_vec + i; + o_acc[i] = mad(p3, CONVERT_KV_ACC4(l_v[j+3][i_abs]), + mad(p2, CONVERT_KV_ACC4(l_v[j+2][i_abs]), + mad(p1, CONVERT_KV_ACC4(l_v[j+1][i_abs]), + mad(p0, CONVERT_KV_ACC4(l_v[j ][i_abs]), + o_acc[i] * scale_prev)))); + } + l_i = l_i * scale_prev + p0 + p1 + p2 + p3; + m_i = m_new; + } + } + barrier(CLK_LOCAL_MEM_FENCE); + } + + // Write output. + if (query_valid) { + if (sinks_void != NULL) { + const global ACC_TYPE * sinks_ptr = + (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset); + const ACC_TYPE m_sink = sinks_ptr[head_idx]; + const ACC_TYPE m_final = max(m_i, m_sink); + const ACC_TYPE scale_o = exp(m_i - m_final); + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) o_acc[i] *= scale_o; + l_i = l_i * scale_o + exp(m_sink - m_final); + m_i = m_final; + } + const ACC_TYPE l_inv = (l_i > 0.0f) ? (1.0f / l_i) : 0.0f; + const ulong o_row_offset = batch_idx * o_nb3 + my_query_row * o_nb2 + head_idx * o_nb1; + global float4 * o_row = (global float4 *) (o_base + o_row_offset); + if (l_inv > 0.0f) { + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) o_row[dv_off_vec + i] = o_acc[i] * l_inv; + } else { + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) o_row[dv_off_vec + i] = (float4)(0.0f); + } + } +} + +// FD Pass 2: merge split partials. Identical across q4_0/q8_0/f16; each FA +// source owns a copy since kernels compile per-source-program. +__kernel void flash_attn_f32_merge( + const global float * partial_void, + global void * o_void, + const ulong o_offset, + const int n_head, + const int n_splits, + const ulong o_nb1, const ulong o_nb2, const ulong o_nb3, + const global void * sinks_void, + const ulong sinks_offset, + const int n_q +) { + const int lane = get_local_id(0); + const int head_batch_idx = get_global_id(1); + const int q_idx = get_global_id(2); + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + + const ulong record_stride = (ulong) FA_PARTIAL_FLOATS; + const ulong record_idx_0 = (((ulong) batch_idx * n_head + head_idx) * n_q + q_idx) * n_splits; + const global float * rec0 = partial_void + record_idx_0 * record_stride; + + __local ACC_TYPE m_final_shared; + __local ACC_TYPE l_final_shared; + if (lane == 0) { + ACC_TYPE m = FA_M_INIT; + for (int c = 0; c < n_splits; ++c) { + const ACC_TYPE m_c = rec0[c * record_stride + 0]; + m = max(m, m_c); + } + ACC_TYPE m_sink = 0.0f; + bool has_sink = false; + if (sinks_void != NULL) { + const global ACC_TYPE * sinks_ptr = + (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset); + m_sink = sinks_ptr[head_idx]; + has_sink = true; + m = max(m, m_sink); + } + ACC_TYPE l = 0.0f; + for (int c = 0; c < n_splits; ++c) { + const ACC_TYPE m_c = rec0[c * record_stride + 0]; + const ACC_TYPE l_c = rec0[c * record_stride + 1]; + if (m_c > FA_M_INIT) { + l += l_c * exp(m_c - m); + } + } + if (has_sink) { + l += exp(m_sink - m); + } + m_final_shared = m; + l_final_shared = l; + } + barrier(CLK_LOCAL_MEM_FENCE); + const ACC_TYPE m_final = m_final_shared; + const ACC_TYPE l_final = l_final_shared; + const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f; + + ACC_TYPE4 o = (ACC_TYPE4)(0.0f); + for (int c = 0; c < n_splits; ++c) { + const global float * rec_c = rec0 + c * record_stride; + const ACC_TYPE m_c = rec_c[0]; + if (m_c <= FA_M_INIT) continue; + const global float4 * rec_oc = (const global float4 *) (rec_c + 2); + const ACC_TYPE scale_c = exp(m_c - m_final); + o = mad((ACC_TYPE4)(scale_c), rec_oc[lane], o); + } + o = o * l_inv; + + const ulong o_row_offset = (ulong) batch_idx * o_nb3 + (ulong) q_idx * o_nb2 + (ulong) head_idx * o_nb1; + global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) ((global char *) o_void + o_offset + o_row_offset); + o_row[lane] = CONVERT_O_DATA4(o); +} diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl new file mode 100644 index 000000000..a25823f00 --- /dev/null +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl @@ -0,0 +1,1049 @@ +#pragma OPENCL EXTENSION cl_khr_fp16 : enable +#ifdef cl_khr_integer_dot_product +#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable +#define FA_HAVE_INT_DOT 1 +#endif + +#ifdef cl_khr_subgroup_shuffle +#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable +#define HAS_SUBGROUP_SHUFFLE 1 +#elif defined(cl_qcom_subgroup_shuffle) +#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable +#define HAS_SUBGROUP_SHUFFLE 1 +#endif + +// Flash attention: Q=f32, K=q8_0, V=q8_0. + +#define ACC_TYPE float +#define ACC_TYPE4 float4 +#define Q_DATA_TYPE4 float4 +#define O_DATA_TYPE4 float4 +#define MASK_DATA_TYPE half +#define CONVERT_Q_ACC4(x) (x) +#define CONVERT_O_DATA4(x) (x) + +#define DK_VEC (DK/4) +#define DV_VEC (DV/4) +#define Q1_WG_SIZE 64 + +// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs, +// infinite operand can cause undefined behavior and miscompilation for exp. +// Therefore, a large negative value is used instead. +#define FA_M_INIT (-3.0e38f) + +// q8_0 block: 2B scale (half) + 32B int8 quants. +#define QK8_0 32 +#define Q8_0_BLOCK_SIZE 34 + +#define DK_Q8_BLOCKS (DK / QK8_0) +#define DV_Q8_BLOCKS (DV / QK8_0) + +inline float dot_q8_0_f32(const global char * block_ptr, ACC_TYPE4 * q_slice) { + float d = vload_half(0, (const global half *)block_ptr); + const global char * qs = block_ptr + 2; + + float sum = 0.0f; + #pragma unroll + for (int i = 0; i < 8; i++) { + float4 qv = (float4)((float)qs[i*4], (float)qs[i*4+1], (float)qs[i*4+2], (float)qs[i*4+3]); + sum += dot(q_slice[i], qv); + } + return sum * d; +} + +#ifdef FA_HAVE_INT_DOT +inline uint pack_i8x4(char a, char b, char c, char d) { + return ((uint)(uchar)a) | + ((uint)(uchar)b) << 8 | + ((uint)(uchar)c) << 16 | + ((uint)(uchar)d) << 24; +} + +inline float quant_q_block_int8_packed(const ACC_TYPE4 * q_block, + uint * out_packed) { + float amax = 0.0f; + #pragma unroll + for (int i = 0; i < 8; ++i) { + float4 av = fabs(q_block[i]); + amax = fmax(amax, fmax(fmax(av.s0, av.s1), fmax(av.s2, av.s3))); + } + float qd = amax / 127.0f; + float qid = (amax > 0.0f) ? 127.0f / amax : 0.0f; + + #pragma unroll + for (int i = 0; i < 8; ++i) { + float4 v = q_block[i] * qid; + char a = (char)((int)round(v.s0)); + char b = (char)((int)round(v.s1)); + char c = (char)((int)round(v.s2)); + char d = (char)((int)round(v.s3)); + out_packed[i] = pack_i8x4(a, b, c, d); + } + return qd; +} + +inline float dot_q8_0_int(const global char * k_block_ptr, + const uint * q_packed, + float q_d) { + float kd = vload_half(0, (const global half *)k_block_ptr); + const global uchar * k_qs = (const global uchar *)(k_block_ptr + 2); + + // k_qs is 2-byte aligned; pack chars per iteration rather than cast to uint*. + int sum = 0; + #pragma unroll + for (int i = 0; i < 8; ++i) { + uint k_packed = + (uint)k_qs[i*4 + 0] | + ((uint)k_qs[i*4 + 1]) << 8 | + ((uint)k_qs[i*4 + 2]) << 16 | + ((uint)k_qs[i*4 + 3]) << 24; + sum = dot_acc_sat_4x8packed_ss_int(q_packed[i], k_packed, sum); + } + return (float)sum * q_d * kd; +} +#endif // FA_HAVE_INT_DOT + +inline void dequant_q8_0_f32(const global char * block_ptr, ACC_TYPE4 * out) { + float d = vload_half(0, (const global half *)block_ptr); + const global char * qs = block_ptr + 2; + + #pragma unroll + for (int i = 0; i < 8; i++) { + out[i] = d * (float4)((float)qs[i*4], (float)qs[i*4+1], (float)qs[i*4+2], (float)qs[i*4+3]); + } +} + +// max_bias<=0 returns 1.0 so score += 1.0 * mask[k] stays a no-op multiplier. +inline float get_alibi_slope(float max_bias, int head_idx, int n_head_log2, float m0, float m1) { + if (max_bias <= 0.0f) return 1.0f; + float base = (head_idx < n_head_log2) ? m0 : m1; + int exph = (head_idx < n_head_log2) ? (head_idx + 1) : (2*(head_idx - n_head_log2) + 1); + return pow(base, (float)exph); +} + +// q1 decode: one query row per WG, threads sweep KV positions. +__kernel void flash_attn_f32_q8_0_q1( + const global void * q_void, ulong q_offset, + const global void * k_void, ulong k_offset, + const global void * v_void, ulong v_offset, + global void * o_void, ulong o_offset, + const float scale, + const int n_q, + const int n_kv, + const int is_causal, + const int n_head, + const ulong q_nb1, const ulong q_nb2, const ulong q_nb3, + const ulong k_nb1, const ulong k_nb2, const ulong k_nb3, + const ulong v_nb1, const ulong v_nb2, const ulong v_nb3, + const ulong o_nb1, const ulong o_nb2, const ulong o_nb3, + const float max_bias, + const float m0, + const float m1, + const int n_head_log2, + const float logit_softcap, + const int n_head_kv, + const global void* mask_void, + const ulong mask_offset, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3, + const global void* sinks_void, + const ulong sinks_offset +) { + const int tid = get_local_id(0); + const int head_batch_idx = get_global_id(1); + + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + + const int gqa_ratio = n_head / n_head_kv; + const int head_kv_idx = head_idx / gqa_ratio; + + const global char* q_base = (const global char*)q_void + q_offset; + const global char* k_base = (const global char*)k_void + k_offset; + const global char* v_base = (const global char*)v_void + v_offset; + global char* o_base = (global char*)o_void + o_offset; + + const global char* mask_base = NULL; + if (mask_void != NULL) { + const int mask_head_idx = head_idx % mask_ne2; + const int mask_batch_idx = batch_idx % mask_ne3; + mask_base = (const global char*)mask_void + mask_offset + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2; + } + + ACC_TYPE4 q_priv[DK_VEC]; + const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2; + const global Q_DATA_TYPE4* q_ptr = (const global Q_DATA_TYPE4*)(q_base + q_row_offset); + #pragma unroll + for (int i = 0; i < DK_VEC; ++i) { + q_priv[i] = CONVERT_Q_ACC4(q_ptr[i]); + } + +#ifdef FA_HAVE_INT_DOT + // Quantise Q once per thread; q_priv stays as fp for the V accumulate. + uint q_packed[DK_Q8_BLOCKS * 8]; + float q_d_scale[DK_Q8_BLOCKS]; + #pragma unroll + for (int b = 0; b < DK_Q8_BLOCKS; ++b) { + q_d_scale[b] = quant_q_block_int8_packed(&q_priv[b * 8], &q_packed[b * 8]); + } +#endif + + float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); + + const global ACC_TYPE* sinks_ptr = NULL; + if (sinks_void != NULL) { + sinks_ptr = (const global ACC_TYPE*)((const global char*)sinks_void + sinks_offset); + } + + // One-pass online softmax: per-thread maintains running (m_i, l_i, o_acc), + // updating each as new K positions are processed. Eliminates the second + // K read of the original two-pass implementation. After the loop, threads + // are merged via the standard FA-2 cross-thread reduction (rescale each + // thread's l_i and o_acc by alpha=exp(m_i_thread - m_final), then sum). + ACC_TYPE m_i = (sinks_ptr != NULL) ? sinks_ptr[head_idx] : FA_M_INIT; + ACC_TYPE l_i = 0.0f; + ACC_TYPE4 o_acc[DV_VEC]; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); + + for (int k_idx = tid; k_idx < n_kv; k_idx += Q1_WG_SIZE) { + const global char* k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; + const global char* v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1; + + ACC_TYPE score = 0.0f; + #pragma unroll + for (int b = 0; b < DK_Q8_BLOCKS; b++) { +#ifdef FA_HAVE_INT_DOT + score += dot_q8_0_int(k_row + b * Q8_0_BLOCK_SIZE, + &q_packed[b * 8], q_d_scale[b]); +#else + score += dot_q8_0_f32(k_row + b * Q8_0_BLOCK_SIZE, &q_priv[b * 8]); +#endif + } + score *= scale; + + if (mask_base != NULL) { + const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base); + score += slope * (ACC_TYPE)mask_ptr[k_idx]; + } + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } + + // Online softmax step. + const ACC_TYPE m_new = max(m_i, score); + const ACC_TYPE alpha = exp(m_i - m_new); + const ACC_TYPE p = exp(score - m_new); + + l_i = alpha * l_i + p; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] *= alpha; + + #pragma unroll + for (int b = 0; b < DV_Q8_BLOCKS; b++) { + ACC_TYPE4 v_dequant[8]; + dequant_q8_0_f32(v_row + b * Q8_0_BLOCK_SIZE, v_dequant); + #pragma unroll + for (int i = 0; i < 8; i++) { + o_acc[b * 8 + i] = mad(p, v_dequant[i], o_acc[b * 8 + i]); + } + } + + m_i = m_new; + } + + // Cross-thread reduce: max(m_i) -> m_final, then rescale per-thread l_i + // and o_acc by alpha = exp(m_i_thread - m_final) before sum-reduce. + __local ACC_TYPE local_m[Q1_WG_SIZE]; + local_m[tid] = m_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]); + barrier(CLK_LOCAL_MEM_FENCE); + } + const ACC_TYPE m_final = local_m[0]; + + const ACC_TYPE alpha_final = exp(m_i - m_final); + l_i *= alpha_final; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] *= alpha_final; + + __local ACC_TYPE local_l[Q1_WG_SIZE]; + __local ACC_TYPE4 local_o_comp[Q1_WG_SIZE]; + local_l[tid] = l_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_l[tid] += local_l[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + + const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1; + global O_DATA_TYPE4 *o_row = (global O_DATA_TYPE4 *)(o_base + o_row_offset); + ACC_TYPE l_final = local_l[0]; + + if (sinks_ptr != NULL) { + l_final += exp(sinks_ptr[head_idx] - m_final); + } + + if (l_final > 0.0f) { + const ACC_TYPE l_inv = 1.0f / l_final; + for (int i = 0; i < DV_VEC; i++) { + local_o_comp[tid] = o_acc[i]; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_o_comp[tid] += local_o_comp[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + if (tid == 0) { + o_row[i] = CONVERT_O_DATA4(local_o_comp[0] * l_inv); + } + } + } else if (tid == 0) { + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_row[i] = (O_DATA_TYPE4)(0.0f); + } +} + +// Flash-decoding split pass for q8_0 KV. Partial record: [m, l, O[DV]]. +// Merge kernel from flash_attn_f32_f16.cl is type-agnostic and reused. +#define FA_PARTIAL_FLOATS (2 + DV) + +__kernel void flash_attn_f32_q8_0_q1_split( + const global void * q_void, ulong q_offset, + const global void * k_void, ulong k_offset, + const global void * v_void, ulong v_offset, + const float scale, + const int n_q, + const int n_kv, + const int n_head, + const ulong q_nb1, const ulong q_nb2, const ulong q_nb3, + const ulong k_nb1, const ulong k_nb2, const ulong k_nb3, + const ulong v_nb1, const ulong v_nb2, const ulong v_nb3, + const float max_bias, + const float m0, + const float m1, + const int n_head_log2, + const float logit_softcap, + const int n_head_kv, + const global void * mask_void, + const ulong mask_offset, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3, + global float * partial_void, + const int n_splits, + const int kv_per_split +) { + const int tid = get_local_id(0); + const int head_batch_idx = get_global_id(1); + const int split_q_idx = get_global_id(2); + const int split_idx = split_q_idx % n_splits; + const int q_idx = split_q_idx / n_splits; + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + const int gqa_ratio = n_head / n_head_kv; + const int head_kv_idx = head_idx / gqa_ratio; + + const int kv_start = split_idx * kv_per_split; + const int kv_end = min(kv_start + kv_per_split, n_kv); + + const ulong record_stride = (ulong) FA_PARTIAL_FLOATS; + const ulong record_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx) + * n_splits + split_idx); + global float * rec = partial_void + record_idx * record_stride; + global float4 * rec_o = (global float4 *) (rec + 2); + + if (kv_start >= kv_end) { + // Empty split: leave sentinel partial for merge. + if (tid == 0) { + rec[0] = FA_M_INIT; + rec[1] = 0.0f; + } + return; + } + + const global char * q_base = (const global char *) q_void + q_offset; + const global char * k_base = (const global char *) k_void + k_offset; + const global char * v_base = (const global char *) v_void + v_offset; + + const global char * mask_base = NULL; + if (mask_void != NULL) { + const int mask_head_idx = head_idx % mask_ne2; + const int mask_batch_idx = batch_idx % mask_ne3; + mask_base = (const global char *) mask_void + mask_offset + + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2 + + (ulong) q_idx * mask_nb1; + } + + ACC_TYPE4 q_priv[DK_VEC]; + const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1; + const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset); + #pragma unroll + for (int i = 0; i < DK_VEC; ++i) { + q_priv[i] = CONVERT_Q_ACC4(q_ptr[i]); + } + +#ifdef FA_HAVE_INT_DOT + uint q_packed[DK_Q8_BLOCKS * 8]; + float q_d_scale[DK_Q8_BLOCKS]; + #pragma unroll + for (int b = 0; b < DK_Q8_BLOCKS; ++b) { + q_d_scale[b] = quant_q_block_int8_packed(&q_priv[b * 8], &q_packed[b * 8]); + } +#endif + + const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); + + // One-pass online softmax (FA-2): single sweep over the split's K range, + // updating per-thread (m_i, l_i, o_acc) per position. Eliminates the + // second K read of the original two-pass implementation. + ACC_TYPE m_i = FA_M_INIT; + ACC_TYPE l_i = 0.0f; + ACC_TYPE4 o_acc[DV_VEC]; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); + + for (int k_idx = kv_start + tid; k_idx < kv_end; k_idx += Q1_WG_SIZE) { + const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1; + const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1; + ACC_TYPE score = 0.0f; + #pragma unroll + for (int b = 0; b < DK_Q8_BLOCKS; ++b) { +#ifdef FA_HAVE_INT_DOT + score += dot_q8_0_int(k_row + b * Q8_0_BLOCK_SIZE, &q_packed[b * 8], q_d_scale[b]); +#else + score += dot_q8_0_f32(k_row + b * Q8_0_BLOCK_SIZE, &q_priv[b * 8]); +#endif + } + score *= scale; + if (mask_base != NULL) { + const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base); + score += slope * (ACC_TYPE) mask_ptr[k_idx]; + } + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } + + // Online softmax step. + const ACC_TYPE m_new = max(m_i, score); + const ACC_TYPE alpha = exp(m_i - m_new); + const ACC_TYPE p = exp(score - m_new); + + l_i = alpha * l_i + p; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] *= alpha; + + #pragma unroll + for (int b = 0; b < DV_Q8_BLOCKS; ++b) { + ACC_TYPE4 v_dequant[8]; + dequant_q8_0_f32(v_row + b * Q8_0_BLOCK_SIZE, v_dequant); + #pragma unroll + for (int i = 0; i < 8; ++i) { + o_acc[b * 8 + i] = mad(p, v_dequant[i], o_acc[b * 8 + i]); + } + } + + m_i = m_new; + } + + // Cross-thread reduce: max(m_i) -> m_c, then rescale per-thread l_i and + // o_acc by alpha = exp(m_i_thread - m_c) before sum-reduce. + __local ACC_TYPE local_m[Q1_WG_SIZE]; + local_m[tid] = m_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]); + barrier(CLK_LOCAL_MEM_FENCE); + } + const ACC_TYPE m_c = local_m[0]; + + const ACC_TYPE alpha_final = exp(m_i - m_c); + l_i *= alpha_final; + #pragma unroll + for (int i = 0; i < DV_VEC; ++i) o_acc[i] *= alpha_final; + + __local ACC_TYPE local_l[Q1_WG_SIZE]; + __local ACC_TYPE4 local_o[Q1_WG_SIZE]; + local_l[tid] = l_i; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_l[tid] += local_l[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + const ACC_TYPE l_c = local_l[0]; + + if (tid == 0) { + rec[0] = (float) m_c; + rec[1] = (float) l_c; + } + for (int i = 0; i < DV_VEC; ++i) { + local_o[tid] = o_acc[i]; + barrier(CLK_LOCAL_MEM_FENCE); + #pragma unroll + for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) local_o[tid] += local_o[tid + s]; + barrier(CLK_LOCAL_MEM_FENCE); + } + if (tid == 0) { + rec_o[i] = local_o[0]; + } + } +} + +// Prefill: q8_0 K/V, n_q > 1. BLOCK_M × BLOCK_N tiling. +// K path keeps packed int8 in local for dp4a QK dot; V path dequant -> half in local. +// Requires DK % QK8_0 == 0 and DV % QK8_0 == 0 (gated in supports_op). +#define KV_DATA_TYPE4 half4 +#define CONVERT_KV_ACC4(x) convert_float4(x) + +#define DK_Q8_BLOCKS_PREFILL (DK / QK8_0) +#define DV_Q8_BLOCKS_PREFILL (DV / QK8_0) + +// N_SPLIT>1 splits DK/DV across N_SPLIT threads per query row; needs +// sub_group_shuffle_xor and DK_Q8_BLOCKS_PREFILL % N_SPLIT == 0. +#ifndef N_SPLIT +#define N_SPLIT 1 +#endif + +#if N_SPLIT > 1 +#define SPLIT_DK_VEC (DK_VEC / N_SPLIT) +#define SPLIT_DV_VEC (DV_VEC / N_SPLIT) +#define SPLIT_DK_Q8_BLOCKS (DK_Q8_BLOCKS_PREFILL / N_SPLIT) +#define WG_SIZE (BLOCK_M * N_SPLIT) +#else +#define SPLIT_DK_VEC DK_VEC +#define SPLIT_DV_VEC DV_VEC +#define SPLIT_DK_Q8_BLOCKS DK_Q8_BLOCKS_PREFILL +#define WG_SIZE BLOCK_M +#endif + +// FA_V_STRATEGY: 0 = dequant V to half in local (default); 2 = keep packed +// int8 in local, dequant in the accumulate loop (smaller local, slightly slower). +#ifndef FA_V_STRATEGY +#define FA_V_STRATEGY 0 +#endif + +__kernel void flash_attn_f32_q8_0( + const global void * q_void, ulong q_offset, + const global void * k_void, ulong k_offset, + const global void * v_void, ulong v_offset, + global void * o_void, ulong o_offset, + const float scale, + const int n_q, + const int n_kv, + const int is_causal, + const int n_head, + const ulong q_nb1, const ulong q_nb2, const ulong q_nb3, + const ulong k_nb1, const ulong k_nb2, const ulong k_nb3, + const ulong v_nb1, const ulong v_nb2, const ulong v_nb3, + const ulong o_nb1, const ulong o_nb2, const ulong o_nb3, + const float max_bias, + const float m0, + const float m1, + const int n_head_log2, + const float logit_softcap, + const int n_head_kv, + const global void* mask_void, + const ulong mask_offset, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3, + const global void* sinks_void, + const ulong sinks_offset, + // blk: per-(qblock,kvblock) class from flash_attn_blk_f16 + // (0=masked, 1=mixed, 2=unmasked). NULL disables the prepass opt. + const global void * blk_void +) { + const int tid = get_local_id(0); + const int block_q_idx = get_group_id(0); + const int head_batch_idx = get_global_id(1); + +#if N_SPLIT > 1 + const int q_lane = tid / N_SPLIT; + const int split_idx = tid % N_SPLIT; +#else + const int q_lane = tid; + const int split_idx = 0; +#endif + const int my_query_row = block_q_idx * BLOCK_M + q_lane; + const int query_valid = my_query_row < n_q; + + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + + const int gqa_ratio = n_head / n_head_kv; + const int head_kv_idx = head_idx / gqa_ratio; + const int mask_head_idx = mask_void != NULL ? head_idx % mask_ne2 : 0; + const int mask_batch_idx = mask_void != NULL ? batch_idx % mask_ne3 : 0; + + const global char * q_base = (const global char *) q_void + q_offset; + const global char * k_base = (const global char *) k_void + k_offset; + const global char * v_base = (const global char *) v_void + v_offset; + global char * o_base = (global char *) o_void + o_offset; + + const global char * mask_base = NULL; + if (mask_void != NULL) { + mask_base = (const global char *) mask_void + mask_offset + + mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2; + } + + // BLK_PREPASS_BM may differ from this kernel's BLOCK_M; scale q-block idx. + #ifndef BLK_PREPASS_BM + #define BLK_PREPASS_BM BLOCK_M + #endif + const global char * blk_base = NULL; + int n_kv_blocks = 0; + if (blk_void != NULL) { + n_kv_blocks = (n_kv + BLOCK_N - 1) / BLOCK_N; + const int n_q_blocks_prepass = (n_q + BLK_PREPASS_BM - 1) / BLK_PREPASS_BM; + const int prepass_q_block = (block_q_idx * BLOCK_M) / BLK_PREPASS_BM; + blk_base = (const global char *) blk_void + + (((mask_batch_idx * mask_ne2) + mask_head_idx) * n_q_blocks_prepass + prepass_q_block) * n_kv_blocks; + } + + const int dk_off_vec = split_idx * SPLIT_DK_VEC; + ACC_TYPE4 q_priv[SPLIT_DK_VEC]; + if (query_valid) { + const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + my_query_row * q_nb1; + const global float4 * q_ptr = (const global float4 *) (q_base + q_row_offset); + #pragma unroll + for (int i = 0; i < SPLIT_DK_VEC; ++i) { + q_priv[i] = q_ptr[dk_off_vec + i]; + } + } else { + #pragma unroll + for (int i = 0; i < SPLIT_DK_VEC; ++i) q_priv[i] = (ACC_TYPE4)(0.0f); + } + +#ifdef FA_HAVE_INT_DOT + uint q_packed_pf[SPLIT_DK_Q8_BLOCKS * 8]; + float q_d_pf[SPLIT_DK_Q8_BLOCKS]; + #pragma unroll + for (int b = 0; b < SPLIT_DK_Q8_BLOCKS; ++b) { + q_d_pf[b] = quant_q_block_int8_packed(&q_priv[b * 8], &q_packed_pf[b * 8]); + } +#endif + + const int dv_off_vec = split_idx * SPLIT_DV_VEC; + ACC_TYPE4 o_acc[SPLIT_DV_VEC]; + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f); + + ACC_TYPE m_i = FA_M_INIT; + ACC_TYPE l_i = 0.0f; + + float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); + +#ifdef FA_HAVE_INT_DOT + __local uint l_k_packed[BLOCK_N][DK_Q8_BLOCKS_PREFILL * 8]; + __local float l_k_scale [BLOCK_N][DK_Q8_BLOCKS_PREFILL]; +#else + __local half4 l_k[BLOCK_N][DK_VEC]; +#endif + +#if FA_V_STRATEGY == 2 + __local uint l_v_packed[BLOCK_N][DV_Q8_BLOCKS_PREFILL * 8]; + __local float l_v_scale [BLOCK_N][DV_Q8_BLOCKS_PREFILL]; +#else + __local half4 l_v[BLOCK_N][DV_VEC]; +#endif + + for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) { + // Skip fully-masked KV tiles (uniform branch across WG). + char blk_cur = 1; + if (blk_base != NULL) { + blk_cur = blk_base[k_start / BLOCK_N]; + if (blk_cur == 0) continue; + } + + { +#ifdef FA_HAVE_INT_DOT + const int k_blocks_per_row = DK_Q8_BLOCKS_PREFILL; + const int n_blocks_total = BLOCK_N * k_blocks_per_row; + for (int i = tid; i < n_blocks_total; i += WG_SIZE) { + const int row = i / k_blocks_per_row; + const int blk = i % k_blocks_per_row; + const int k_row_idx = k_start + row; + if (k_row_idx < n_kv) { + const ulong k_row_off = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_row_idx * k_nb1; + const global char * blk_ptr = k_base + k_row_off + blk * Q8_0_BLOCK_SIZE; + const float df = (float) vload_half(0, (const global half *) blk_ptr); + const global uchar * qs = (const global uchar *)(blk_ptr + 2); + l_k_scale[row][blk] = df; + #pragma unroll + for (int j = 0; j < 8; ++j) { + uint k_packed = + (uint) qs[j*4 + 0] | + ((uint) qs[j*4 + 1]) << 8 | + ((uint) qs[j*4 + 2]) << 16 | + ((uint) qs[j*4 + 3]) << 24; + l_k_packed[row][blk * 8 + j] = k_packed; + } + } else { + l_k_scale[row][blk] = 0.0f; + #pragma unroll + for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u; + } + } +#else + // Fallback: dequant q8_0 -> half in local memory. + const int k_blocks_per_row = DK / QK8_0; + const int n_blocks_total = BLOCK_N * k_blocks_per_row; + for (int i = tid; i < n_blocks_total; i += WG_SIZE) { + const int row = i / k_blocks_per_row; + const int blk = i % k_blocks_per_row; + const int k_row_idx = k_start + row; + if (k_row_idx < n_kv) { + const ulong k_row_off = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_row_idx * k_nb1; + const global char * blk_ptr = k_base + k_row_off + blk * Q8_0_BLOCK_SIZE; + const float df = (float) vload_half(0, (const global half *) blk_ptr); + const global char * qs = blk_ptr + 2; + #pragma unroll + for (int j = 0; j < 8; ++j) { + const float4 v = df * (float4)((float) qs[j*4 + 0], + (float) qs[j*4 + 1], + (float) qs[j*4 + 2], + (float) qs[j*4 + 3]); + l_k[row][blk * 8 + j] = (half4)((half) v.s0, (half) v.s1, (half) v.s2, (half) v.s3); + } + } else { + #pragma unroll + for (int j = 0; j < 8; ++j) l_k[row][blk * 8 + j] = (half4)(0.0h); + } + } +#endif + } + // V tile load — strategy-dependent. +#if FA_V_STRATEGY == 2 + { + // Int8 packed V in local memory + per-block scale. Accumulate + // step unpacks inline. + const int v_blocks_per_row = DV_Q8_BLOCKS_PREFILL; + const int n_blocks_total = BLOCK_N * v_blocks_per_row; + for (int i = tid; i < n_blocks_total; i += WG_SIZE) { + const int row = i / v_blocks_per_row; + const int blk = i % v_blocks_per_row; + const int v_row_idx = k_start + row; + if (v_row_idx < n_kv) { + const ulong v_row_off = batch_idx * v_nb3 + head_kv_idx * v_nb2 + v_row_idx * v_nb1; + const global char * blk_ptr = v_base + v_row_off + blk * Q8_0_BLOCK_SIZE; + const float df = (float) vload_half(0, (const global half *) blk_ptr); + const global uchar * qs = (const global uchar *)(blk_ptr + 2); + l_v_scale[row][blk] = df; + #pragma unroll + for (int j = 0; j < 8; ++j) { + uint v_packed = + (uint) qs[j*4 + 0] | + ((uint) qs[j*4 + 1]) << 8 | + ((uint) qs[j*4 + 2]) << 16 | + ((uint) qs[j*4 + 3]) << 24; + l_v_packed[row][blk * 8 + j] = v_packed; + } + } else { + l_v_scale[row][blk] = 0.0f; + #pragma unroll + for (int j = 0; j < 8; ++j) l_v_packed[row][blk * 8 + j] = 0u; + } + } + } +#else + { + // Default: dequant V -> half in local memory. + const int v_blocks_per_row = DV / QK8_0; + const int n_blocks_total = BLOCK_N * v_blocks_per_row; + for (int i = tid; i < n_blocks_total; i += WG_SIZE) { + const int row = i / v_blocks_per_row; + const int blk = i % v_blocks_per_row; + const int v_row_idx = k_start + row; + if (v_row_idx < n_kv) { + const ulong v_row_off = batch_idx * v_nb3 + head_kv_idx * v_nb2 + v_row_idx * v_nb1; + const global char * blk_ptr = v_base + v_row_off + blk * Q8_0_BLOCK_SIZE; + const float df = (float) vload_half(0, (const global half *) blk_ptr); + const global char * qs = blk_ptr + 2; + #pragma unroll + for (int j = 0; j < 8; ++j) { + const float4 v = df * (float4)((float) qs[j*4 + 0], + (float) qs[j*4 + 1], + (float) qs[j*4 + 2], + (float) qs[j*4 + 3]); + l_v[row][blk * 8 + j] = (half4)((half) v.s0, (half) v.s1, (half) v.s2, (half) v.s3); + } + } else { + #pragma unroll + for (int j = 0; j < 8; ++j) l_v[row][blk * 8 + j] = (half4)(0.0h); + } + } + } +#endif + barrier(CLK_LOCAL_MEM_FENCE); + + // QK dot + online softmax. N_SPLIT>1 reduces per-thread partials via shuffle_xor. +#if N_SPLIT > 1 + { +#else + if (query_valid) { +#endif + const int k_blk_base = split_idx * SPLIT_DK_Q8_BLOCKS; + for (int j = 0; j < BLOCK_N; j += 4) { + const int k_row0 = k_start + j; + const int k_row1 = k_start + j + 1; + const int k_row2 = k_start + j + 2; + const int k_row3 = k_start + j + 3; + + ACC_TYPE s0, s1, s2, s3; +#ifdef FA_HAVE_INT_DOT + // dp4a-accelerated QK dot over owned blocks. + s0 = 0.0f; s1 = 0.0f; s2 = 0.0f; s3 = 0.0f; + #pragma unroll + for (int b_local = 0; b_local < SPLIT_DK_Q8_BLOCKS; ++b_local) { + const int b = k_blk_base + b_local; + int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0; + #pragma unroll + for (int g = 0; g < 8; ++g) { + const uint qp = q_packed_pf[b_local * 8 + g]; + sum0 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j ][b * 8 + g], sum0); + sum1 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+1][b * 8 + g], sum1); + sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2); + sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3); + } + const float qd = q_d_pf[b_local]; + s0 += (float)sum0 * qd * l_k_scale[j ][b]; + s1 += (float)sum1 * qd * l_k_scale[j+1][b]; + s2 += (float)sum2 * qd * l_k_scale[j+2][b]; + s3 += (float)sum3 * qd * l_k_scale[j+3][b]; + } +#else + ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc1 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc2 = (ACC_TYPE4)(0.0f); + ACC_TYPE4 dot_acc3 = (ACC_TYPE4)(0.0f); + #pragma unroll + for (int k = 0; k < SPLIT_DK_VEC; ++k) { + const ACC_TYPE4 qk = q_priv[k]; + const int k_abs = dk_off_vec + k; + dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j ][k_abs]), dot_acc0); + dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k_abs]), dot_acc1); + dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k_abs]), dot_acc2); + dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k_abs]), dot_acc3); + } + s0 = dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3; + s1 = dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3; + s2 = dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3; + s3 = dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3; +#endif + +#if N_SPLIT > 1 + // Power-of-2 N_SPLIT: shuffle_xor butterfly. N_SPLIT=3 (DK=96): 3-way shuffle. + #if (N_SPLIT & (N_SPLIT - 1)) == 0 + #pragma unroll + for (int step = 1; step < N_SPLIT; step <<= 1) { + s0 += sub_group_shuffle_xor(s0, step); + s1 += sub_group_shuffle_xor(s1, step); + s2 += sub_group_shuffle_xor(s2, step); + s3 += sub_group_shuffle_xor(s3, step); + } + #else + const uint tri_base = (get_sub_group_local_id() / N_SPLIT) * N_SPLIT; + s0 = sub_group_shuffle(s0, tri_base + 0) + sub_group_shuffle(s0, tri_base + 1) + sub_group_shuffle(s0, tri_base + 2); + s1 = sub_group_shuffle(s1, tri_base + 0) + sub_group_shuffle(s1, tri_base + 1) + sub_group_shuffle(s1, tri_base + 2); + s2 = sub_group_shuffle(s2, tri_base + 0) + sub_group_shuffle(s2, tri_base + 1) + sub_group_shuffle(s2, tri_base + 2); + s3 = sub_group_shuffle(s3, tri_base + 0) + sub_group_shuffle(s3, tri_base + 1) + sub_group_shuffle(s3, tri_base + 2); + #endif + if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; } +#endif + s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale; + + if (is_causal) { + const int causal_limit = n_kv - n_q + my_query_row; + if (k_row0 > causal_limit) s0 = FA_M_INIT; + if (k_row1 > causal_limit) s1 = FA_M_INIT; + if (k_row2 > causal_limit) s2 = FA_M_INIT; + if (k_row3 > causal_limit) s3 = FA_M_INIT; + } + if (k_row0 >= n_kv) s0 = FA_M_INIT; + if (k_row1 >= n_kv) s1 = FA_M_INIT; + if (k_row2 >= n_kv) s2 = FA_M_INIT; + if (k_row3 >= n_kv) s3 = FA_M_INIT; + + if (query_valid && mask_base != NULL && blk_cur != 2) { + const global MASK_DATA_TYPE * mask_ptr = + (const global MASK_DATA_TYPE *) (mask_base + my_query_row * mask_nb1); + if (k_row0 < n_kv) s0 += slope * (ACC_TYPE) mask_ptr[k_row0]; + if (k_row1 < n_kv) s1 += slope * (ACC_TYPE) mask_ptr[k_row1]; + if (k_row2 < n_kv) s2 += slope * (ACC_TYPE) mask_ptr[k_row2]; + if (k_row3 < n_kv) s3 += slope * (ACC_TYPE) mask_ptr[k_row3]; + } + if (logit_softcap > 0.0f) { + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); + } + + const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); + // Whole tile masked (m_new == FA_M_INIT): force the exp() args + // far negative so the tile contributes 0, not exp(0)=1. + const ACC_TYPE m_exp = (m_new == FA_M_INIT) ? 0.0f : m_new; + const ACC_TYPE scale_prev = native_exp(m_i - m_exp); + const ACC_TYPE p0 = native_exp(s0 - m_exp); + const ACC_TYPE p1 = native_exp(s1 - m_exp); + const ACC_TYPE p2 = native_exp(s2 - m_exp); + const ACC_TYPE p3 = native_exp(s3 - m_exp); + +#if FA_V_STRATEGY == 2 + #pragma unroll + for (int b_local = 0; b_local < DV_Q8_BLOCKS_PREFILL / N_SPLIT; ++b_local) { + const int b_abs = split_idx * (DV_Q8_BLOCKS_PREFILL / N_SPLIT) + b_local; + const float d0 = l_v_scale[j ][b_abs]; + const float d1 = l_v_scale[j+1][b_abs]; + const float d2 = l_v_scale[j+2][b_abs]; + const float d3 = l_v_scale[j+3][b_abs]; + #pragma unroll + for (int g = 0; g < 8; ++g) { + const int lane_abs = b_abs * 8 + g; + const int lane_local = b_local * 8 + g; + uint pk0 = l_v_packed[j ][lane_abs]; + uint pk1 = l_v_packed[j+1][lane_abs]; + uint pk2 = l_v_packed[j+2][lane_abs]; + uint pk3 = l_v_packed[j+3][lane_abs]; + float4 v0 = d0 * (float4)((float)(char)(pk0 & 0xff), (float)(char)((pk0>>8)&0xff), (float)(char)((pk0>>16)&0xff), (float)(char)((pk0>>24)&0xff)); + float4 v1 = d1 * (float4)((float)(char)(pk1 & 0xff), (float)(char)((pk1>>8)&0xff), (float)(char)((pk1>>16)&0xff), (float)(char)((pk1>>24)&0xff)); + float4 v2 = d2 * (float4)((float)(char)(pk2 & 0xff), (float)(char)((pk2>>8)&0xff), (float)(char)((pk2>>16)&0xff), (float)(char)((pk2>>24)&0xff)); + float4 v3 = d3 * (float4)((float)(char)(pk3 & 0xff), (float)(char)((pk3>>8)&0xff), (float)(char)((pk3>>16)&0xff), (float)(char)((pk3>>24)&0xff)); + o_acc[lane_local] = mad(p3, v3, + mad(p2, v2, + mad(p1, v1, + mad(p0, v0, + o_acc[lane_local] * scale_prev)))); + } + } +#else // FA_V_STRATEGY == 0 + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) { + const int i_abs = dv_off_vec + i; + o_acc[i] = mad(p3, CONVERT_KV_ACC4(l_v[j+3][i_abs]), + mad(p2, CONVERT_KV_ACC4(l_v[j+2][i_abs]), + mad(p1, CONVERT_KV_ACC4(l_v[j+1][i_abs]), + mad(p0, CONVERT_KV_ACC4(l_v[j ][i_abs]), + o_acc[i] * scale_prev)))); + } +#endif + l_i = l_i * scale_prev + p0 + p1 + p2 + p3; + m_i = m_new; + } + } + barrier(CLK_LOCAL_MEM_FENCE); + } + + // Write output. With N_SPLIT>1 each thread writes its SPLIT_DV_VEC slice. + if (query_valid) { + if (sinks_void != NULL) { + const global ACC_TYPE * sinks_ptr = + (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset); + const ACC_TYPE m_sink = sinks_ptr[head_idx]; + const ACC_TYPE m_final = max(m_i, m_sink); + const ACC_TYPE scale_o = exp(m_i - m_final); + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) o_acc[i] *= scale_o; + l_i = l_i * scale_o + exp(m_sink - m_final); + m_i = m_final; + } + const ACC_TYPE l_inv = (l_i > 0.0f) ? (1.0f / l_i) : 0.0f; + const ulong o_row_offset = batch_idx * o_nb3 + my_query_row * o_nb2 + head_idx * o_nb1; + global float4 * o_row = (global float4 *) (o_base + o_row_offset); + if (l_inv > 0.0f) { + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) o_row[dv_off_vec + i] = o_acc[i] * l_inv; + } else { + #pragma unroll + for (int i = 0; i < SPLIT_DV_VEC; ++i) o_row[dv_off_vec + i] = (float4)(0.0f); + } + } +} + +// FD Pass 2: merge split partials. Identical across q4_0/q8_0/f16; each FA +// source owns a copy since kernels compile per-source-program. +__kernel void flash_attn_f32_merge( + const global float * partial_void, + global void * o_void, + const ulong o_offset, + const int n_head, + const int n_splits, + const ulong o_nb1, const ulong o_nb2, const ulong o_nb3, + const global void * sinks_void, + const ulong sinks_offset, + const int n_q +) { + const int lane = get_local_id(0); + const int head_batch_idx = get_global_id(1); + const int q_idx = get_global_id(2); + const int batch_idx = head_batch_idx / n_head; + const int head_idx = head_batch_idx % n_head; + + const ulong record_stride = (ulong) FA_PARTIAL_FLOATS; + const ulong record_idx_0 = (((ulong) batch_idx * n_head + head_idx) * n_q + q_idx) * n_splits; + const global float * rec0 = partial_void + record_idx_0 * record_stride; + + __local ACC_TYPE m_final_shared; + __local ACC_TYPE l_final_shared; + if (lane == 0) { + ACC_TYPE m = FA_M_INIT; + for (int c = 0; c < n_splits; ++c) { + const ACC_TYPE m_c = rec0[c * record_stride + 0]; + m = max(m, m_c); + } + ACC_TYPE m_sink = 0.0f; + bool has_sink = false; + if (sinks_void != NULL) { + const global ACC_TYPE * sinks_ptr = + (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset); + m_sink = sinks_ptr[head_idx]; + has_sink = true; + m = max(m, m_sink); + } + ACC_TYPE l = 0.0f; + for (int c = 0; c < n_splits; ++c) { + const ACC_TYPE m_c = rec0[c * record_stride + 0]; + const ACC_TYPE l_c = rec0[c * record_stride + 1]; + if (m_c > FA_M_INIT) { + l += l_c * exp(m_c - m); + } + } + if (has_sink) { + l += exp(m_sink - m); + } + m_final_shared = m; + l_final_shared = l; + } + barrier(CLK_LOCAL_MEM_FENCE); + const ACC_TYPE m_final = m_final_shared; + const ACC_TYPE l_final = l_final_shared; + const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f; + + ACC_TYPE4 o = (ACC_TYPE4)(0.0f); + for (int c = 0; c < n_splits; ++c) { + const global float * rec_c = rec0 + c * record_stride; + const ACC_TYPE m_c = rec_c[0]; + if (m_c <= FA_M_INIT) continue; + const global float4 * rec_oc = (const global float4 *) (rec_c + 2); + const ACC_TYPE scale_c = exp(m_c - m_final); + o = mad((ACC_TYPE4)(scale_c), rec_oc[lane], o); + } + o = o * l_inv; + + const ulong o_row_offset = (ulong) batch_idx * o_nb3 + (ulong) q_idx * o_nb2 + (ulong) head_idx * o_nb1; + global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) ((global char *) o_void + o_offset + o_row_offset); + o_row[lane] = CONVERT_O_DATA4(o); +} diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_pre_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_pre_f16.cl new file mode 100644 index 000000000..88ead4bcb --- /dev/null +++ b/ggml/src/ggml-opencl/kernels/flash_attn_pre_f16.cl @@ -0,0 +1,156 @@ +#pragma OPENCL EXTENSION cl_khr_fp16 : enable + +__kernel void flash_attn_kv_pad_f16( + const global void * k_void, ulong k_offset, + const global void * v_void, ulong v_offset, + global void * k_pad_void, + global void * v_pad_void, + const int n_kv, + const int n_head_kv, + const int n_batch, + const ulong k_nb1, const ulong k_nb2, const ulong k_nb3, + const ulong v_nb1, const ulong v_nb2, const ulong v_nb3 +) { + const int row_idx = get_global_id(0); + const int head_kv_idx = get_global_id(1); + const int batch_idx = get_global_id(2); + + if (row_idx >= BLOCK_N || head_kv_idx >= n_head_kv || batch_idx >= n_batch) { + return; + } + + const int tail_start = n_kv - (n_kv % BLOCK_N); + const int src_row_idx = tail_start + row_idx; + + const global char * k_src = (const global char *) k_void + k_offset; + const global char * v_src = (const global char *) v_void + v_offset; + global char * k_pad = (global char *) k_pad_void; + global char * v_pad = (global char *) v_pad_void; + + const ulong k_dst_offset = ((ulong) batch_idx * (ulong) n_head_kv + (ulong) head_kv_idx) * ((ulong) BLOCK_N * k_nb1) + (ulong) row_idx * k_nb1; + const ulong v_dst_offset = ((ulong) batch_idx * (ulong) n_head_kv + (ulong) head_kv_idx) * ((ulong) BLOCK_N * v_nb1) + (ulong) row_idx * v_nb1; + + if (src_row_idx < n_kv) { + const ulong k_src_offset = (ulong) batch_idx * k_nb3 + (ulong) head_kv_idx * k_nb2 + (ulong) src_row_idx * k_nb1; + const ulong v_src_offset = (ulong) batch_idx * v_nb3 + (ulong) head_kv_idx * v_nb2 + (ulong) src_row_idx * v_nb1; + + for (ulong i = 0; i < k_nb1; ++i) { + k_pad[k_dst_offset + i] = k_src[k_src_offset + i]; + } + for (ulong i = 0; i < v_nb1; ++i) { + v_pad[v_dst_offset + i] = v_src[v_src_offset + i]; + } + } else { + for (ulong i = 0; i < k_nb1; ++i) { + k_pad[k_dst_offset + i] = 0; + } + for (ulong i = 0; i < v_nb1; ++i) { + v_pad[v_dst_offset + i] = 0; + } + } +} + +__kernel void flash_attn_mask_pad_f16( + const global void * mask_void, ulong mask_offset, + global void * mask_pad_void, + const int n_q, + const int n_kv, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3 +) { + const int col_idx = get_global_id(0); + const int q_row = get_global_id(1); + const int mask_slice = get_global_id(2); + + if (col_idx >= BLOCK_N || q_row >= n_q || mask_slice >= mask_ne2 * mask_ne3) { + return; + } + + const int tail_start = n_kv - (n_kv % BLOCK_N); + const int src_col_idx = tail_start + col_idx; + const int mask_head_idx = mask_slice % mask_ne2; + const int mask_batch_idx = mask_slice / mask_ne2; + + const global char * mask_src_base = (const global char *) mask_void + mask_offset + + (ulong) mask_batch_idx * mask_nb3 + + (ulong) mask_head_idx * mask_nb2 + + (ulong) q_row * mask_nb1; + const global half * mask_src = (const global half *) mask_src_base; + + global half * mask_pad = (global half *) mask_pad_void; + const ulong dst_idx = + (((ulong) mask_batch_idx * (ulong) mask_ne2 + (ulong) mask_head_idx) * (ulong) n_q + (ulong) q_row) * (ulong) BLOCK_N + + (ulong) col_idx; + + mask_pad[dst_idx] = src_col_idx < n_kv ? mask_src[src_col_idx] : (half) (-INFINITY); +} + +// Per-KV-tile mask class. 0=all -inf (skip tile), 1=mixed (apply mask), +// 2=all zero, no -inf (skip mask lookup). Causal diagonal tiles are class 1. +__kernel void flash_attn_blk_f16( + const global void * mask_void, ulong mask_offset, + global char * blk, + const int n_q, + const int n_kv, + const ulong mask_nb1, + const ulong mask_nb2, + const ulong mask_nb3, + const int mask_ne2, + const int mask_ne3 +) { + const int kv_block_idx = get_global_id(0); + const int q_block_idx = get_global_id(1); + const int mask_slice = get_global_id(2); + + const int n_q_blocks = (n_q + BLOCK_M - 1) / BLOCK_M; + const int n_kv_blocks = (n_kv + BLOCK_N - 1) / BLOCK_N; + if (kv_block_idx >= n_kv_blocks || q_block_idx >= n_q_blocks || mask_slice >= mask_ne2 * mask_ne3) { + return; + } + + const int mask_head_idx = mask_slice % mask_ne2; + const int mask_batch_idx = mask_slice / mask_ne2; + const int q_start = q_block_idx * BLOCK_M; + const int k_start = kv_block_idx * BLOCK_N; + const int q_count = min(BLOCK_M, n_q - q_start); + const int k_count = min(BLOCK_N, n_kv - k_start); + + const half neg_max_half = (half) (-65504.0f); + char has_unmasked = 0; + char has_masked = 0; + char has_nonzero = 0; + + const global char * mask_base = (const global char *) mask_void + mask_offset + + (ulong) mask_batch_idx * mask_nb3 + + (ulong) mask_head_idx * mask_nb2; + + for (int qi = 0; qi < q_count; ++qi) { + const global half * mask_row = (const global half *) (mask_base + (ulong) (q_start + qi) * mask_nb1) + k_start; + for (int ki = 0; ki < k_count; ++ki) { + const half v = mask_row[ki]; + if (v <= neg_max_half) { + has_masked = 1; + } else { + has_unmasked = 1; + if (v != (half) 0.0f) { + has_nonzero = 1; + } + } + } + if (has_masked && has_unmasked) break; // mixed tile — short-circuit. + } + + char res; + if (has_unmasked == 0) { + res = 0; + } else if (has_masked || has_nonzero) { + res = 1; + } else { + res = 2; + } + + blk[((ulong) mask_slice * (ulong) n_q_blocks + (ulong) q_block_idx) * (ulong) n_kv_blocks + (ulong) kv_block_idx] = res; +} diff --git a/ggml/src/ggml-opencl/kernels/set_rows.cl b/ggml/src/ggml-opencl/kernels/set_rows.cl index fc3ff7aa1..4ad5af13f 100644 --- a/ggml/src/ggml-opencl/kernels/set_rows.cl +++ b/ggml/src/ggml-opencl/kernels/set_rows.cl @@ -158,6 +158,239 @@ kernel void kernel_set_rows_f32_i32( } } +// f32 -> q8_0 quantize set_rows. Block = half d + char qs[32]. +#define QK8_0 32 + +inline void quantize_q8_0_block(global float * x, global char * qs, global half * d_out) { + float amax = 0.0f; + for (int j = 0; j < QK8_0; j++) { + amax = fmax(amax, fabs(x[j])); + } + + float d = amax / 127.0f; + float id = (d != 0.0f) ? 127.0f / amax : 0.0f; + + vstore_half(d, 0, d_out); + + for (int j = 0; j < QK8_0; j++) { + qs[j] = (char)((int)round(x[j] * id)); + } +} + +kernel void kernel_set_rows_q8_0_i64( + global char * src0, + ulong offset0, + global char * src1, + ulong offset1, + global char * dst, + ulong offsetd, + int ne01, + ulong nb01, + ulong nb02, + ulong nb03, + uint4 ne11, + uint4 ne12, + ulong nb10, + ulong nb11, + ulong nb12, + int nblk0, + ulong nb1, + ulong nb2, + ulong nb3 +) { + src0 = src0 + offset0; + src1 = src1 + offset1; + dst = dst + offsetd; + + int i03 = get_group_id(2); + int i02 = get_group_id(1); + int i01 = get_group_id(0)*get_local_size(1) + get_local_id(1); + + if (i01 >= ne01) { + return; + } + + int i12 = fastmod(i03, ne12); + int i11 = fastmod(i02, ne11); + + int i10 = i01; + long i1 = ((global long *)(src1 + i10*nb10 + i11*nb11 + i12*nb12))[0]; + + global char * dst_row = (global char *) (dst + i1*nb1 + i02*nb2 + i03*nb3); + global float * src_row = (global float *) (src0 + i01*nb01 + i02*nb02 + i03*nb03); + + for (int blk = get_local_id(0); blk < nblk0; blk += get_local_size(0)) { + global float * x = src_row + blk * QK8_0; + global char * y = dst_row + blk * (2 + QK8_0); + + quantize_q8_0_block(x, y + 2, (global half *)y); + } +} + +kernel void kernel_set_rows_q8_0_i32( + global char * src0, + ulong offset0, + global char * src1, + ulong offset1, + global char * dst, + ulong offsetd, + int ne01, + ulong nb01, + ulong nb02, + ulong nb03, + uint4 ne11, + uint4 ne12, + ulong nb10, + ulong nb11, + ulong nb12, + int nblk0, + ulong nb1, + ulong nb2, + ulong nb3 +) { + src0 = src0 + offset0; + src1 = src1 + offset1; + dst = dst + offsetd; + + int i03 = get_group_id(2); + int i02 = get_group_id(1); + int i01 = get_group_id(0)*get_local_size(1) + get_local_id(1); + + if (i01 >= ne01) { + return; + } + + int i12 = fastmod(i03, ne12); + int i11 = fastmod(i02, ne11); + + int i10 = i01; + int i1 = ((global int *)(src1 + i10*nb10 + i11*nb11 + i12*nb12))[0]; + + global char * dst_row = (global char *) (dst + i1*nb1 + i02*nb2 + i03*nb3); + global float * src_row = (global float *) (src0 + i01*nb01 + i02*nb02 + i03*nb03); + + for (int blk = get_local_id(0); blk < nblk0; blk += get_local_size(0)) { + global float * x = src_row + blk * QK8_0; + global char * y = dst_row + blk * (2 + QK8_0); + + quantize_q8_0_block(x, y + 2, (global half *)y); + } +} + +// SoA q8_0 variants. dst_q: int8[QK8_0] per block; dst_d: fp16 scale per block. +// Layout matches kernel_convert_block_q8_0; block index follows dst element order. +kernel void kernel_set_rows_q8_0_soa_i64( + global char * src0, + ulong offset0, + global char * src1, + ulong offset1, + global char * dst_q, + ulong offset_q, + global char * dst_d, + ulong offset_d, + int ne01, + ulong nb01, + ulong nb02, + ulong nb03, + uint4 ne11, + uint4 ne12, + ulong nb10, + ulong nb11, + ulong nb12, + int nblk0, + int ne1_dst, + int ne2_dst, + int ne3_dst +) { + src0 = src0 + offset0; + src1 = src1 + offset1; + dst_q = dst_q + offset_q; + dst_d = dst_d + offset_d; + + int i03 = get_group_id(2); + int i02 = get_group_id(1); + int i01 = get_group_id(0)*get_local_size(1) + get_local_id(1); + + if (i01 >= ne01) { + return; + } + + int i12 = fastmod(i03, ne12); + int i11 = fastmod(i02, ne11); + + int i10 = i01; + long i1 = ((global long *)(src1 + i10*nb10 + i11*nb11 + i12*nb12))[0]; + + long row_blk_base = ((long)i03 * ne2_dst * ne1_dst + (long)i02 * ne1_dst + i1) * nblk0; + + global half * d_row = (global half *)(dst_d) + row_blk_base; + global char * q_row = (global char *)(dst_q) + row_blk_base * QK8_0; + global float * src_row = (global float *)(src0 + i01*nb01 + i02*nb02 + i03*nb03); + + for (int blk = get_local_id(0); blk < nblk0; blk += get_local_size(0)) { + global float * x = src_row + blk * QK8_0; + global char * q = q_row + blk * QK8_0; + + quantize_q8_0_block(x, q, d_row + blk); + } +} + +kernel void kernel_set_rows_q8_0_soa_i32( + global char * src0, + ulong offset0, + global char * src1, + ulong offset1, + global char * dst_q, + ulong offset_q, + global char * dst_d, + ulong offset_d, + int ne01, + ulong nb01, + ulong nb02, + ulong nb03, + uint4 ne11, + uint4 ne12, + ulong nb10, + ulong nb11, + ulong nb12, + int nblk0, + int ne1_dst, + int ne2_dst, + int ne3_dst +) { + src0 = src0 + offset0; + src1 = src1 + offset1; + dst_q = dst_q + offset_q; + dst_d = dst_d + offset_d; + + int i03 = get_group_id(2); + int i02 = get_group_id(1); + int i01 = get_group_id(0)*get_local_size(1) + get_local_id(1); + + if (i01 >= ne01) { + return; + } + + int i12 = fastmod(i03, ne12); + int i11 = fastmod(i02, ne11); + + int i10 = i01; + int i1 = ((global int *)(src1 + i10*nb10 + i11*nb11 + i12*nb12))[0]; + + long row_blk_base = ((long)i03 * ne2_dst * ne1_dst + (long)i02 * ne1_dst + i1) * nblk0; + + global half * d_row = (global half *)(dst_d) + row_blk_base; + global char * q_row = (global char *)(dst_q) + row_blk_base * QK8_0; + global float * src_row = (global float *)(src0 + i01*nb01 + i02*nb02 + i03*nb03); + + for (int blk = get_local_id(0); blk < nblk0; blk += get_local_size(0)) { + global float * x = src_row + blk * QK8_0; + global char * q = q_row + blk * QK8_0; + + quantize_q8_0_block(x, q, d_row + blk); + } +} + kernel void kernel_set_rows_f16_i32( global char * src0, ulong offset0, @@ -206,3 +439,270 @@ kernel void kernel_set_rows_f16_i32( dst_row[ind] = src_row[ind]; } } + +// f32 -> q4_0 quantize set_rows. Block = half d + uchar qs[16] (shuffled +// nibbles: qs[j] low/high = elem j / j+16). +// Dequant: val[i] = d * (nibble_i - 8) +// nblk0 = number of q4_0 blocks per row = ne00 / 32. +#define QK4_0 32 +#define Q4_0_BLOCK_SIZE 18 + +inline void quantize_q4_0_block(global float * x, global uchar * qs, global half * d_out) { + // Find the signed value with the largest absolute magnitude (matches ggml ref). + float max = 0.0f; + float amax = 0.0f; + for (int j = 0; j < QK4_0; j++) { + float v = x[j]; + float a = fabs(v); + if (a > amax) { + amax = a; + max = v; + } + } + + float d = max / -8.0f; + float id = (d != 0.0f) ? 1.0f / d : 0.0f; + + vstore_half(d, 0, d_out); + + for (int j = 0; j < QK4_0/2; j++) { + float x0 = x[j] * id; + float x1 = x[j + QK4_0/2] * id; + + int i0 = (int)(x0 + 8.5f); + int i1 = (int)(x1 + 8.5f); + if (i0 < 0) i0 = 0; + if (i0 > 15) i0 = 15; + if (i1 < 0) i1 = 0; + if (i1 > 15) i1 = 15; + + qs[j] = (uchar)i0 | ((uchar)i1 << 4); + } +} + +kernel void kernel_set_rows_q4_0_i64( + global char * src0, + ulong offset0, + global char * src1, + ulong offset1, + global char * dst, + ulong offsetd, + int ne01, + ulong nb01, + ulong nb02, + ulong nb03, + uint4 ne11, + uint4 ne12, + ulong nb10, + ulong nb11, + ulong nb12, + int nblk0, + ulong nb1, + ulong nb2, + ulong nb3 +) { + src0 = src0 + offset0; + src1 = src1 + offset1; + dst = dst + offsetd; + + int i03 = get_group_id(2); + int i02 = get_group_id(1); + int i01 = get_group_id(0)*get_local_size(1) + get_local_id(1); + + if (i01 >= ne01) { + return; + } + + int i12 = fastmod(i03, ne12); + int i11 = fastmod(i02, ne11); + + int i10 = i01; + long i1 = ((global long *)(src1 + i10*nb10 + i11*nb11 + i12*nb12))[0]; + + global char * dst_row = (global char *) (dst + i1*nb1 + i02*nb2 + i03*nb3); + global float * src_row = (global float *) (src0 + i01*nb01 + i02*nb02 + i03*nb03); + + for (int blk = get_local_id(0); blk < nblk0; blk += get_local_size(0)) { + global float * x = src_row + blk * QK4_0; + global char * y = dst_row + blk * Q4_0_BLOCK_SIZE; + global half * yd = (global half *)(y); + global uchar * yqs = (global uchar *)(y + 2); + + quantize_q4_0_block(x, yqs, yd); + } +} + +kernel void kernel_set_rows_q4_0_i32( + global char * src0, + ulong offset0, + global char * src1, + ulong offset1, + global char * dst, + ulong offsetd, + int ne01, + ulong nb01, + ulong nb02, + ulong nb03, + uint4 ne11, + uint4 ne12, + ulong nb10, + ulong nb11, + ulong nb12, + int nblk0, + ulong nb1, + ulong nb2, + ulong nb3 +) { + src0 = src0 + offset0; + src1 = src1 + offset1; + dst = dst + offsetd; + + int i03 = get_group_id(2); + int i02 = get_group_id(1); + int i01 = get_group_id(0)*get_local_size(1) + get_local_id(1); + + if (i01 >= ne01) { + return; + } + + int i12 = fastmod(i03, ne12); + int i11 = fastmod(i02, ne11); + + int i10 = i01; + int i1 = ((global int *)(src1 + i10*nb10 + i11*nb11 + i12*nb12))[0]; + + global char * dst_row = (global char *) (dst + i1*nb1 + i02*nb2 + i03*nb3); + global float * src_row = (global float *) (src0 + i01*nb01 + i02*nb02 + i03*nb03); + + for (int blk = get_local_id(0); blk < nblk0; blk += get_local_size(0)) { + global float * x = src_row + blk * QK4_0; + global char * y = dst_row + blk * Q4_0_BLOCK_SIZE; + global half * yd = (global half *)(y); + global uchar * yqs = (global uchar *)(y + 2); + + quantize_q4_0_block(x, yqs, yd); + } +} + +// SoA variants for q4_0 dst. Used when the backend has split block_q4_0 records +// into separate quant (dst_q) and scale (dst_d) sub-buffers — same pattern as +// the q8_0 SoA variants above. +// +// Layout (matches kernel_convert_block_q4_0, the "shuffled" variant): +// dst_q: contiguous 16 packed nibbles per block, block i at offset i * 16 bytes. +// dst_d: contiguous fp16 scales, block i at offset i * 2 bytes. +// Nibble layout inside each byte is unchanged from AoS: qs[j] low nibble = element j, +// qs[j] high nibble = element j+16. kernel_restore_block_q4_0 copies bytes as-is. +kernel void kernel_set_rows_q4_0_soa_i64( + global char * src0, + ulong offset0, + global char * src1, + ulong offset1, + global char * dst_q, + ulong offset_q, + global char * dst_d, + ulong offset_d, + int ne01, + ulong nb01, + ulong nb02, + ulong nb03, + uint4 ne11, + uint4 ne12, + ulong nb10, + ulong nb11, + ulong nb12, + int nblk0, + int ne1_dst, + int ne2_dst, + int ne3_dst +) { + src0 = src0 + offset0; + src1 = src1 + offset1; + dst_q = dst_q + offset_q; + dst_d = dst_d + offset_d; + + int i03 = get_group_id(2); + int i02 = get_group_id(1); + int i01 = get_group_id(0)*get_local_size(1) + get_local_id(1); + + if (i01 >= ne01) { + return; + } + + int i12 = fastmod(i03, ne12); + int i11 = fastmod(i02, ne11); + + int i10 = i01; + long i1 = ((global long *)(src1 + i10*nb10 + i11*nb11 + i12*nb12))[0]; + + long row_blk_base = ((long)i03 * ne2_dst * ne1_dst + (long)i02 * ne1_dst + i1) * nblk0; + + global half * d_row = (global half *)(dst_d) + row_blk_base; + global uchar * q_row = (global uchar *)(dst_q) + row_blk_base * (QK4_0/2); + global float * src_row = (global float *)(src0 + i01*nb01 + i02*nb02 + i03*nb03); + + for (int blk = get_local_id(0); blk < nblk0; blk += get_local_size(0)) { + global float * x = src_row + blk * QK4_0; + global uchar * qs = q_row + blk * (QK4_0/2); + global half * d_bk = d_row + blk; + + quantize_q4_0_block(x, qs, d_bk); + } +} + +kernel void kernel_set_rows_q4_0_soa_i32( + global char * src0, + ulong offset0, + global char * src1, + ulong offset1, + global char * dst_q, + ulong offset_q, + global char * dst_d, + ulong offset_d, + int ne01, + ulong nb01, + ulong nb02, + ulong nb03, + uint4 ne11, + uint4 ne12, + ulong nb10, + ulong nb11, + ulong nb12, + int nblk0, + int ne1_dst, + int ne2_dst, + int ne3_dst +) { + src0 = src0 + offset0; + src1 = src1 + offset1; + dst_q = dst_q + offset_q; + dst_d = dst_d + offset_d; + + int i03 = get_group_id(2); + int i02 = get_group_id(1); + int i01 = get_group_id(0)*get_local_size(1) + get_local_id(1); + + if (i01 >= ne01) { + return; + } + + int i12 = fastmod(i03, ne12); + int i11 = fastmod(i02, ne11); + + int i10 = i01; + int i1 = ((global int *)(src1 + i10*nb10 + i11*nb11 + i12*nb12))[0]; + + long row_blk_base = ((long)i03 * ne2_dst * ne1_dst + (long)i02 * ne1_dst + i1) * nblk0; + + global half * d_row = (global half *)(dst_d) + row_blk_base; + global uchar * q_row = (global uchar *)(dst_q) + row_blk_base * (QK4_0/2); + global float * src_row = (global float *)(src0 + i01*nb01 + i02*nb02 + i03*nb03); + + for (int blk = get_local_id(0); blk < nblk0; blk += get_local_size(0)) { + global float * x = src_row + blk * QK4_0; + global uchar * qs = q_row + blk * (QK4_0/2); + global half * d_bk = d_row + blk; + + quantize_q4_0_block(x, qs, d_bk); + } +} From 27c8bb4f63ad9f20bf5901067810a4be5ffe20c4 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Sun, 28 Jun 2026 08:52:15 +0300 Subject: [PATCH 08/17] logs : reduce v2 (#25078) * server : reduce logs * cont : common * cont : spec * cont : CMN_ -> COM_ --- common/common.cpp | 94 +++++++++++----------- common/common.h | 7 ++ common/fit.cpp | 2 +- common/reasoning-budget.cpp | 20 ++--- common/speculative.cpp | 133 +++++++++++++++++--------------- src/llama-context.cpp | 2 +- tools/server/server-context.cpp | 86 +++++++++++---------- tools/server/server-http.cpp | 16 ++-- tools/server/server-schema.cpp | 2 +- tools/server/server-stream.cpp | 11 ++- tools/server/server-task.cpp | 12 +-- tools/server/server.cpp | 12 +-- 12 files changed, 203 insertions(+), 194 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index a14e7bbed..0dd9ede5e 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -225,7 +225,7 @@ bool set_process_priority(enum ggml_sched_priority prio) { } if (!SetPriorityClass(GetCurrentProcess(), p)) { - LOG_WRN("failed to set process priority class %d : (%d)\n", prio, (int) GetLastError()); + COM_WRN("failed to set process priority class %d : (%d)\n", prio, (int) GetLastError()); return false; } @@ -251,7 +251,7 @@ bool set_process_priority(enum ggml_sched_priority prio) { } if (setpriority(PRIO_PROCESS, 0, p) != 0) { - LOG_WRN("failed to set process priority %d : %s (%d)\n", prio, strerror(errno), errno); + COM_WRN("failed to set process priority %d : %s (%d)\n", prio, strerror(errno), errno); return false; } return true; @@ -284,14 +284,14 @@ void postprocess_cpu_params(common_cpu_params & cpuparams, const common_cpu_para if (n_set && n_set < cpuparams.n_threads) { // Not enough set bits, may experience performance issues. - LOG_WRN("Not enough set bits in CPU mask (%d) to satisfy requested thread count: %d\n", n_set, cpuparams.n_threads); + COM_WRN("Not enough set bits in CPU mask (%d) to satisfy requested thread count: %d\n", n_set, cpuparams.n_threads); } } bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THREADS]) { size_t dash_loc = range.find('-'); if (dash_loc == std::string::npos) { - LOG_ERR("Format of CPU range is invalid! Expected []-[].\n"); + COM_ERR("%s", "Format of CPU range is invalid! Expected []-[].\n"); return false; } @@ -303,7 +303,7 @@ bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THRE } else { start_i = std::stoull(range.substr(0, dash_loc)); if (start_i >= GGML_MAX_N_THREADS) { - LOG_ERR("Start index out of bounds!\n"); + COM_ERR("%s", "Start index out of bounds!\n"); return false; } } @@ -313,7 +313,7 @@ bool parse_cpu_range(const std::string & range, bool (&boolmask)[GGML_MAX_N_THRE } else { end_i = std::stoull(range.substr(dash_loc + 1)); if (end_i >= GGML_MAX_N_THREADS) { - LOG_ERR("End index out of bounds!\n"); + COM_ERR("%s", "End index out of bounds!\n"); return false; } } @@ -333,7 +333,7 @@ bool parse_cpu_mask(const std::string & mask, bool (&boolmask)[GGML_MAX_N_THREAD } size_t num_digits = mask.length() - start_i; - if (num_digits > 128) num_digits = 128; + num_digits = std::min(num_digits, 128); size_t end_i = num_digits + start_i; @@ -348,7 +348,7 @@ bool parse_cpu_mask(const std::string & mask, bool (&boolmask)[GGML_MAX_N_THREAD } else if (c >= 'A' && c <= 'F') { id -= 'A' - 10; } else { - LOG_ERR("Invalid hex character '%c' at position %d\n", c, int32_t(i)); + COM_ERR("Invalid hex character '%c' at position %d\n", c, int32_t(i)); return false; } @@ -379,21 +379,21 @@ void common_params_print_info(const common_params & params, bool print_devices) #else const char * build_type = " (debug)"; #endif - LOG_TRC("%s: build %d (%s) with %s for %s%s\n", __func__, llama_build_number(), llama_commit(), llama_compiler(), llama_build_target(), build_type); + COM_TRC("%s: build %d (%s) with %s for %s%s\n", __func__, llama_build_number(), llama_commit(), llama_compiler(), llama_build_target(), build_type); - LOG_INF("log_info: verbosity = %d (adjust with the `-lv N` CLI arg)\n", common_log_get_verbosity_thold()); + COM_INF("%s: verbosity = %d (adjust with the `-lv N` CLI arg)\n", __func__, common_log_get_verbosity_thold()); // device enumeration creates a primary context on CUDA backends, skip it when the caller does not own any device if (print_devices) { - LOG_INF("device_info:\n"); + COM_TRC("%s", "device_info:\n"); for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { auto * dev = ggml_backend_dev_get(i); size_t free, total; ggml_backend_dev_memory(dev, &free, &total); - LOG_INF(" - %-8s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024); + COM_TRC(" - %-8s: %s (%zu MiB, %zu MiB free)\n", ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024); } } - LOG_INF("%s\n", common_params_get_system_info(params).c_str()); + COM_TRC("%s\n", common_params_get_system_info(params).c_str()); } std::string common_params_get_system_info(const common_params & params) { @@ -660,7 +660,7 @@ void string_process_escapes(std::string & input) { bool string_parse_kv_override(const char * data, std::vector & overrides) { const char * sep = strchr(data, '='); if (sep == nullptr || sep - data >= 128) { - LOG_ERR("%s: malformed KV override '%s'\n", __func__, data); + COM_ERR("%s: malformed KV override '%s'\n", __func__, data); return false; } llama_model_kv_override kvo; @@ -683,20 +683,20 @@ bool string_parse_kv_override(const char * data, std::vector 127) { - LOG_ERR("%s: malformed KV override '%s', value cannot exceed 127 chars\n", __func__, data); + COM_ERR("%s: malformed KV override '%s', value cannot exceed 127 chars\n", __func__, data); return false; } strncpy(kvo.val_str, sep, 127); kvo.val_str[127] = '\0'; } else { - LOG_ERR("%s: invalid type for KV override '%s'\n", __func__, data); + COM_ERR("%s: invalid type for KV override '%s'\n", __func__, data); return false; } overrides.emplace_back(std::move(kvo)); @@ -1199,8 +1199,8 @@ common_init_result::common_init_result(common_params & params, bool model_only) auto cparams = common_context_params_to_llama(params); if (params.fit_params) { - LOG_INF("%s: fitting params to device memory ...\n", __func__); - LOG_INF("%s: (for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n", __func__); + COM_TRC("%s", "fitting params to device memory ...\n"); + COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n"); common_fit_params(params.model.path.c_str(), &mparams, &cparams, params.tensor_split, params.tensor_buft_overrides.data(), @@ -1227,7 +1227,7 @@ common_init_result::common_init_result(common_params & params, bool model_only) llama_adapter_lora_ptr lora; lora.reset(llama_adapter_lora_init(model, la.path.c_str())); if (lora == nullptr) { - LOG_ERR("%s: failed to load lora adapter '%s'\n", __func__, la.path.c_str()); + COM_ERR("failed to load lora adapter '%s'\n", la.path.c_str()); pimpl->model.reset(model); return; } @@ -1246,14 +1246,14 @@ common_init_result::common_init_result(common_params & params, bool model_only) common_init_sampler_from_model(model, params.sampling); if (params.sampling.ignore_eos && llama_vocab_eos(vocab) == LLAMA_TOKEN_NULL) { - LOG_WRN("%s: warning: vocab does not have an EOS token, ignoring --ignore-eos\n", __func__); + COM_WRN("%s", "vocab does not have an EOS token, ignoring --ignore-eos\n"); params.sampling.ignore_eos = false; } // initialize once for (llama_token i = 0; i < llama_vocab_n_tokens(vocab); i++) { if (llama_vocab_is_eog(vocab, i)) { - LOG_TRC("%s: added %s logit bias = %f\n", __func__, common_token_to_piece(vocab, i).c_str(), -INFINITY); + COM_TRC("added %s logit bias = %f\n", common_token_to_piece(vocab, i).c_str(), -INFINITY); params.sampling.logit_bias_eog.push_back({i, -INFINITY}); } } @@ -1291,7 +1291,7 @@ common_init_result::common_init_result(common_params & params, bool model_only) llama_context * lctx = llama_init_from_model(model, cparams); if (lctx == NULL) { - LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str()); + COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str()); return; } @@ -1328,7 +1328,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode llama_model * model = res->model(); if (model == NULL) { - LOG_ERR("%s: failed to load model '%s'\n", __func__, params.model.path.c_str()); + COM_ERR("failed to load model '%s'\n", params.model.path.c_str()); return res; } @@ -1338,14 +1338,14 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode llama_context * lctx = res->context(); if (lctx == NULL) { - LOG_ERR("%s: failed to create context with model '%s'\n", __func__, params.model.path.c_str()); + COM_ERR("failed to create context with model '%s'\n", params.model.path.c_str()); return res; } const llama_vocab * vocab = llama_model_get_vocab(model); if (params.ctx_shift && !llama_memory_can_shift(llama_get_memory(lctx))) { - LOG_WRN("%s: KV cache shifting is not supported for this context, disabling KV cache shifting\n", __func__); + COM_WRN("%s", "KV cache shifting is not supported for this context, disabling KV cache shifting\n"); params.ctx_shift = false; } @@ -1374,7 +1374,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode bool ok = true; if (llama_vocab_bos(vocab) == LLAMA_TOKEN_NULL) { - LOG_WRN("%s: warning: vocab does not have a BOS token, reranking will not work\n", __func__); + COM_WRN("%s", "vocab does not have a BOS token, reranking will not work\n"); ok = false; } @@ -1383,10 +1383,10 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode bool has_rerank_prompt = llama_model_chat_template(model, "rerank") != NULL; if (!has_eos && !has_sep && !has_rerank_prompt) { - LOG_WRN("%s: warning: vocab does not have an EOS token, SEP token, or rerank prompt. Reranking will not work\n", __func__); + COM_WRN("%s", "vocab does not have an EOS token, SEP token, or rerank prompt. Reranking will not work\n"); ok = false; } else if (!has_eos) { - LOG_WRN("%s: warning: vocab does not have an EOS token, using SEP token as fallback\n", __func__); + COM_WRN("%s", "vocab does not have an EOS token, using SEP token as fallback\n"); } if (!ok) { @@ -1399,7 +1399,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode } if (params.warmup) { - LOG_INF("%s: warming up the model with an empty run - please wait ... (--no-warmup to disable)\n", __func__); + COM_TRC("%s", "warming up the model with an empty run - please wait ... (--no-warmup to disable)\n"); std::vector tmp; llama_token bos = llama_vocab_bos(vocab); @@ -1473,20 +1473,20 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) { int ret = llama_decode(ctx, llama_batch_get_one(tmp.data(), tmp.size())); if (ret != 0) { - LOG_ERR("%s: llama_decode() failed: %d\n", __func__, ret); + COM_ERR("llama_decode() failed: %d\n", ret); res = COMMON_CONTEXT_SEQ_RM_TYPE_NO; goto done; } if (llama_n_rs_seq(ctx) > 0) { - LOG_INF("%s: the context supports bounded partial sequence removal\n", __func__); + COM_TRC("%s", "the context supports bounded partial sequence removal\n"); res = COMMON_CONTEXT_SEQ_RM_TYPE_RS; goto done; } // try to remove the last tokens if (!llama_memory_seq_rm(mem, 0, 1, -1)) { - LOG_TRC("%s: the context does not support partial sequence removal\n", __func__); + COM_TRC("%s", "the context does not support partial sequence removal\n"); res = COMMON_CONTEXT_SEQ_RM_TYPE_FULL; goto done; } @@ -1803,13 +1803,13 @@ static common_control_vector_data common_control_vector_load_one(const common_co }; struct gguf_context * ctx_gguf = gguf_init_from_file(load_info.fname.c_str(), meta_gguf_params); if (!ctx_gguf) { - LOG_ERR("%s: failed to load control vector file from %s\n", __func__, load_info.fname.c_str()); + COM_ERR("failed to load control vector file from %s\n", load_info.fname.c_str()); return result; } int32_t n_tensors = gguf_get_n_tensors(ctx_gguf); if (n_tensors == 0) { - LOG_WRN("%s: no direction tensors found in %s\n", __func__, load_info.fname.c_str()); + COM_WRN("no direction tensors found in %s\n", load_info.fname.c_str()); } for (int i = 0; i < n_tensors; i++) { @@ -1827,23 +1827,23 @@ static common_control_vector_data common_control_vector_load_one(const common_co } } if (layer_idx < 0) { - LOG_ERR("%s: invalid/unparsable direction tensor layer index in %s\n", __func__, load_info.fname.c_str()); + COM_ERR("invalid/unparsable direction tensor layer index in %s\n", load_info.fname.c_str()); result.n_embd = -1; break; } else if (layer_idx == 0) { - LOG_ERR("%s: invalid (zero) direction tensor layer index in %s\n", __func__, load_info.fname.c_str()); + COM_ERR("invalid (zero) direction tensor layer index in %s\n", load_info.fname.c_str()); result.n_embd = -1; break; } struct ggml_tensor * tensor = ggml_get_tensor(ctx, name.c_str()); if (tensor->type != GGML_TYPE_F32) { - LOG_ERR("%s: invalid (non-F32) direction tensor type in %s\n", __func__, load_info.fname.c_str()); + COM_ERR("invalid (non-F32) direction tensor type in %s\n", load_info.fname.c_str()); result.n_embd = -1; break; } if (ggml_n_dims(tensor) != 1) { - LOG_ERR("%s: invalid (non-1D) direction tensor shape in %s\n", __func__, load_info.fname.c_str()); + COM_ERR("invalid (non-1D) direction tensor shape in %s\n", load_info.fname.c_str()); result.n_embd = -1; break; } @@ -1851,7 +1851,7 @@ static common_control_vector_data common_control_vector_load_one(const common_co if (result.n_embd == -1) { result.n_embd = ggml_nelements(tensor); } else if (ggml_nelements(tensor) != result.n_embd) { - LOG_ERR("%s: direction tensor in %s does not match previous dimensions\n", __func__, load_info.fname.c_str()); + COM_ERR("direction tensor in %s does not match previous dimensions\n", load_info.fname.c_str()); result.n_embd = -1; break; } @@ -1868,7 +1868,7 @@ static common_control_vector_data common_control_vector_load_one(const common_co } if (result.n_embd == -1) { - LOG_WRN("%s: skipping %s due to invalid direction tensors\n", __func__, load_info.fname.c_str()); + COM_WRN("skipping %s due to invalid direction tensors\n", load_info.fname.c_str()); result.data.clear(); } @@ -1889,7 +1889,7 @@ common_control_vector_data common_control_vector_load(const std::vector(all_tokens.data() + offset), n_tokens_before_last))) { - LOG_ERR("%s : failed to eval\n", __func__); + COM_ERR("%s", "failed to eval\n"); return false; } n_past += n_tokens_before_last; llama_state_save_file(ctx, state_path.data(), all_tokens.data(), all_tokens.size()); - LOG_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size()); + COM_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size()); llama_token last_token = all_tokens.back(); llama_batch batch = llama_batch_get_one(&last_token, 1); @@ -2030,13 +2030,13 @@ bool common_prompt_batch_decode( batch.pos = &pos; if (llama_decode(ctx, batch)) { - LOG_ERR("%s : failed to eval last token\n", __func__); + COM_ERR("%s", "failed to eval last token\n"); return false; } n_past++; } else { if (llama_decode(ctx, llama_batch_get_one(const_cast(all_tokens.data() + offset), n_new))) { - LOG_ERR("%s : failed to eval\n", __func__); + COM_ERR("%s", "failed to eval\n"); return false; } n_past += n_new; diff --git a/common/common.h b/common/common.h index 94147d5d8..d56f6064b 100644 --- a/common/common.h +++ b/common/common.h @@ -25,6 +25,13 @@ #define DIRECTORY_SEPARATOR '/' #endif // _WIN32 +#define COM_DBG(fmt, ...) LOG_DBG("cmn %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define COM_TRC(fmt, ...) LOG_TRC("cmn %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define COM_INF(fmt, ...) LOG_INF("cmn %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define COM_WRN(fmt, ...) LOG_WRN("cmn %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define COM_ERR(fmt, ...) LOG_ERR("cmn %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define COM_CNT(fmt, ...) LOG_CNT("" fmt, __VA_ARGS__) + #define die(msg) do { fputs("error: " msg "\n", stderr); exit(1); } while (0) #define die_fmt(fmt, ...) do { fprintf(stderr, "error: " fmt "\n", __VA_ARGS__); exit(1); } while (0) diff --git a/common/fit.cpp b/common/fit.cpp index a8565bfc9..afbf0b10f 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -233,7 +233,7 @@ static void common_params_fit_impl( sum_projected_used = dmds_full.back().mb.total(); sum_free = dmds_full.back().total; sum_projected_free = sum_free - sum_projected_used; - LOG_INF("%s: projected to use %" PRId64 " MiB of host memory vs. %" PRId64 " MiB of total host memory\n", + LOG_TRC("%s: projected to use %" PRId64 " MiB of host memory vs. %" PRId64 " MiB of total host memory\n", __func__, sum_projected_used/MiB, sum_free/MiB); if (sum_projected_free >= margins[0]) { LOG_TRC("%s: will leave %" PRId64 " >= %" PRId64 " MiB of system memory, no changes needed\n", diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index ce41d029b..7da0bb1c5 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -65,12 +65,12 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to if (ctx->start_matcher.advance(token)) { ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; - LOG_INF("reasoning-budget: activated, budget=%d tokens\n", ctx->budget); + COM_TRC("activated, budget=%d tokens\n", ctx->budget); if (ctx->remaining <= 0) { ctx->state = REASONING_BUDGET_FORCING; ctx->force_pos = 0; - LOG_INF("reasoning-budget: budget=0, forcing immediately\n"); + COM_TRC("%s", "budget=0, forcing immediately\n"); } } break; @@ -80,7 +80,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to { if (ctx->end_matcher.advance(token)) { ctx->state = REASONING_BUDGET_DONE; - LOG_INF("reasoning-budget: deactivated (natural end)\n"); + COM_TRC("%s", "deactivated (natural end)\n"); break; } @@ -95,7 +95,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to ctx->state = REASONING_BUDGET_FORCING; ctx->force_pos = 0; ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: UTF-8 complete, now forcing end sequence\n"); + COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n"); } } else if (ctx->state == REASONING_BUDGET_COUNTING) { ctx->remaining--; @@ -104,11 +104,11 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to ctx->state = REASONING_BUDGET_FORCING; ctx->force_pos = 0; ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: budget exhausted, forcing end sequence\n"); + COM_TRC("%s", "budget exhausted, forcing end sequence\n"); } else { ctx->state = REASONING_BUDGET_WAITING_UTF8; ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: budget exhausted, waiting for UTF-8 completion\n"); + COM_TRC("%s", "budget exhausted, waiting for UTF-8 completion\n"); } } } @@ -118,7 +118,7 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to ctx->force_pos++; if (ctx->force_pos >= ctx->forced_tokens.size()) { ctx->state = REASONING_BUDGET_DONE; - LOG_INF("reasoning-budget: forced sequence complete, done\n"); + COM_TRC("%s", "forced sequence complete, done\n"); } break; case REASONING_BUDGET_DONE: @@ -128,12 +128,12 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to ctx->state = REASONING_BUDGET_COUNTING; ctx->remaining = ctx->budget; ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: re-activated on new start tag, budget=%d tokens\n", ctx->budget); + COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget); if (ctx->remaining <= 0) { ctx->state = REASONING_BUDGET_FORCING; ctx->force_pos = 0; - LOG_INF("reasoning-budget: budget=0, forcing immediately\n"); + COM_TRC("%s", "budget=0, forcing immediately\n"); } } break; @@ -264,7 +264,7 @@ bool common_reasoning_budget_force(struct llama_sampler * smpl) { ctx->state = REASONING_BUDGET_FORCING; ctx->force_pos = 0; ctx->end_matcher.reset(); - LOG_INF("reasoning-budget: forced into forcing state (manual transition)\n"); + COM_TRC("%s", "forced into forcing state (manual transition)\n"); return true; } diff --git a/common/speculative.cpp b/common/speculative.cpp index c922a3f59..a3495c3a1 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -18,6 +18,13 @@ #include #include +#define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define SPC_TRC(fmt, ...) LOG_TRC("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define SPC_INF(fmt, ...) LOG_INF("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define SPC_WRN(fmt, ...) LOG_WRN("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define SPC_ERR(fmt, ...) LOG_ERR("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__) +#define SPC_CNT(fmt, ...) LOG_CNT("" fmt, __VA_ARGS__) + #define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 @@ -60,21 +67,20 @@ static bool common_speculative_are_compatible( const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); const auto vocab_type_tgt = llama_vocab_type(vocab_tgt); - LOG_DBG("%s: vocab_type tgt: %d\n", __func__, vocab_type_tgt); + SPC_DBG("vocab_type tgt: %d\n", vocab_type_tgt); const auto vocab_type_dft = llama_vocab_type(vocab_dft); - LOG_DBG("%s: vocab_type dft: %d\n", __func__, vocab_type_dft); + SPC_DBG("vocab_type dft: %d\n", vocab_type_dft); if (vocab_type_tgt != vocab_type_dft) { - LOG_WRN("%s: draft model vocab type must match target model to use speculation but " - "vocab_type_dft = %d while vocab_type_tgt = %d\n", __func__, vocab_type_dft, vocab_type_tgt); + SPC_WRN("draft model vocab type must match target model to use speculation but " + "vocab_type_dft = %d while vocab_type_tgt = %d\n", vocab_type_dft, vocab_type_tgt); return false; } if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || (llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) { - LOG_WRN("%s: draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", - __func__, + SPC_WRN("draft model bos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft), llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft)); return false; @@ -82,8 +88,7 @@ static bool common_speculative_are_compatible( if (llama_vocab_get_add_eos(vocab_tgt) != llama_vocab_get_add_eos(vocab_dft) || (llama_vocab_get_add_eos(vocab_tgt) && llama_vocab_eos(vocab_tgt) != llama_vocab_eos(vocab_dft))) { - LOG_WRN("%s: draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", - __func__, + SPC_WRN("draft model eos tokens must match target model to use speculation. add: %d - %d, id: %d - %d)\n", llama_vocab_get_add_eos(vocab_tgt), llama_vocab_get_add_eos(vocab_dft), llama_vocab_eos(vocab_tgt), llama_vocab_eos(vocab_dft)); return false; @@ -97,8 +102,8 @@ static bool common_speculative_are_compatible( : n_vocab_dft - n_vocab_tgt; if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { - LOG_DBG("%s: draft model vocab must closely match target model to use speculation but ", __func__); - LOG_DBG("target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", + SPC_DBG("draft model vocab must closely match target model to use speculation but " + "target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", n_vocab_tgt, llama_vocab_n_tokens(vocab_dft), vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); return false; } @@ -108,8 +113,8 @@ static bool common_speculative_are_compatible( const char * token_text_dft = llama_vocab_get_text(vocab_dft, i); if (std::strcmp(token_text_tgt, token_text_dft) != 0) { - LOG_DBG("%s: draft model vocab must match target model to use speculation but ", __func__); - LOG_DBG("token %d content differs - target '%s', draft '%s'\n", i, + SPC_DBG("draft model vocab must match target model to use speculation but " + "token %d content differs - target '%s', draft '%s'\n", i, common_token_to_piece(vocab_tgt, i).c_str(), common_token_to_piece(vocab_dft, i).c_str()); return false; @@ -186,9 +191,9 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { auto * ctx_dft = this->params.ctx_dft; auto * ctx_tgt = this->params.ctx_tgt; - LOG_INF("%s: adding speculative implementation 'draft-simple'\n", __func__); - LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); - LOG_INF("%s: - gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", __func__, + SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n"); + SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min); + SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", this->params.n_gpu_layers, ggml_type_name(this->params.cache_type_k), ggml_type_name(this->params.cache_type_v), @@ -228,16 +233,16 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { } const bool vocab_cmpt = common_speculative_are_compatible(llama_get_model(ctx_tgt), llama_get_model(ctx_dft)); - LOG_DBG("%s: vocab_cmpt = %d\n", __func__, vocab_cmpt); + SPC_DBG("vocab_cmpt = %d\n", vocab_cmpt); if (!vocab_cmpt) { - LOG_ERR("%s: the target and draft vocabs are not compatible\n", __func__); + SPC_ERR("%s", "the target and draft vocabs are not compatible\n"); throw std::runtime_error("draft model vocab type must match target model to use speculation"); } if (n_seq != llama_n_seq_max(ctx_dft)) { - LOG_ERR("%s: n_seq mismatch: %d != %d\n", __func__, n_seq, llama_n_seq_max(ctx_dft)); + SPC_ERR("n_seq mismatch: %d != %d\n", n_seq, llama_n_seq_max(ctx_dft)); throw std::runtime_error("the draft model number of sequences is incompatible with the speculative n_seq"); } @@ -257,7 +262,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { const int ret = llama_decode(ctx_dft, batch); if (ret != 0) { - LOG_ERR("%s: failed to decode draft batch, ret = %d\n", __func__, ret); + SPC_ERR("failed to decode draft batch, ret = %d\n", ret); return false; } @@ -290,7 +295,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { int ret = llama_decode(ctx_dft, batch); if (ret != 0) { - LOG_WRN("%s: llama_decode returned %d\n", __func__, ret); + SPC_ERR("llama_decode returned %d\n", ret); return; } @@ -314,7 +319,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { const auto * cur_p = common_sampler_get_candidates(smpl, true); for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + SPC_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); } @@ -354,7 +359,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { // evaluate the drafted tokens on the draft model ret = llama_decode(ctx_dft, batch); if (ret != 0) { - LOG_WRN("%s: llama_decode[%d] returned %d\n", __func__, i, ret); + SPC_ERR("llama_decode[%d] returned %d\n", i, ret); break; } @@ -449,8 +454,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq) , params(params.draft) { - LOG_INF("%s: adding speculative implementation 'draft-eagle3'\n", __func__); - LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%f, backend_sampling=%d\n", __func__, params.draft.n_max, params.draft.n_min, params.draft.p_min, (int) params.draft.backend_sampling); + SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n"); + SPC_TRC("- n_max=%d, n_min=%d, p_min=%f, backend_sampling=%d\n", params.draft.n_max, params.draft.n_min, params.draft.p_min, (int) params.draft.backend_sampling); auto * ctx_tgt = this->params.ctx_tgt; auto * ctx_dft = this->params.ctx_dft; @@ -493,7 +498,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { llama_sampler_chain_add(chain, llama_sampler_init_top_k(10)); if (!llama_set_sampler(ctx_dft, seq_id, chain)) { - LOG_WRN("%s: backend offload failed for seq_id=%d; using CPU sampler\n", __func__, (int) seq_id); + SPC_WRN("backend offload failed for seq_id=%d; using CPU sampler\n", (int) seq_id); llama_sampler_free(chain); chain = nullptr; } @@ -548,9 +553,9 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { auto * ctx_dft = this->params.ctx_dft; const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id); if (pos_max < N - 2) { - LOG_WRN("%s: ctx_dft pos_max=%d < N-2=%d — process() did not run on every prefill ubatch. " + SPC_WRN("ctx_dft pos_max=%d < N-2=%d — process() did not run on every prefill ubatch. " "Drafts may degrade.\n", - __func__, (int) pos_max, N - 2); + (int) pos_max, N - 2); } } @@ -621,8 +626,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { }; const int32_t rc = llama_encode(ctx_dft, enc_batch); if (rc != 0) { - LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n", - __func__, rc, (int) n_chunk, (int) i); + SPC_ERR("llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n", + rc, (int) n_chunk, (int) i); return false; } @@ -692,8 +697,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { if (batch.n_tokens > 0) { const int32_t rc = llama_decode(ctx_dft, batch); if (rc != 0) { - LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, ubatch_pos[0]=%d)\n", - __func__, rc, (int) batch.n_tokens, (int) batch_in.pos[0]); + SPC_ERR("llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, ubatch_pos[0]=%d)\n", + rc, (int) batch.n_tokens, (int) batch_in.pos[0]); return false; } } @@ -744,7 +749,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { int ret = llama_decode(ctx_dft, batch); if (ret != 0) { - LOG_WRN("%s: llama_decode returned %d\n", __func__, ret); + SPC_ERR("llama_decode returned %d\n", ret); return; } @@ -770,7 +775,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { const auto * cur_p = common_sampler_get_candidates(smpl, true); for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + SPC_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); } @@ -809,7 +814,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { ret = llama_decode(ctx_dft, batch); if (ret != 0) { - LOG_WRN("%s: llama_decode[%d] returned %d\n", __func__, i, ret); + SPC_ERR("llama_decode[%d] returned %d\n", i, ret); break; } @@ -942,9 +947,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { "MTP input row width must match the target h_nextn width"); n_mtp_layers = std::max(1, (int) llama_model_n_layer_nextn(llama_get_model(ctx_dft))); - LOG_INF("%s: adding speculative implementation 'draft-mtp'\n", __func__); - LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling); - LOG_INF("%s: - gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", __func__, + SPC_TRC("%s", "adding speculative implementation 'draft-mtp'\n"); + SPC_TRC("- n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling); + SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", this->params.n_gpu_layers, ggml_type_name(this->params.cache_type_k), ggml_type_name(this->params.cache_type_v), @@ -975,7 +980,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { llama_sampler_chain_add(chain, llama_sampler_init_top_k(10)); if (!llama_set_sampler(ctx_dft, seq_id, chain)) { - LOG_WRN("%s: backend offload failed for seq_id=%d; using CPU sampler\n", __func__, (int) seq_id); + SPC_WRN("backend offload failed for seq_id=%d; using CPU sampler\n", (int) seq_id); llama_sampler_free(chain); chain = nullptr; } @@ -1038,11 +1043,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id); if (pos_max < N - 1 && !is_mem_shared) { - LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - " + SPC_WRN("ctx_dft pos_max=%d < N-1=%d - " "process() hook may not have run on every prefill ubatch " "(need_embd / logits=1 on every prompt position?). " "Drafts may degrade.\n", - __func__, (int) pos_max, N - 1); + (int) pos_max, N - 1); } } @@ -1128,8 +1133,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const int32_t rc = llama_decode(ctx_dft, batch); if (rc != 0) { - LOG_ERR("%s: llama_decode(ctx_dft) head=%d failed rc=%d (pos=%d)\n", - __func__, head, (int) rc, (int) batch_in.pos[0]); + SPC_ERR("llama_decode(ctx_dft) head=%d failed rc=%d (pos=%d)\n", + head, (int) rc, (int) batch_in.pos[0]); ok = false; break; } @@ -1217,7 +1222,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { int ret = llama_decode(ctx_dft, batch); if (ret != 0) { - LOG_WRN("%s: llama_decode[%d] returned %d\n", __func__, i, ret); + SPC_ERR("llama_decode[%d] returned %d\n", i, ret); break; } @@ -1239,7 +1244,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { const auto * cur_p = common_sampler_get_candidates(smpl, true); for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { - LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + SPC_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p, common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); } @@ -1353,8 +1358,8 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl { , params(params.ngram_simple) , config(config) { - LOG_INF("%s: adding speculative implementation 'ngram-simple'\n", __func__); - LOG_INF("%s: - size_n=%d, size_m=%d, min_hits=%d\n", __func__, + SPC_TRC("%s", "adding speculative implementation 'ngram-simple'\n"); + SPC_TRC("- size_n=%d, size_m=%d, min_hits=%d\n", this->params.size_n, this->params.size_m, this->params.min_hits); } @@ -1403,8 +1408,8 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl { this->config.push_back(config); } - LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(this->type).c_str()); - LOG_INF("%s: - size_key=%d, size_value=%d, key_only=%d, min_hits=%d\n", __func__, + SPC_TRC("adding speculative implementation '%s'\n", common_speculative_type_to_str(this->type).c_str()); + SPC_TRC("- size_key=%d, size_value=%d, key_only=%d, min_hits=%d\n", config.size_key, config.size_value, config.key_only, config.min_hits); } @@ -1478,15 +1483,15 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { , verbose(std::getenv("LLAMA_TRACE") != nullptr) { static_assert(sizeof(llama_token) == sizeof(common_ngram_mod::entry_t)); - LOG_INF("%s: adding speculative implementation 'ngram-mod'\n", __func__); - LOG_INF("%s: - n_match=%d, n_max=%d, n_min=%d\n", __func__, + SPC_TRC("%s", "adding speculative implementation 'ngram-mod'\n"); + SPC_TRC("- n_match=%d, n_max=%d, n_min=%d\n", this->params.n_match, this->params.n_max, this->params.n_min); - LOG_INF("%s: - mod size=%zu (%.3f MB)\n", __func__, + SPC_TRC("- mod size=%zu (%.3f MB)\n", mod.size(), (float)(mod.size_bytes())/1024/1024); if (this->params.n_match < 16) { - LOG_WRN("%s: ngram_mod n_match=%d is too small - poor quality is possible, " - "see: https://github.com/ggml-org/llama.cpp/pull/19164\n", __func__, this->params.n_match); + SPC_WRN("ngram_mod n_match=%d is too small - poor quality is possible, " + "see: https://github.com/ggml-org/llama.cpp/pull/19164\n", this->params.n_match); } sinfos.resize(n_seq); @@ -1510,11 +1515,11 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { sinfo.i_last = prompt.size() - n; const double f = (double)mod.get_used() / (double)mod.size(); - LOG_INF("%s: ngram_mod occupancy = %zu/%zu (%.2f)\n", __func__, mod.get_used(), mod.size(), f); + SPC_TRC("ngram_mod occupancy = %zu/%zu (%.2f)\n", mod.get_used(), mod.size(), f); constexpr double f_thold = 0.25; if (f > f_thold) { - LOG_WRN("%s: ngram_mod occupancy %.2f exceeds threshold (%.2f) - resetting\n", __func__, f, f_thold); + SPC_WRN("ngram_mod occupancy %.2f exceeds threshold (%.2f) - resetting\n", f, f_thold); mod.reset(); } @@ -1608,7 +1613,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { sinfo.n_low++; if (sinfo.n_low >= 5) { if (verbose) { - LOG_WRN("%s: low acceptance streak (%d) - resetting ngram_mod\n", __func__, sinfo.n_low); + SPC_TRC("low acceptance streak (%d) - resetting ngram_mod\n", sinfo.n_low); } mod.reset(); @@ -1658,8 +1663,8 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { , save_dynamic(save_dynamic) , save_static(save_static) { - LOG_INF("%s: adding speculative implementation 'ngram-cache'\n", __func__); - LOG_INF("%s: - n_draft=%d, cache_static=%s, cache_dynamic=%s\n", __func__, + SPC_TRC("%s", "adding speculative implementation 'ngram-cache'\n"); + SPC_TRC("- n_draft=%d, cache_static=%s, cache_dynamic=%s\n", n_draft, path_static.empty() ? "none" : path_static.c_str(), path_dynamic.empty() ? "none" : path_dynamic.c_str()); @@ -1674,7 +1679,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { sinfo.ngram_cache_static = ngram_cache_static; } } catch (...) { - LOG_ERR("failed to open static lookup cache: %s", path_static.c_str()); + SPC_ERR("failed to open static lookup cache: %s", path_static.c_str()); GGML_ABORT("Couldn't read static lookup cache"); } } @@ -1687,7 +1692,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { sinfo.ngram_cache_dynamic = ngram_cache_dynamic; } } catch (...) { - LOG_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str()); + SPC_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str()); GGML_ABORT("Couldn't read dynamic lookup cache"); } } @@ -2034,7 +2039,7 @@ common_speculative * common_speculative_init(common_params_speculative & params, } if (impls.empty()) { - LOG_WRN("%s: no implementations specified for speculative decoding\n", __func__); + SPC_TRC("%s", "no implementations specified for speculative decoding\n"); return nullptr; } @@ -2161,13 +2166,13 @@ void common_speculative_draft(common_speculative * spec) { if (dp.n_max > 0) { if (!result.empty() && (int) result.size() > dp.n_max) { - LOG_DBG("%s: truncating draft to %d tokens\n", __func__, dp.n_max); + SPC_DBG("truncating draft to %d tokens\n", dp.n_max); result.resize(dp.n_max); } } if (!result.empty()) { - LOG_DBG("%s: called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n", __func__, + SPC_DBG("called impl %s, hist size = %zu, call_count = %zu, gen = %zu\n", common_speculative_type_to_str(impl.get()->type).c_str(), dp.prompt->size(), impl.get()->n_call_draft, result.size()); @@ -2291,7 +2296,7 @@ void common_speculative_print_stats(const common_speculative * spec) { str_stats = ", #mean acc len = " + oss.str() + ", #acc rate/pos = (" + tmp.str() + ")"; } - LOG_INF("statistics %16s: #calls(b,g,a) = %4zu %6zu %6zu, #gen drafts = %6zu, #acc drafts = %5zu, #gen tokens = %6zu, #acc tokens = %5zu%s%s\n", + SPC_TRC("statistics %16s: #calls(b,g,a) = %4zu %6zu %6zu, #gen drafts = %6zu, #acc drafts = %5zu, #gen tokens = %6zu, #acc tokens = %5zu%s%s\n", common_speculative_type_to_str(impl->type).c_str(), impl->n_call_begin, impl->n_call_draft, impl->n_call_accept, impl->n_gen_drafts, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 220240ea9..9f8a8fdb8 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -256,7 +256,7 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max); if (cparams.n_ctx_seq < hparams.n_ctx_train) { - LLAMA_LOG_WARN("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n", + LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n", __func__, cparams.n_ctx_seq, hparams.n_ctx_train); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5c33a418f..bb1c236cb 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -106,7 +106,6 @@ struct server_batch { if ((int32_t)tokens.size() >= n_tokens_alloc) { return false; } - // LOG_INF("adding token to batch: slot=%d, token=%d, pos=%d, output=%d\n", id_slot, token, pos, output); tokens.push_back({ id_slot, token, pos, output }); return true; } @@ -228,7 +227,7 @@ struct server_slot { const size_t cur_size = cur_size_tgt + cur_size_dft; - SRV_WRN(" - saving prompt with length %d, total state size = %.3f MiB (draft: %.3f MiB)\n", + SRV_TRC(" - saving prompt with length %d, total state size = %.3f MiB (draft: %.3f MiB)\n", (int) prompt.tokens.size(), cur_size / (1024.0 * 1024.0), cur_size_dft / (1024.0 * 1024.0)); auto * cur = prompt_cache.alloc(prompt, cur_size_tgt, cur_size_dft); @@ -258,7 +257,7 @@ struct server_slot { GGML_ASSERT(!is_processing()); } - SLT_INF(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); + SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); common_context_seq_rm(ctx_tgt, id, -1, -1); if (ctx_dft) { @@ -627,8 +626,10 @@ struct server_slot { } SLT_INF(*this, - "draft acceptance = %0.5f (%5d accepted / %5d generated), mean acceptance length = %5.2f, acceptance rate per position = (%s)\n", - draft_ratio, n_draft_accepted, n_draft_total, mean_acc_len, acceptance_rates_per_pos.c_str()); + "draft acceptance = %0.5f (%5d accepted / %5d generated), mean len = %5.2f\n", + draft_ratio, n_draft_accepted, n_draft_total, mean_acc_len); + SLT_TRC(*this, + " acc per pos = (%s)\n", acceptance_rates_per_pos.c_str()); } common_speculative_print_stats(spec); @@ -771,7 +772,7 @@ struct server_slot { } // TODO @ngxson : move this log line to debug when it become more stable - SLT_INF(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added); + SLT_TRC(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added); res = mtmd_batch_encode(mbatch.get()); if (res != 0) { @@ -1032,7 +1033,8 @@ private: } - SRV_INF("loading model '%s'\n", params.model.path.c_str()); + SRV_INF("loading model '%s'\n", params.model.get_name().c_str()); + SRV_TRC("local path '%s'\n", params.model.path.c_str()); std::string & mmproj_path = params_base.mmproj.path; mtmd_context_params mparams = mtmd_context_params_default(); @@ -1061,7 +1063,7 @@ private: for (auto & [dev, size] : mmproj_mem) { total += size; } - SRV_INF("[mtmd] estimated worst-case memory usage of mmproj is %.2f MiB (took %.2f ms)\n", total / (1024.0 * 1024.0), t_elapsed / 1000.0); + SRV_TRC("[mtmd] estimated worst-case memory usage of mmproj is %.2f MiB (took %.2f ms)\n", total / (1024.0 * 1024.0), t_elapsed / 1000.0); GGML_ASSERT(!params_base.fit_params_target.empty()); for (auto & [dev, size] : mmproj_mem) { for (size_t i = 0; i < ggml_backend_dev_count(); i++) { @@ -1141,7 +1143,7 @@ private: } } } - SRV_INF("[spec] estimated memory usage of %s is %.2f MiB\n", + SRV_TRC("[spec] estimated memory usage of %s is %.2f MiB\n", has_draft ? "draft model" : "MTP context", total / (1024.0 * 1024.0)); } catch (const std::exception & e) { @@ -1177,7 +1179,7 @@ private: // TODO speculative: move to common/speculative.cpp? const auto & params_spec = params_base.speculative.draft; - SRV_INF("loading draft model '%s'\n", params_spec.mparams.path.c_str()); + SRV_TRC("loading draft model '%s'\n", params_spec.mparams.path.c_str()); auto params_dft = params_base; @@ -1229,7 +1231,7 @@ private: // no new model load, so we simply report 0.0 and 1.0 progress load_progress_callback(0.0f, &load_progress_spec); - SRV_INF("creating MTP draft context against the target model '%s'\n", + SRV_TRC("creating MTP draft context against the target model '%s'\n", params_base.model.path.c_str()); auto cparams_mtp = common_context_params_to_llama(params_base); @@ -1303,9 +1305,6 @@ private: // Necessary similarity of prompt for slot selection slot_prompt_similarity = params_base.slot_prompt_similarity; - // setup slots - SRV_INF("initializing slots, n_slots = %d\n", params_base.n_parallel); - const int n_ctx_train = llama_model_n_ctx_train(model_tgt); int n_ctx_slot = llama_n_ctx_seq(ctx_tgt); @@ -1322,9 +1321,13 @@ private: } if (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL) { - SRV_WRN("%s", "speculative decoding will use checkpoints\n"); + SRV_TRC("%s", "speculative decoding will use checkpoints\n"); } + // 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"); + // initialize slots for (int i = 0; i < params_base.n_parallel; i++) { slots.emplace_back(); @@ -1344,7 +1347,7 @@ private: } if (spec) { - SRV_INF("%s", "speculative decoding context initialized\n"); + SRV_TRC("%s", "speculative decoding context initialized\n"); } else { ctx_dft.reset(); } @@ -1361,7 +1364,7 @@ private: slot.mctx = mctx; slot.prompt.tokens.has_mtmd = mctx != nullptr; - SLT_INF(slot, "new slot, n_ctx = %d\n", slot.n_ctx); + SLT_TRC(slot, "new slot, n_ctx = %d\n", slot.n_ctx); slot.callback_on_release = [this](int id_slot) { queue_tasks.pop_deferred_task(id_slot); @@ -1397,23 +1400,23 @@ private: if (params_base.cache_ram_mib != 0) { if (params_base.cache_ram_mib < 0) { - SRV_INF("prompt cache is enabled, size limit: %s\n", "no limit"); + SRV_TRC("prompt cache is enabled, size limit: %s\n", "no limit"); } else { - SRV_INF("prompt cache is enabled, size limit: %d MiB\n", params_base.cache_ram_mib); + SRV_TRC("prompt cache is enabled, size limit: %d MiB\n", params_base.cache_ram_mib); } - SRV_INF("%s", "use `--cache-ram 0` to disable the prompt cache\n"); + SRV_TRC("%s", "use `--cache-ram 0` to disable the prompt cache\n"); prompt_cache = std::make_unique(params_base.cache_ram_mib, n_ctx); } else { - SRV_INF("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n"); + SRV_TRC("%s", "prompt cache is disabled - use `--cache-ram N` to enable it\n"); } - SRV_INF("%s", "for more info see https://github.com/ggml-org/llama.cpp/pull/16391\n"); + SRV_TRC("%s", "for more info see https://github.com/ggml-org/llama.cpp/pull/16391\n"); if (params_base.n_ctx_checkpoints > 0) { - SRV_INF("context checkpoints enabled, max = %d, min spacing = %d\n", + SRV_TRC("context checkpoints enabled, max = %d, min spacing = %d\n", params_base.n_ctx_checkpoints, params_base.checkpoint_min_step); } else { - SRV_INF("%s", "context checkpoints disabled\n"); + SRV_TRC("%s", "context checkpoints disabled\n"); } if (!params_base.model_alias.empty()) { @@ -1470,11 +1473,11 @@ private: params_base.cache_idle_slots = false; } else { if (params_base.kv_unified) { - SRV_INF("%s", "idle slots will be saved to prompt cache and cleared upon starting a new task\n"); + SRV_TRC("%s", "idle slots will be saved to prompt cache and cleared upon starting a new task\n"); } else { // without a unified KV cache, clearing a slot frees no reusable room, so we only // publish a RAM-cache copy of idle slots (their KV stays in VRAM) [TAG_IDLE_SLOT_CLEAR] - SRV_INF("%s", "idle slots will be saved to prompt cache upon starting a new task\n"); + SRV_TRC("%s", "idle slots will be saved to prompt cache upon starting a new task\n"); } SRV_DBG("%s", "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__\n"); } @@ -1500,7 +1503,7 @@ private: try { chat_templates = common_chat_templates_init(model_tgt, params_base.chat_template); - LOG_INF("%s: chat template, example_format: '%s'\n", __func__, + SRV_TRC("%s: chat template, example_format: '%s'\n", __func__, common_chat_format_example(chat_templates.get(), params_base.use_jinja, params_base.default_template_kwargs).c_str()); } catch (const std::exception & e) { @@ -1515,7 +1518,7 @@ private: // 2. The chat template supports it const bool template_supports_thinking = params_base.use_jinja && common_chat_templates_support_enable_thinking(chat_templates.get()); const bool enable_thinking = params_base.enable_reasoning != 0 && template_supports_thinking; - SRV_INF("%s: chat template, thinking = %d\n", __func__, enable_thinking); + SRV_TRC("%s: chat template, thinking = %d\n", __func__, enable_thinking); // IMPORTANT: chat_params is reused across sleeping / resuming states, // never store llama_context/llama_model pointers in chat_params, @@ -1658,7 +1661,7 @@ private: update_cache = update_cache && task.type == SERVER_TASK_TYPE_COMPLETION; if (update_cache) { - SRV_INF("%s", "updating prompt cache\n"); + SRV_TRC("%s", "updating prompt cache\n"); const int64_t t_start = ggml_time_us(); @@ -1670,7 +1673,7 @@ private: prompt_cache->update(); - SRV_INF("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0); + SRV_TRC("prompt cache update took %.2f ms\n", (ggml_time_us() - t_start) / 1000.0); } } @@ -2290,7 +2293,7 @@ private: int id_parent = parent_task.id; - SRV_INF("launching slots for parent task id_task = %d with %zu child tasks\n", id_parent, parent_task.child_tasks.size()); + SRV_TRC("launching slots for parent task id_task = %d with %zu child tasks\n", id_parent, parent_task.child_tasks.size()); // to be called in case of failure to release all launched slots auto release_slots = [this, id_parent]() { @@ -2351,7 +2354,7 @@ private: // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.id, cur.data_spec); - SLT_INF(slot, + SLT_TRC(slot, "created context checkpoint %d of %d (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", size = %.3f MiB)\n", (int) slot.prompt.checkpoints.size(), params_base.n_ctx_checkpoints, cur.pos_min, cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024); @@ -2415,7 +2418,7 @@ private: if (params_base.cache_idle_slots) { for (auto & slot : slots) { if (!slot.is_processing()) { - SLT_INF(slot, "%s", "saving idle slot to prompt cache\n"); + SLT_TRC(slot, "%s", "saving idle slot to prompt cache\n"); if (slot.prompt_save(*prompt_cache)) { SLT_DBG(slot, "%s", "__TEST_TAG_CACHE_IDLE_SLOT__\n"); @@ -2671,7 +2674,7 @@ private: auto new_loras = construct_lora_list(task.set_lora); // logging for (size_t i = 0; i < new_loras.size(); ++i) { - SRV_INF("set lora adapter idx=%zu scale=%f\n", i, new_loras[i].scale); + SRV_TRC("set lora adapter idx=%zu scale=%f\n", i, new_loras[i].scale); } // TODO @ngxson : make lora_adapters a dedicated member of server_context params_base.lora_adapters = new_loras; @@ -2771,7 +2774,7 @@ private: } if (all_idle) { - SRV_INF("%s", "all slots are idle\n"); + SRV_TRC("%s", "all slots are idle\n"); return; // skip further processing } else { @@ -3287,10 +3290,9 @@ private: const auto it = std::find_if( slot.prompt.checkpoints.rbegin(), slot.prompt.checkpoints.rend(), - [&, func_name = __func__](const auto & cur) { + [&](const auto & cur) { // guarantee that a checkpoint will result in at least one token being processed [TAG_PROMPT_LOGITS] - LOG_INF("slot %12.*s: id %2d | task %d | Checking checkpoint with [%d, %d] against %d...\n", 12, - func_name, (slot).id, ((slot).task ? (slot).task->id : -1), cur.pos_min, cur.pos_max, pos_min_thold); + SLT_TRC(slot, "checking checkpoint with [%d, %d] against %d...\n", cur.pos_min, cur.pos_max, pos_min_thold); // workaround for [TAG_CHECKPOINTS_FIX_POS_MIN] if (cur.pos_max > pos_next) { return false; @@ -3310,11 +3312,11 @@ private: pos_next = std::min(pos_next, std::max(it->pos_min + 1, it->pos_max)); n_past = std::min(slot.prompt.tokens.size_up_to_pos(pos_next), (size_t) it->n_tokens); - SLT_WRN(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, n_past, (float) it->size() / 1024 / 1024); + SLT_TRC(slot, "restored context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_past = %d, size = %.3f MiB)\n", it->pos_min, it->pos_max, it->n_tokens, n_past, (float) it->size() / 1024 / 1024); } if (do_reset) { - SLT_WRN(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n", + SLT_TRC(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA or hybrid/recurrent memory, see %s)\n", "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055"); pos_next = 0; n_past = 0; @@ -3327,7 +3329,7 @@ private: for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end();) { const auto & cur = *it; if (cur.pos_max > pos_next) { - SLT_WRN(slot, "erased invalidated context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_swa = %d, pos_next = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_swa, pos_next, (float) cur.size() / 1024 / 1024); + SLT_TRC(slot, "erased invalidated context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_swa = %d, pos_next = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_swa, pos_next, (float) cur.size() / 1024 / 1024); it = slot.prompt.checkpoints.erase(it); } else { ++it; @@ -3674,7 +3676,7 @@ private: // all children slots should already launched by launch_slots_with_parent_task() // copy state to the child slots for (auto & child : children) { - SLT_INF(slot, " - copying state to child %d\n", child->id); + SLT_TRC(slot, " - copying state to child %d\n", child->id); GGML_ASSERT(child->state == SLOT_STATE_WAIT_OTHER); diff --git a/tools/server/server-http.cpp b/tools/server/server-http.cpp index 82f34edac..21bed64c9 100644 --- a/tools/server/server-http.cpp +++ b/tools/server/server-http.cpp @@ -83,7 +83,7 @@ bool server_http_context::init(const common_params & params) { hostname = params.hostname; if (gcp.enabled) { - SRV_INF("Google Cloud Platform compat: health route = %s, predict route = %s, port = %d\n", gcp.path_health.c_str(), gcp.path_predict.c_str(), gcp.port); + SRV_TRC("Google Cloud Platform compat: health route = %s, predict route = %s, port = %d\n", gcp.path_health.c_str(), gcp.path_predict.c_str(), gcp.port); if (port != gcp.port) { SRV_WRN("Google Cloud Platform compat: overriding server port %d with AIP_HTTP_PORT %d\n", port, gcp.port); @@ -96,13 +96,13 @@ bool server_http_context::init(const common_params & params) { #ifdef CPPHTTPLIB_OPENSSL_SUPPORT if (!params.ssl_file_key.empty() && !params.ssl_file_cert.empty()) { - SRV_INF("running with SSL: key = %s, cert = %s\n", params.ssl_file_key.c_str(), params.ssl_file_cert.c_str()); + SRV_TRC("running with SSL: key = %s, cert = %s\n", params.ssl_file_key.c_str(), params.ssl_file_cert.c_str()); srv = std::make_unique( params.ssl_file_cert.c_str(), params.ssl_file_key.c_str() ); is_ssl = true; } else { - SRV_INF("%s", "running without SSL\n"); + SRV_TRC("%s", "running without SSL\n"); srv = std::make_unique(); } #else @@ -165,9 +165,9 @@ bool server_http_context::init(const common_params & params) { if (params.api_keys.size() == 1) { const auto key = params.api_keys[0]; const std::string substr = key.substr(std::max(static_cast(key.length() - 4), 0)); - SRV_INF("api_keys: ****%s\n", substr.c_str()); + SRV_TRC("api_keys: ****%s\n", substr.c_str()); } else if (params.api_keys.size() > 1) { - SRV_INF("api_keys: %zu keys loaded\n", params.api_keys.size()); + SRV_TRC("api_keys: %zu keys loaded\n", params.api_keys.size()); } // @@ -293,7 +293,7 @@ bool server_http_context::init(const common_params & params) { // +4 threads for monitoring, health and some threads reserved for MCP and other tasks in the future n_threads_http = std::max(params.n_parallel + 4, static_cast(std::thread::hardware_concurrency() - 1)); } - SRV_INF("using %d threads for HTTP server\n", n_threads_http); + SRV_TRC("using %d threads for HTTP server\n", n_threads_http); srv->new_task_queue = [n_threads_http] { // spawn n_threads_http fixed thread (always alive), while allow up to 1024 max possible additional threads // when n_threads_http is used, server will create new "dynamic" threads that will be destroyed after processing each request @@ -412,13 +412,13 @@ bool server_http_context::start() { auto is_sock = false; if (string_ends_with(std::string(hostname), ".sock")) { is_sock = true; - SRV_INF("%s", "setting address family to AF_UNIX\n"); + SRV_TRC("%s", "setting address family to AF_UNIX\n"); srv->set_address_family(AF_UNIX); // bind_to_port requires a second arg, any value other than 0 should // simply get ignored was_bound = srv->bind_to_port(hostname, 8080); } else { - SRV_INF("%s", "binding port with default address family\n"); + SRV_TRC("%s", "binding port with default address family\n"); // bind HTTP listen port if (port == 0) { const auto bound_port = srv->bind_to_any_port(hostname); diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index ed4bda241..07a842bd6 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -287,7 +287,7 @@ std::vector> make_llama_cmpl_schema(const common_params & ->set_desc("Chat format used internally by the server") ->set_handler([&](field_eval_context & ctx, const json & data) { ctx.params.chat_parser_params.format = static_cast(data.at("chat_format").get()); - SRV_INF("Chat format: %s\n", common_chat_format_name(ctx.params.chat_parser_params.format)); + SRV_TRC("chat format: %s\n", common_chat_format_name(ctx.params.chat_parser_params.format)); })); add((new field_str("reasoning_format")) diff --git a/tools/server/server-stream.cpp b/tools/server/server-stream.cpp index 757c36ad2..785c28b3a 100644 --- a/tools/server/server-stream.cpp +++ b/tools/server/server-stream.cpp @@ -339,11 +339,11 @@ void stream_pipe_producer::close() { // httplib bails its content provider the moment is_peer_alive() goes false, so pump the rest // of the generation into the ring buffer here. a DELETE flips is_cancelled and cuts it short if (done_ || session_->is_cancelled()) { - SRV_INF("stream_pipe close: skip drain (done=%d cancelled=%d) conv=%s\n", + SRV_TRC("stream_pipe close: skip drain (done=%d cancelled=%d) conv=%s\n", done_ ? 1 : 0, session_->is_cancelled() ? 1 : 0, session_->conversation_id.c_str()); return; } - SRV_INF("stream_pipe close: draining conv=%s\n", session_->conversation_id.c_str()); + SRV_TRC("stream_pipe close: draining conv=%s\n", session_->conversation_id.c_str()); size_t drained = 0; std::string chunk; while (true) { @@ -357,7 +357,7 @@ void stream_pipe_producer::close() { break; } } - SRV_INF("stream_pipe close: drain ended conv=%s bytes=%zu\n", session_->conversation_id.c_str(), drained); + SRV_TRC("stream_pipe close: drain ended conv=%s bytes=%zu\n", session_->conversation_id.c_str(), drained); } std::shared_ptr stream_pipe_producer::create(stream_session_ptr session, @@ -520,7 +520,7 @@ server_http_context::handler_t make_stream_delete_handler() { if (conv_id.empty()) { return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST); } - SRV_INF("DELETE /v1/stream/%s -> evict_and_cancel\n", conv_id.c_str()); + SRV_TRC("DELETE /v1/stream/%s -> evict_and_cancel\n", conv_id.c_str()); g_stream_sessions.evict_and_cancel(conv_id); auto res = std::make_unique(); res->status = 204; @@ -550,8 +550,7 @@ std::string stream_conv_id_from_headers(const std::map void stream_session_attach_pipe(server_http_res & res, const std::map & headers) { std::string conversation_id = stream_conv_id_from_headers(headers); - SRV_INF("stream_session_attach_pipe: conv_id=%s (empty=%d)\n", - conversation_id.c_str(), conversation_id.empty() ? 1 : 0); + SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0); if (conversation_id.empty()) { return; } diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index a9ebac013..775f50baf 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1626,7 +1626,7 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t const int cur_lcp_len = it->tokens.get_common_prefix(prompt.tokens); if (cur_lcp_len == (int) prompt.tokens.size()) { - SRV_INF("%s", " - prompt is already in the cache, skipping\n"); + SRV_TRC("%s", " - prompt is already in the cache, skipping\n"); return nullptr; } } @@ -1636,7 +1636,7 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t const int len = it->tokens.get_common_prefix(prompt.tokens); if (len == (int) it->tokens.size()) { - SRV_WRN(" - removing obsolete cached prompt with length %d\n", len); + SRV_TRC(" - removing obsolete cached prompt with length %d\n", len); it = states.erase(it); } else { @@ -1681,7 +1681,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok float f_keep_best = prompt.tokens.size() > 0 ? float(lcp_best) / prompt.tokens.size() : -1.0f; // empty slot: any cache entry wins float sim_best = float(lcp_best) / tokens_new.size(); - SRV_INF(" - looking for better prompt, base f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); + SRV_TRC(" - looking for better prompt, base f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); auto it_best = states.end(); @@ -1706,7 +1706,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok } if (it_best != states.end()) { - SRV_INF(" - found better prompt with f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); + SRV_TRC(" - found better prompt with f_keep = %.3f, sim = %.3f\n", f_keep_best, sim_best); { auto & data = it_best->data.main; @@ -1783,11 +1783,11 @@ void server_prompt_cache::update() { } } - SRV_INF(" - cache state: %zu prompts, %.3f MiB (limits: %.3f MiB, %zu tokens, %zu est)\n", + SRV_TRC(" - cache state: %zu prompts, %.3f MiB (limits: %.3f MiB, %zu tokens, %zu est)\n", states.size(), size() / (1024.0 * 1024.0), limit_size / (1024.0 * 1024.0), limit_tokens, limit_tokens_cur); for (const auto & state : states) { - SRV_INF(" - prompt %p: %7d tokens, checkpoints: %2zu, %9.3f MiB\n", + SRV_TRC(" - prompt %p: %7d tokens, checkpoints: %2zu, %9.3f MiB\n", (const void *)&state, state.n_tokens(), state.checkpoints.size(), state.size() / (1024.0 * 1024.0)); } } diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 1bbc99d89..eafef86ba 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -124,7 +124,7 @@ int llama_server(int argc, char ** argv) { } if (params.n_parallel < 0) { - SRV_INF("%s", "n_parallel is set to auto, using n_parallel = 4 and kv_unified = true\n"); + SRV_TRC("%s", "n_parallel is set to auto, using n_parallel = 4 and kv_unified = true\n"); params.n_parallel = 4; params.kv_unified = true; @@ -338,7 +338,7 @@ int llama_server(int argc, char ** argv) { std::function clean_up; if (is_router_server) { - SRV_INF("%s", "starting router server, no model will be loaded in this process\n"); + SRV_INF("%s", "starting server in router mode. models will be automatically loaded on-demand\n"); clean_up = [&models_routes]() { SRV_INF("%s: cleaning up before exit...\n", __func__); @@ -391,9 +391,6 @@ int llama_server(int argc, char ** argv) { }); } - // load the model - SRV_INF("%s", "loading model\n"); - if (!ctx_server.load_model(params)) { clean_up(); if (ctx_http.thread.joinable()) { @@ -429,8 +426,9 @@ int llama_server(int argc, char ** argv) { SetConsoleCtrlHandler(reinterpret_cast(console_ctrl_handler), true); #endif + SRV_INF("listening on %s\n", ctx_http.listening_address.c_str()); + if (is_router_server) { - SRV_INF("router server is listening on %s\n", ctx_http.listening_address.c_str()); SRV_WRN("%s", "NOTE: router mode is experimental\n"); SRV_WRN("%s", " it is not recommended to use this mode in untrusted environments\n"); @@ -446,8 +444,6 @@ int llama_server(int argc, char ** argv) { // when the HTTP server stops, clean up and exit clean_up(); } else { - SRV_INF("server is listening on %s\n", ctx_http.listening_address.c_str()); - // optionally, notify router server that this instance is ready std::thread monitor_thread; if (child.is_child()) { From c1a1c8ee94e37ff7ba2c872e783aaf7f77e0f320 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= Date: Sun, 28 Jun 2026 12:34:11 +0200 Subject: [PATCH 09/17] common : allow --offline in llama download (#25091) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expose the existing --offline flag to `llama download` so a script can run it to check whether a model is already cached and ready to be served without touching the network. Also fix a latent use-after-free in the URL-task on_done callback: first_path is block-scoped and was captured by reference, but invoked after the block ends. Signed-off-by: Adrien Gallouët --- common/arg.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 841ca3ce2..c289ff713 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -467,7 +467,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params // the first part is what gets loaded, so point params.model.path at it if (!url_tasks.empty()) { std::string first_path = url_tasks.front().local_path; - url_tasks.front().on_done = [&]() { params.model.path = first_path; }; + url_tasks.front().on_done = [&, first_path]() { params.model.path = first_path; }; } for (auto & task : url_tasks) { tasks.push_back(std::move(task)); @@ -3471,7 +3471,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex [](common_params & params) { params.offline = true; } - ).set_env("LLAMA_ARG_OFFLINE")); + ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD}).set_env("LLAMA_ARG_OFFLINE")); add_opt(common_arg( {"-lv", "--verbosity", "--log-verbosity"}, "N", string_format("Set the verbosity threshold. Messages with a higher verbosity will be ignored. Values:\n" From d1b34251bc57b696a5c91968069f8a0e6be13ef4 Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Sun, 28 Jun 2026 15:01:34 +0200 Subject: [PATCH 10/17] spec : add DFlash support (#22105) * spec: add DFlash v2 support * dflash: support sliding window attention per layer_types * docs: add dflash section --------- Co-authored-by: Kashif Rasul --- common/common.h | 3 +- common/speculative.cpp | 303 ++++++++++++++++++++++++++++++++++++- conversion/__init__.py | 1 + conversion/qwen.py | 52 +++++++ docs/speculative.md | 29 +++- gguf-py/gguf/constants.py | 18 +++ src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-context.cpp | 4 +- src/llama-graph.cpp | 7 +- src/llama-model.cpp | 6 +- src/models/dflash.cpp | 276 +++++++++++++++++++++++++++++++++ src/models/models.h | 16 ++ tests/test-llama-archs.cpp | 4 +- 14 files changed, 712 insertions(+), 9 deletions(-) create mode 100644 src/models/dflash.cpp diff --git a/common/common.h b/common/common.h index d56f6064b..2adb310b8 100644 --- a/common/common.h +++ b/common/common.h @@ -169,6 +169,7 @@ enum common_speculative_type { COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, // standalone draft model speculative decoding COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, // Eagle3 speculative decoding COMMON_SPECULATIVE_TYPE_DRAFT_MTP, // Multi-token prediction + COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, // DFlash speculative decoding COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding based on n-grams COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values @@ -384,7 +385,7 @@ struct common_params_speculative { uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { - return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3; + return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH; }); return needs_rs_seq ? draft.n_max : 0u; diff --git a/common/speculative.cpp b/common/speculative.cpp index a3495c3a1..3951bbed5 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -33,6 +33,7 @@ const std::map common_speculative_type_fro {"draft-simple", COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE}, {"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3}, {"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP}, + {"draft-dflash", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH}, {"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE}, {"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K}, {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V}, @@ -898,6 +899,296 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { } }; +// DFlash: block-diffusion drafting with a draft-side KV cache injection +struct common_speculative_impl_draft_dflash : public common_speculative_impl { + common_params_speculative_draft params; + + llama_batch batch; // noise tokens + llama_batch batch_inject; // target features for KV cache injection + + std::vector smpls; + + int32_t n_embd_dec = 0; // draft hidden size + int32_t n_embd_enc = 0; // target_layer_ids_n * target_hidden_size + int32_t n_embd_tgt = 0; // target model hidden size + + int32_t block_size = 0; + llama_token mask_token_id = 0; + + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices + uint32_t target_layer_ids_n = 0; + + // scratch buffer for concatenated target features [n_tokens, n_embd_enc] + std::vector features_buf; + + common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, n_seq) + , params(params.draft) + { + auto * ctx_tgt = this->params.ctx_tgt; + auto * ctx_dft = this->params.ctx_dft; + GGML_ASSERT(ctx_tgt && ctx_dft && "DFlash requires ctx_tgt and ctx_dft to be set"); + + const llama_model * model_dft = llama_get_model(ctx_dft); + const llama_model * model_tgt = llama_get_model(ctx_tgt); + + target_layer_ids = llama_model_target_layer_ids (model_dft); + target_layer_ids_n = llama_model_target_layer_ids_n(model_dft); + GGML_ASSERT(target_layer_ids_n > 0 && "DFlash model has no target_layer_ids"); + + n_embd_tgt = llama_model_n_embd(model_tgt); + n_embd_dec = llama_model_n_embd(model_dft); + n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt; + + // read the trained block size from the dflash.block_size metadata key + block_size = 16; + { + char buf[32] = {}; + if (llama_model_meta_val_str(model_dft, "dflash.block_size", buf, sizeof(buf)) >= 0) { + block_size = std::atoi(buf); + } + } + mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); + + LOG_INF("%s: adding speculative implementation 'draft-dflash'\n", __func__); + 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\n", __func__, block_size, mask_token_id, target_layer_ids_n); + + // DFlash input is [id_last, * (block_size-1)], so it can draft at most block_size-1 tokens per step + if (this->params.n_max > block_size - 1) { + LOG_WRN("%s: requested draft size %d exceeds the trained DFlash block size %d -- clamping to %d draft tokens per step\n", + __func__, this->params.n_max, block_size - 1, block_size - 1); + this->params.n_max = block_size - 1; + } + + batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); + batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq); + + smpls.resize(n_seq); + for (auto & s : smpls) { + common_params_sampling sparams; + sparams.no_perf = false; + sparams.top_k = 1; + sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K }; + s.reset(common_sampler_init(model_dft, sparams)); + } + + // turn on extraction of the target layers' input embeddings + for (uint32_t k = 0; k < target_layer_ids_n; ++k) { + llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true); + } + + llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true); + llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention + } + + ~common_speculative_impl_draft_dflash() override { + llama_batch_free(batch); + llama_batch_free(batch_inject); + } + + void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { + if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { + return; + } + + const int32_t N = (int32_t) prompt.size(); + if (N <= 0) { + return; + } + + const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(params.ctx_dft), seq_id); + if (pos_max < N - 1) { + LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - process() did not run on every prefill ubatch. " + "Drafts may degrade.\n", + __func__, (int) pos_max, N - 1); + } + } + + bool process(const llama_batch & batch_in) override { + if (batch_in.n_tokens <= 0) { + return true; + } + + if (batch_in.token == nullptr || batch_in.embd != nullptr) { + return true; + } + + const int32_t n_tokens = batch_in.n_tokens; + + // per-seq inclusive batch range (assumes each seq's tokens are contiguous in the batch) + std::vector i_batch_beg(n_seq, -1); + std::vector i_batch_end(n_seq, -1); + for (int32_t k = 0; k < n_tokens; ++k) { + GGML_ASSERT(batch_in.n_seq_id[k] == 1); + const llama_seq_id seq_id = batch_in.seq_id[k][0]; + if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { + continue; + } + i_batch_end[seq_id] = k; + if (i_batch_beg[seq_id] < 0) { + i_batch_beg[seq_id] = k; + } + } + + auto * ctx_tgt = this->params.ctx_tgt; + auto * ctx_dft = this->params.ctx_dft; + + const int32_t n_ubatch = (int32_t) llama_n_ubatch(ctx_dft); + + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + if (i_batch_beg[seq_id] < 0) { + continue; + } + const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1; + + for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) { + const int32_t n_chunk = std::min(n_ubatch, n_rows - offset); + + // gather this chunk's target features, interleaved by extract layer + features_buf.resize((size_t) n_chunk * n_embd_enc); + for (uint32_t k = 0; k < target_layer_ids_n; ++k) { + const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]); + if (!layer) { + GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]); + } + for (int32_t i = 0; i < n_chunk; ++i) { + float * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt; + const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt; + std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float)); + } + } + + // fuse extracted features through DFlash encoder + llama_batch enc_batch = { + /*.n_tokens =*/ n_chunk, + /*.token =*/ nullptr, + /*.embd =*/ features_buf.data(), + /*.pos =*/ nullptr, + /*.n_seq_id =*/ nullptr, + /*.seq_id =*/ nullptr, + /*.logits =*/ nullptr, + }; + + int32_t rc = llama_encode(ctx_dft, enc_batch); + if (rc != 0) { + LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n", + __func__, rc, (int) n_chunk, (int) offset); + return false; + } + + const float * inp_g = llama_get_embeddings_nextn(ctx_dft); + GGML_ASSERT(inp_g && "DFlash encoder produced no output."); + + // inject the DFlash decoder K/V cache at the tokens' target positions + batch_inject.n_tokens = n_chunk; + std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float)); + + for (int32_t i = 0; i < n_chunk; ++i) { + batch_inject.pos[i] = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + batch_inject.n_seq_id[i] = 1; + batch_inject.seq_id[i][0] = seq_id; + batch_inject.logits[i] = false; + } + rc = llama_decode(ctx_dft, batch_inject); + if (rc != 0) { + LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n", + __func__, rc, (int) n_chunk, (int) offset); + return false; + } + } + } + + return true; + } + + void draft(common_speculative_draft_params_vec & dparams) override { + auto & ctx_dft = params.ctx_dft; + + common_batch_clear(batch); + + // build one batch holding every drafting sequence's noise block into a single decode) + // record where each block starts and its size + std::vector i_block_beg(n_seq, -1); + std::vector n_block (n_seq, 0); + + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + auto & dp = dparams[seq_id]; + if (!dp.drafting) { + continue; + } + + common_sampler_reset(smpls[seq_id].get()); + + const int32_t n = (int32_t) dp.n_past; + + int32_t n_draft = params.n_max; + if (dp.n_max > 0) { + n_draft = std::min(n_draft, dp.n_max); + } + + const int32_t n_block_tokens = n_draft + 1; // id_last + n_draft * + i_block_beg[seq_id] = batch.n_tokens; + n_block [seq_id] = n_block_tokens; + for (int32_t i = 0; i < n_block_tokens; ++i) { + common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true); + } + } + + if (batch.n_tokens == 0) { + return; + } + + // decode all sequence's noise block in a single batch + int ret = llama_decode(ctx_dft, batch); + if (ret != 0) { + LOG_WRN("%s: llama_decode returned %d\n", __func__, ret); + return; + } + + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + if (i_block_beg[seq_id] < 0) { + continue; + } + auto & dp = dparams[seq_id]; + + const int32_t beg = i_block_beg[seq_id]; + const int32_t n_block_tokens = n_block[seq_id]; + + auto * smpl = smpls[seq_id].get(); + + auto & result = *dp.result; + + // greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1 + for (int32_t i = 1; i < n_block_tokens; ++i) { + common_sampler_sample(smpl, ctx_dft, beg + i, true); + + const auto * cur_p = common_sampler_get_candidates(smpl, true); + + for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", + seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p, + common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); + } + + const llama_token id = cur_p->data[0].id; + + common_sampler_accept(smpl, id, true); + + result.push_back(id); + } + } + } + + void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { + // noop + } + + bool need_embd() const override { + return false; + } +}; + struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft) @@ -1841,6 +2132,7 @@ std::string common_speculative_type_to_str(common_speculative_type type) { case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: return "draft-simple"; case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3"; case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: return "draft-mtp"; + case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: return "draft-dflash"; case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple"; case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k"; case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v"; @@ -1893,6 +2185,7 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) { case COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE: case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: + case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: n_max = std::max(n_max, std::max(0, spec->draft.n_max)); break; case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: @@ -1930,6 +2223,7 @@ common_speculative * common_speculative_init(common_params_speculative & params, bool has_draft_simple = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE)); bool has_draft_eagle3 = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3)) && params.draft.ctx_dft != nullptr; bool has_draft_mtp = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP)) && params.draft.ctx_dft != nullptr; + bool has_draft_dflash = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)) && params.draft.ctx_dft != nullptr; @@ -1940,7 +2234,7 @@ common_speculative * common_speculative_init(common_params_speculative & params, bool has_ngram_mod = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_MOD)); // when adding a new type - update here the logic above - static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 9); + static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 10); // this list here defines the priority of the speculators // the one with highest priority are listed first @@ -1970,6 +2264,9 @@ common_speculative * common_speculative_init(common_params_speculative & params, if (has_draft_mtp) { configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, params)); } + if (has_draft_dflash) { + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, params)); + } } std::vector> impls = {}; @@ -1990,6 +2287,10 @@ common_speculative * common_speculative_init(common_params_speculative & params, impls.push_back(std::make_unique(config.params, n_seq)); break; } + case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: { + impls.push_back(std::make_unique(config.params, n_seq)); + break; + } case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple); diff --git a/conversion/__init__.py b/conversion/__init__.py index 5aad203e5..4a1fd5bb7 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -50,6 +50,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "DeepseekV2ForCausalLM": "deepseek", "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", + "DFlashDraftModel": "qwen", "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", "DistilBertModel": "bert", diff --git a/conversion/qwen.py b/conversion/qwen.py index 6b85eb9aa..81f450e40 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -625,3 +625,55 @@ class Qwen3_5TextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReor @ModelBase.register("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM") class Qwen3_5MoeTextModel(_Qwen35MtpMixin, _Qwen35MRopeMixin, _LinearAttentionVReorderBase): model_arch = gguf.MODEL_ARCH.QWEN35MOE + + +@ModelBase.register("DFlashDraftModel") +class DFlashModel(Qwen3Model): + model_arch = gguf.MODEL_ARCH.DFLASH + + def set_vocab(self): + if self.target_model_dir is None: + raise ValueError( + "DFlash draft model requires --target-model-dir to be specified. " + "Please provide the path to the target model directory containing the tokenizer." + ) + logger.info(f"DFlash: Using tokenizer from target model: {self.target_model_dir}") + original_dir = self.dir_model + self.dir_model = self.target_model_dir + super().set_vocab() + self.dir_model = original_dir + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + block_size = self.hparams.get("block_size", 16) + self.gguf_writer.add_uint32(f"{self.gguf_writer.arch}.block_size", block_size) + dflash_config = self.hparams.get("dflash_config", {}) + + target_layer_ids = dflash_config.get("target_layer_ids", []) + if target_layer_ids: + extract_layer_ids = [i + 1 for i in target_layer_ids] + self.gguf_writer.add_array(f"{self.gguf_writer.arch}.target_layers", extract_layer_ids) + + mask_token_id = dflash_config.get("mask_token_id", None) + if mask_token_id is not None: + self.gguf_writer.add_mask_token_id(mask_token_id) + + use_sliding_window = self.hparams.get("use_sliding_window", False) + sliding_window = 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) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name == "fc.weight": + yield (name, data_torch) + return + if name == "hidden_norm.weight": + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch) + return + if not name.startswith("model."): + name = "model." + name + yield from super().modify_tensors(data_torch, name, bid) diff --git a/docs/speculative.md b/docs/speculative.md index 8f91256c4..4100b92f8 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -52,6 +52,32 @@ Supported EAGLE-3 draft models include: For the full and up-to-date list of supported models, see #18039. +### DFlash (`draft-dflash`) + +DFlash produces an entire block of draft tokens in a single forward pass (block diffusion) and +injects the target model's hidden states into the draft model's attention, instead of drafting one +token at a time. This keeps the draft model small while making drafting GPU-friendly. Unlike EAGLE-3 +(a single-layer autoregressive draft), the DFlash draft uses several transformer layers but emits a +whole block per draft step. + +The draft is a small block-diffusion model trained for a specific target (for example +`z-lab/Qwen3-4B-DFlash` for `Qwen/Qwen3-4B`). Convert it with `--target-model-dir` so it inherits the +target's tokenizer and token embeddings: + +```bash +python convert_hf_to_gguf.py z-lab/Qwen3-4B-DFlash \ + --target-model-dir Qwen/Qwen3-4B --outtype bf16 --outfile Qwen3-4B-DFlash.gguf + +llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DFlash.gguf \ + --spec-type draft-dflash --spec-draft-n-max 15 -fa on --jinja +``` + +`--spec-draft-n-max` is clamped to the draft model's trained block size. + +See: + +- #22105 + ### n-gram Cache (`ngram-cache`) An n-gram is a sequence of n tokens. The n-gram cache implementation maintains statistics about short n-gram sequences. @@ -147,7 +173,7 @@ If a draft model is combined with a draftless decoding the draftless decoding ha ### General Speculative Parameters ``` ---spec-type [none|draft-simple|draft-eagle3|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod] +--spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod] comma-separated list of types of speculative decoding to use (default: none) (env: LLAMA_ARG_SPEC_TYPE) @@ -287,6 +313,7 @@ Specifies a comma-separated list of speculative decoding types to use. | `none` | No speculative decoding (default) | | `draft-simple` | Use a simple draft model for speculation | | `draft-eagle3` | Use an EAGLE-3 draft model that reads the target's hidden states | +| `draft-dflash` | Use a DFlash block-diffusion draft model that emits a block per step | | `draft-mtp` | Use Multi Token Prediction (MTP) heads from the main model | | `ngram-cache` | Use n-gram cache lookup | | `ngram-simple` | Use simple n-gram pattern matching | diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 1bda9452d..bcd10beb0 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -517,6 +517,7 @@ class MODEL_ARCH(IntEnum): PANGU_EMBED = auto() MISTRAL3 = auto() EAGLE3 = auto() + DFLASH = auto() MISTRAL4 = auto() PADDLEOCR = auto() MIMO2 = auto() @@ -1074,6 +1075,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.PANGU_EMBED: "pangu-embedded", MODEL_ARCH.MISTRAL3: "mistral3", MODEL_ARCH.EAGLE3: "eagle3", + MODEL_ARCH.DFLASH: "dflash", MODEL_ARCH.MISTRAL4: "mistral4", MODEL_ARCH.PADDLEOCR: "paddleocr", MODEL_ARCH.MIMO2: "mimo2", @@ -4086,6 +4088,22 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FC, MODEL_TENSOR.D2T, ], + MODEL_ARCH.DFLASH: [ + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FC, + MODEL_TENSOR.ENC_OUTPUT_NORM, + ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 4a52d9772..d80915ffd 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -129,6 +129,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_PANGU_EMBED, "pangu-embedded" }, { LLM_ARCH_MISTRAL3, "mistral3" }, { LLM_ARCH_EAGLE3, "eagle3" }, + { LLM_ARCH_DFLASH, "dflash" }, { LLM_ARCH_MISTRAL4, "mistral4" }, { LLM_ARCH_PADDLEOCR, "paddleocr" }, { LLM_ARCH_MIMO2, "mimo2" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 989da06d8..946518d5f 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -143,6 +143,7 @@ enum llm_arch { LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, + LLM_ARCH_DFLASH, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9f8a8fdb8..029141e2a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -100,10 +100,10 @@ llama_context::llama_context( cparams.ctx_other = params.ctx_other; } - if (model.arch == LLM_ARCH_EAGLE3) { + if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH) { if (model.tok_embd == nullptr || model.output == nullptr) { if (params.ctx_other == nullptr) { - throw std::runtime_error("EAGLE3 requires ctx_other to be set (this warning is normal during memory fitting)"); + throw std::runtime_error(model.arch_name() + " requires ctx_other to be set (this warning is normal during memory fitting)"); } cparams.ctx_other = params.ctx_other; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 68c9e606c..3ded70bc0 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -486,7 +486,11 @@ void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { mctx->set_input_k_idxs(self_k_idxs, ubatch); mctx->set_input_v_idxs(self_v_idxs, ubatch); - mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + // the mask is left unallocated when the graph only stores K/V without attending + // (e.g. DFlash's KV-injection pass) + if (self_kq_mask && self_kq_mask->buffer) { + mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + } if (self_k_rot) { mctx->set_input_k_rot(self_k_rot); @@ -904,6 +908,7 @@ void llm_graph_result::reset() { t_logits = nullptr; t_embd = nullptr; t_embd_pooled = nullptr; + t_h_nextn = nullptr; t_layer_inp.resize(LLAMA_MAX_LAYERS); std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 6cb0ec379..eaf29505c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -291,6 +291,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_mistral3(params); case LLM_ARCH_EAGLE3: return new llama_model_eagle3(params); + case LLM_ARCH_DFLASH: + return new llama_model_dflash(params); case LLM_ARCH_MIMO2: return new llama_model_mimo2(params); case LLM_ARCH_KIMI_LINEAR: @@ -2494,6 +2496,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_STEP35: case LLM_ARCH_TALKIE: case LLM_ARCH_MELLUM: + case LLM_ARCH_DFLASH: return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_QWEN2VL: @@ -2617,7 +2620,8 @@ bool llama_model_has_encoder(const llama_model * model) { switch (model->arch) { case LLM_ARCH_T5: case LLM_ARCH_T5ENCODER: - case LLM_ARCH_EAGLE3: return true; + case LLM_ARCH_EAGLE3: + case LLM_ARCH_DFLASH: return true; default: return false; } } diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp new file mode 100644 index 000000000..a7b4f4435 --- /dev/null +++ b/src/models/dflash.cpp @@ -0,0 +1,276 @@ +#include "models.h" + +#include "llama-kv-cache.h" +#include "llama-kv-cache-iswa.h" + +void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { + + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) { + throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata"); + } + + hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd; + + LLAMA_LOG_INFO("%s: DFlash extract_layers = [", __func__); + for (size_t i = 0; i < target_layer_ids.size(); ++i) { + LLAMA_LOG_INFO("%d%s", target_layer_ids[i], i + 1 < target_layer_ids.size() ? ", " : ""); + } + LLAMA_LOG_INFO("]\n"); + + // optional interleaved sliding-window attention with per-layer pattern array. + // DFlash has a single rope, so the SWA rope == main rope. + if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) { + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + } + + type = LLM_TYPE_UNKNOWN; +} + +void llama_model_dflash::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_embd_inp = hparams.n_embd_inp_enc(); + + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); + output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); + + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head }, 0); + layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0); + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); + + 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); + + 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); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); + } +} + +std::unique_ptr llama_model_dflash::build_arch_graph(const llm_graph_params & params) const { + switch (params.gtype) { + case LLM_GRAPH_TYPE_ENCODER: + return std::make_unique>(*this, params); + case LLM_GRAPH_TYPE_DEFAULT: + case LLM_GRAPH_TYPE_DECODER: + return std::make_unique>(*this, params); + default: + GGML_ABORT("invalid graph type"); + }; +} + +template <> +ggml_tensor * llama_model_dflash::graph::build_inp_embd_enc() const { + auto inp_target = std::make_unique(hparams.n_embd_inp_enc()); + + inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp_enc(), n_tokens); + ggml_set_input(inp_target->embd); + + ggml_tensor * cur = inp_target->embd; + cb(cur, "inp_embd", -1); + + res->add_input(std::move(inp_target)); + + return cur; +} + +// DFlash Encoder: processes target model features through feature fusion layer +template <> +llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + ggml_tensor * cur = build_inp_embd_enc(); + + cur = build_lora_mm(model.fc, cur); + cb(cur, "fc_out", -1); + + cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1); + cb(cur, "enc_norm_out", -1); + + ggml_set_output(cur); + res->t_h_nextn = cur; + + ggml_build_forward_expand(gf, cur); +} + +// DFlash decoder, dual-mode by batch type: +// * embd batch -> fused target features: project + inject K/V into the cache. +// * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens +template <> +llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + ggml_tensor * inp_pos = build_inp_pos(); + + // optional iSWA: pick the matching attention input + const bool use_iswa = hparams.swa_type != LLAMA_SWA_TYPE_NONE; + + llm_graph_input_attn_kv * inp_attn = nullptr; + llm_graph_input_attn_kv_iswa * inp_attn_iswa = nullptr; + if (use_iswa) { + inp_attn_iswa = build_attn_inp_kv_iswa(); + } else { + inp_attn = build_attn_inp_kv(); + } + + const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + + // KV cache injection + if (ubatch.embd) { + auto inp = std::make_unique(n_embd); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * inp_g = inp->embd; + cb(inp_g, "inp_g_embeddings", -1); + + res->add_input(std::move(inp)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g); + ggml_tensor * Vcur = build_lora_mm(layer.wv, inp_g); + + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + cb(Kcur, "Kcur_injected", il); + cb(Vcur, "Vcur_injected", il); + + if (use_iswa) { + // route each layer's K/V to its sub-cache: SWA layers -> sliding cache, full -> dense + const bool is_swa = hparams.is_swa(il); + const auto * kv = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base(); + ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs(); + ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs(); + ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il)); + ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il)); + } else { + ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il)); + ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il)); + } + } + + res->t_embd = inp_g; + + ggml_build_forward_expand(gf, inp_g); + return; + } + + // tok_embd from the target model (shared via ctx_other) + auto * tok_embd = model.tok_embd; + if (tok_embd == nullptr) { + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + + GGML_ASSERT(model_other->tok_embd != nullptr && "DFlash decoder requires the target model's token embeddings"); + tok_embd = model_other->tok_embd; + } + + auto inp = std::make_unique(n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); + cb(inpL, "inp_noise_embd", -1); + + res->add_input(std::move(inp)); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(noise_norm, "noise_norm", il); + + ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); + ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); + ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm); + + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + // 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); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + layer.ffn_up, NULL, NULL, + layer.ffn_gate, NULL, NULL, + layer.ffn_down, NULL, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "l_out", il); + + inpL = cur; + } + + ggml_tensor * cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + + res->t_embd = cur; + + // lm_head from the target model (shared via ctx_other) + auto * output = model.output; + if (output == nullptr) { + GGML_ASSERT(cparams.ctx_other != nullptr); + const auto * model_other = llama_get_model(cparams.ctx_other); + GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection"); + output = model_other->output; + } + + cur = build_lora_mm(output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index 2ac8415a3..d89ab96d0 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1122,6 +1122,22 @@ struct llama_model_eagle3 : public llama_model_base { }; +struct llama_model_dflash : public llama_model_base { + llama_model_dflash(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + template + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + ggml_tensor * build_inp_embd_enc() const; + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_mistral4 : public llama_model_deepseek2 { llama_model_mistral4(const struct llama_model_params & params) : llama_model_deepseek2(params) {} // reuse load_arch_hparams and load_arch_tensors from llama_model_deepseek2 diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 524971ae4..c781d2903 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -451,7 +451,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const ggml if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } - if (arch == LLM_ARCH_EAGLE3) { + if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { continue; } for (bool moe : {false, true}) { @@ -557,7 +557,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } - if (arch == LLM_ARCH_EAGLE3) { + if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { continue; } From f68a788b0bf41d88cb90fd93c7027471a2ebe30a Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sun, 28 Jun 2026 15:50:31 +0200 Subject: [PATCH 11/17] jinja: add --dump-prog for debugging (#25086) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * jinja: add --dump-prog for debugging * Update common/jinja/runtime.cpp Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> --------- Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> --- common/jinja/runtime.cpp | 46 +++++++++++++ common/jinja/runtime.h | 127 +++++++++++++++++++++++++++++++++++ tests/test-chat-template.cpp | 20 +++++- 3 files changed, 190 insertions(+), 3 deletions(-) diff --git a/common/jinja/runtime.cpp b/common/jinja/runtime.cpp index f98cb0876..474129df2 100644 --- a/common/jinja/runtime.cpp +++ b/common/jinja/runtime.cpp @@ -954,4 +954,50 @@ value keyword_argument_expression::execute_impl(context & ctx) { return mk_val(k, v); } +std::string runtime::debug_dump_program(const program & prog, const std::string & src) { + std::ostringstream oss; + size_t lvl = 0; + context ctx; + ctx.src.reset(new std::string(src)); + + auto indent = [](size_t lvl) -> std::string { + return std::string(lvl * 2, ' '); + }; + + ctx.visitor = [&](bool is_leaf, statement * node, std::vector children) { + oss << indent(lvl) << node->type() << ":\n"; + lvl++; + if (is_leaf) { + const auto & pos = node->pos; + oss << indent(lvl) << "(leaf) at " << get_line_col(src, pos) << " in source:\n"; + std::string snippet = peak_source(src, pos); + string_replace_all(snippet, "\n", "\n" + indent(lvl)); + oss << indent(lvl) << snippet << "\n"; + } else { + for (auto & [label, children_vec] : children) { + oss << indent(lvl) << label << ":\n"; + lvl++; + if (children_vec.empty()) { + oss << indent(lvl) << "\n\n"; + } else { + for (auto * child : children_vec) { + if (!child) { + continue; + } + child->visit(ctx); + } + } + lvl--; + } + } + lvl--; + }; + + for (const auto & stmt : prog.body) { + stmt->visit(ctx); + } + + return oss.str(); +} + } // namespace jinja diff --git a/common/jinja/runtime.h b/common/jinja/runtime.h index 37b4c35ca..0884a1592 100644 --- a/common/jinja/runtime.h +++ b/common/jinja/runtime.h @@ -47,12 +47,19 @@ const T * cast_stmt(const statement_ptr & ptr) { // not thread-safe void enable_debug(bool enable); +// for visiting AST nodes +// function signature: void(bool is_leaf, statement * node, pair of ) +using visitor_pair = std::pair>; +using visitor_fn = std::function)>; + struct context { std::shared_ptr src; // for debugging; use shared_ptr to avoid copying on scope creation std::time_t current_time; // for functions that need current time bool is_get_stats = false; // whether to collect stats + visitor_fn visitor; + // src is optional, used for error reporting context(std::string src = "") : src(std::make_shared(std::move(src))) { env = mk_val(); @@ -99,6 +106,15 @@ private: value_object env; }; +// utils for visiting AST nodes +static std::vector stmts_to_ptr(const statements & stmts) { + std::vector children; + for (const auto & stmt : stmts) { + children.push_back(stmt.get()); + } + return children; +} + /** * Base class for all nodes in the AST. */ @@ -106,6 +122,7 @@ struct statement { size_t pos; // position in source, for debugging virtual ~statement() = default; virtual std::string type() const { return "Statement"; } + virtual void visit(context & ctx) { ctx.visitor(true, this, {}); } // execute_impl must be overridden by derived classes virtual value execute_impl(context &) { throw_exec_error(); } @@ -166,6 +183,13 @@ struct if_statement : public statement { std::string type() const override { return "If"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"test", {test.get()}}, + {"body", stmts_to_ptr(body)}, + {"alternate", stmts_to_ptr(alternate)} + }); + } }; struct identifier; @@ -190,6 +214,14 @@ struct for_statement : public statement { std::string type() const override { return "For"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"loopvar", {loopvar.get()}}, + {"iterable", {iterable.get()}}, + {"body", stmts_to_ptr(body)}, + {"default_block", stmts_to_ptr(default_block)} + }); + } }; struct break_statement : public statement { @@ -241,6 +273,13 @@ struct set_statement : public statement { std::string type() const override { return "Set"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"assignee", {assignee.get()}}, + {"value", {val.get()}}, + {"body", stmts_to_ptr(body)} + }); + } }; struct macro_statement : public statement { @@ -256,6 +295,13 @@ struct macro_statement : public statement { std::string type() const override { return "Macro"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"name", {name.get()}}, + {"args", stmts_to_ptr(args)}, + {"body", stmts_to_ptr(body)} + }); + } }; struct comment_statement : public statement { @@ -289,6 +335,12 @@ struct member_expression : public expression { } std::string type() const override { return "MemberExpression"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"object", {object.get()}}, + {"property", {property.get()}} + }); + } }; struct call_expression : public expression { @@ -302,6 +354,12 @@ struct call_expression : public expression { } std::string type() const override { return "CallExpression"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"callee", {callee.get()}}, + {"args", stmts_to_ptr(args)} + }); + } }; /** @@ -405,6 +463,12 @@ struct binary_expression : public expression { } std::string type() const override { return "BinaryExpression"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"left", {left.get()}}, + {"right", {right.get()}} + }); + } }; /** @@ -431,6 +495,12 @@ struct filter_expression : public expression { std::string type() const override { return "FilterExpression"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"operand", {operand.get()}}, + {"filter", {filter.get()}} + }); + } }; struct filter_statement : public statement { @@ -443,6 +513,12 @@ struct filter_statement : public statement { } std::string type() const override { return "FilterStatement"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"filter", {filter.get()}}, + {"body", stmts_to_ptr(body)} + }); + } }; /** @@ -468,6 +544,12 @@ struct select_expression : public expression { } return lhs->execute_impl(ctx); } + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"lhs", {lhs.get()}}, + {"test", {test.get()}} + }); + } }; /** @@ -486,6 +568,12 @@ struct test_expression : public expression { } std::string type() const override { return "TestExpression"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"operand", {operand.get()}}, + {"test", {test.get()}} + }); + } }; /** @@ -501,6 +589,11 @@ struct unary_expression : public expression { } std::string type() const override { return "UnaryExpression"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"argument", {argument.get()}} + }); + } }; struct slice_expression : public expression { @@ -518,6 +611,13 @@ struct slice_expression : public expression { [[noreturn]] value execute_impl(context &) override { throw std::runtime_error("must be handled by MemberExpression"); } + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"start_expr", {start_expr.get()}}, + {"stop_expr", {stop_expr.get()}}, + {"step_expr", {step_expr.get()}} + }); + } }; struct keyword_argument_expression : public expression { @@ -531,6 +631,12 @@ struct keyword_argument_expression : public expression { } std::string type() const override { return "KeywordArgumentExpression"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"key", {key.get()}}, + {"val", {val.get()}} + }); + } }; struct spread_expression : public expression { @@ -539,6 +645,11 @@ struct spread_expression : public expression { chk_type(this->argument); } std::string type() const override { return "SpreadExpression"; } + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"argument", {argument.get()}} + }); + } }; struct call_statement : public statement { @@ -553,6 +664,13 @@ struct call_statement : public statement { } std::string type() const override { return "CallStatement"; } value execute_impl(context & ctx) override; + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"call", {call.get()}}, + {"caller_args", stmts_to_ptr(caller_args)}, + {"body", stmts_to_ptr(body)} + }); + } }; struct ternary_expression : public expression { @@ -575,6 +693,13 @@ struct ternary_expression : public expression { return false_expr->execute(ctx); } } + void visit(context & ctx) override { + ctx.visitor(false, this, { + {"condition", {condition.get()}}, + {"true_expr", {true_expr.get()}}, + {"false_expr", {false_expr.get()}} + }); + } }; struct raised_exception : public std::exception { @@ -648,6 +773,8 @@ struct runtime { } return parts; } + + static std::string debug_dump_program(const program & prog, const std::string & src); }; } // namespace jinja diff --git a/tests/test-chat-template.cpp b/tests/test-chat-template.cpp index d971b2374..6a6292cd0 100644 --- a/tests/test-chat-template.cpp +++ b/tests/test-chat-template.cpp @@ -25,7 +25,7 @@ using json = nlohmann::ordered_json; static int main_automated_tests(void); static void run_multiple(const std::string& dir_path, bool stop_on_first_failure, const json& input, bool use_common = false); -static void run_single(const std::string& contents, json input, bool use_common = false, const std::string & output_path = ""); +static void run_single(const std::string& contents, json input, bool use_common = false, bool dump_prog = false, const std::string & output_path = ""); static std::string HELP = R"( Usage: test-chat-template [OPTIONS] PATH_TO_TEMPLATE @@ -35,6 +35,7 @@ Options: --json Path to the JSON input file. --stop-on-first-fail Stop testing on the first failure (default: false). --no-common Use direct Jinja engine instead of common chat templates (default: use common). + --dump-prog Dump the parsed program for debugging (only for single template runs). --output Path to output results (only for single template runs). If PATH_TO_TEMPLATE is a file, runs that single template. If PATH_TO_TEMPLATE is a directory, runs all .jinja files in that directory. @@ -118,6 +119,7 @@ int main(int argc, char ** argv) { std::string & json_to_use = DEFAULT_JSON; bool stop_on_first_fail = false; bool use_common = true; + bool dump_prog = false; for (size_t i = 1; i < args.size(); i++) { if (args[i] == "--help" || args[i] == "-h") { @@ -136,6 +138,8 @@ int main(int argc, char ** argv) { i++; } else if (args[i] == "--no-common") { use_common = false; + } else if (args[i] == "--dump-prog") { + dump_prog = true; } else if (tmpl_path.empty()) { tmpl_path = args[i]; } else { @@ -172,7 +176,7 @@ int main(int argc, char ** argv) { std::string contents = std::string( std::istreambuf_iterator(infile), std::istreambuf_iterator()); - run_single(contents, input_json, use_common, output_path); + run_single(contents, input_json, use_common, dump_prog, output_path); } else { std::cerr << "Error: PATH_TO_TEMPLATE is not a valid file or directory: " << tmpl_path << "\n"; return 1; @@ -276,11 +280,21 @@ static jinja::value_string format_using_direct_engine( } -void run_single(const std::string& contents, json input, bool use_common, const std::string & output_path) { +void run_single(const std::string& contents, json input, bool use_common, bool dump_prog, const std::string & output_path) { jinja::enable_debug(true); jinja::value_string output_parts; + if (dump_prog) { + jinja::lexer lexer; + auto lexer_res = lexer.tokenize(contents); + jinja::program ast = jinja::parse_from_tokens(lexer_res); + std::string prog_dump = jinja::runtime::debug_dump_program(ast, contents); + std::cout << "\n=== DUMPED PROGRAM ===\n"; + std::cout << prog_dump << "\n"; + return; + } + if (use_common) { std::string bos_token = ""; std::string eos_token = ""; From c818263f2a5ddab028dea5f169ea2b2266421125 Mon Sep 17 00:00:00 2001 From: Aldehir Rojas Date: Sun, 28 Jun 2026 09:53:32 -0500 Subject: [PATCH 12/17] chat : implement minicpm5 parser (#24889) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add minicpm5 tool call parser * Refactor MiniCPM5 PEG parser per review feedback * Fix jinja min/max API to match Jinja2 * modify by review * MiniCPM5: use autoparser for XML tool calls and fix grammar preserved-token triggers * MiniCPM5: fix streaming tool-arg placeholder and remove alt XML markers * skip min/max attribute tests in -py mode * test-jinja: use real expected output for min/max attribute tests * MiniCPM5: revert shared mapper and history fallbacks per review Drop streaming tool-arg placeholder workarounds from the generic PEG mapper and restore strict tool-call argument JSON parsing so MiniCPM5 support stays limited to autoparser/diff-analyzer changes. * chat : refactor minicpm5 back to dedicated parser * cont : simplify grammar * cont : refactor * cont : fixes * cont : rename template to openbmb-MiniCPM5-1B.jinja * cont : add message delimiters * cont : fix tests --------- Co-authored-by: zhangtao Co-authored-by: 张涛 <> --- common/chat.cpp | 151 +++++++++++++++++ common/jinja/value.cpp | 44 +++++ models/templates/openbmb-MiniCPM5-1B.jinja | 179 +++++++++++++++++++++ tests/test-chat.cpp | 78 +++++++++ tests/test-jinja.cpp | 30 ++++ 5 files changed, 482 insertions(+) create mode 100644 models/templates/openbmb-MiniCPM5-1B.jinja diff --git a/common/chat.cpp b/common/chat.cpp index 0cee80434..a38c62f5f 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -2376,6 +2376,149 @@ static void func_args_not_string(json & messages) { } +// MiniCPM5 format: +// - Reasoning: {reasoning} (optional) +// - Tool calls: value +static common_chat_params common_chat_params_init_minicpm5(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + data.preserved_tokens = { + "", + "", + "", + "", + }; + + data.thinking_start_tag = ""; + data.thinking_end_tag = ""; + + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "<|im_start|>assistant" }, + { COMMON_CHAT_ROLE_TOOL, "<|im_start|>user\n" }, + { COMMON_CHAT_ROLE_USER, "<|im_start|>user" }, + { COMMON_CHAT_ROLE_SYSTEM, "<|im_start|>system" }, + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = inputs.json_schema.is_object() && !inputs.json_schema.empty(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = "<|im_start|>assistant\n\n" + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += "\n\n\n" + msg.render_content(); + } + + data.prompt += data.generation_prompt; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.literal("<|im_start|>assistant\n"); + + auto reasoning = p.eps(); + if (extract_reasoning) { + reasoning = ("" << p.reasoning(p.until("")) << "") + p.space(); + } + + // Response format parser + if (has_response_format) { + return generation_prompt + reasoning + p.content(p.schema(p.json(), "response-format", inputs.json_schema)); + } + + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + // CDATA lets a value carry characters that would otherwise close the tag (e.g. + // ); capture the inner text only, excluding the CDATA markers. + auto string_value = p.choice({ + p.literal("")) + p.literal("]]>"), "]]>") + p.tool_arg_close(p.literal("")), + p.negate(p.literal("")) + p.tool_arg_close(p.literal("")), "") + }); + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + const std::string name = function.at("name"); + auto params = function.contains("parameters") ? function.at("parameters") : json::object(); + + auto args = p.eps(); + if (params.contains("properties") && params.at("properties").is_object() && !params.at("properties").empty()) { + auto schema_info = common_schema_info(); + schema_info.resolve_refs(params); + + auto arg_choice = p.choice(); + for (const auto & [prop_name, prop_schema] : params.at("properties").items()) { + auto value_parser = p.eps(); + if (schema_info.resolves_to_string(prop_schema)) { + value_parser = string_value; + } else { + value_parser = p.tool_arg_json_value( + p.schema(p.json(), "tool-" + name + "-arg-" + prop_name + "-schema", prop_schema, false) + ) + p.tool_arg_close(p.literal("")); + } + + auto arg_rule = p.tool_arg( + p.tool_arg_open(p.literal("")) + + value_parser + ); + + arg_choice |= arg_rule; + } + args = p.zero_or_more(arg_choice + p.space()); + } + + auto tool_parser = p.tool( + p.tool_open(p.literal("")) + << p.tool_args(args) + << p.tool_close(p.literal(""))); + + tool_choice |= p.rule("tool-" + name, tool_parser); + }); + + auto max_calls = inputs.parallel_tool_calls ? -1 : 1; + auto tool_calls = p.trigger_rule("tool-call", p.repeat(tool_choice + p.space(), 1, max_calls)); + + auto content = p.content(p.until(" common_chat_try_specialized_template( return common_chat_params_init_gemma4(tmpl, params); } + // MiniCPM5 - XML tool calls with ... + if (src.find("Tool usage guidelines:") != std::string::npos && + src.find("(val) ? mk_val(std::move(arr)) : mk_val(std::move(arr)); }}, + {"min", [](const func_args & args) -> value { + args.ensure_count(1, 4); + args.ensure_vals(); + value val_case = args.get_kwarg_or_pos("case_sensitive", 1); + value attribute = args.get_kwarg_or_pos("attribute", 2); + if (!attribute->is_undefined()) { + throw not_implemented_exception("min: attribute not implemented"); + } + // FIXME: min is currently always case sensitive + (void) val_case; + const auto & arr = args.get_pos(0)->as_array(); + if (arr.empty()) { + return mk_val(); + } + value result = arr[0]; + for (size_t i = 1; i < arr.size(); ++i) { + if (value_compare(arr[i], result, value_compare_op::lt)) { + result = arr[i]; + } + } + return result; + }}, + {"max", [](const func_args & args) -> value { + args.ensure_count(1, 4); + args.ensure_vals(); + value val_case = args.get_kwarg_or_pos("case_sensitive", 1); + value attribute = args.get_kwarg_or_pos("attribute", 2); + if (!attribute->is_undefined()) { + throw not_implemented_exception("max: attribute not implemented"); + } + // FIXME: max is currently always case sensitive + (void) val_case; + const auto & arr = args.get_pos(0)->as_array(); + if (arr.empty()) { + return mk_val(); + } + value result = arr[0]; + for (size_t i = 1; i < arr.size(); ++i) { + if (value_compare(arr[i], result, value_compare_op::gt)) { + result = arr[i]; + } + } + return result; + }}, {"unique", array_unique_not_implemented}, }; return builtins; diff --git a/models/templates/openbmb-MiniCPM5-1B.jinja b/models/templates/openbmb-MiniCPM5-1B.jinja new file mode 100644 index 000000000..cb2934c45 --- /dev/null +++ b/models/templates/openbmb-MiniCPM5-1B.jinja @@ -0,0 +1,179 @@ +{{- bos_token }}{%- if tools %} + {%- set tool_definitions %} + {{- "# Tools\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson(ensure_ascii=False) }} + {%- endfor %} + {{- '\n\n\nTool usage guidelines:\n- You may call zero or more functions. If no function calls are needed, just answer normally and do not include any .\n- When calling a function, return an XML object within using:\nparam-value\n- param-value may be multi-line. If it contains <, & or newline characters, wrap it in a CDATA block: ' }} + {%- endset %} + + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {%- if '' in messages[0].content %} + {{- messages[0].content.replace('', tool_definitions) }} + {%- else %} + {{- messages[0].content + '\n\n' + tool_definitions }} + {%- endif %} + {%- else %} + {{- tool_definitions.lstrip() }} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('') and message.content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} +{%- endfor %} +{%- for message in messages %} + {%- if message.content is string %} + {%- set content = message.content %} + {%- else %} + {%- set content = '' %} + {%- endif %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + + {%- if message.tool_calls %} + {%- set content_parts = content.split('') %} + {%- set processed_content = content_parts[0] %} + {%- set tool_calls_count = message.tool_calls|length %} + {%- set tool_sep_count = content_parts|length - 1 %} + {%- set min_count = [tool_calls_count, tool_sep_count]|min %} + + {%- for i in range(1, content_parts|length) %} + {%- set tool_index = i - 1 %} + {%- if tool_index < tool_calls_count %} + {%- set tool_call = message.tool_calls[tool_index] %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- set single_tool_xml %} + {{- '' }} + {%- if tool_call.arguments %} + {%- set args_dict = tool_call.arguments %} + {%- for param_name, param_value in args_dict.items() %} + {{- '' }} + {%- if param_value is string and ('<' in param_value or '&' in param_value or '\n' in param_value) %} + {{- '' }} + {%- else %} + {{- param_value }} + {%- endif %} + {{- '' }} + {%- endfor %} + {%- endif %} + {{- '' }} + {%- endset %} + {%- set processed_content = processed_content + single_tool_xml + content_parts[i] %} + {%- else %} + {%- set processed_content = processed_content + content_parts[i] %} + {%- endif %} + {%- endfor %} + + {%- if tool_calls_count > tool_sep_count %} + {%- for remaining_index in range(tool_sep_count, tool_calls_count) %} + {%- set tool_call = message.tool_calls[remaining_index] %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- set remaining_tool_xml %} + {{- '' }} + {%- if tool_call.arguments %} + {%- set args_dict = tool_call.arguments %} + {%- for param_name, param_value in args_dict.items() %} + {{- '' }} + {%- if param_value is string and ('<' in param_value or '&' in param_value or '\n' in param_value) %} + {{- '' }} + {%- else %} + {{- param_value }} + {%- endif %} + {{- '' }} + {%- endfor %} + {%- endif %} + {{- '' }} + {%- endset %} + {%- set processed_content = processed_content + remaining_tool_xml %} + {%- endfor %} + {%- endif %} + + {%- set content = processed_content %} + {%- endif %} + + {%- if loop.index0 > ns.last_query_index %} + {%- if reasoning_content %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + + {%- if message.tool_calls and not has_tool_sep %} + {%- for tool_call in message.tool_calls %} + {%- if (loop.first and content) or (not loop.first) %} + {{- '\n' }} + {%- endif %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- '' }} + {%- if tool_call.arguments %} + {%- set args_dict = tool_call.arguments %} + {%- for param_name, param_value in args_dict.items() %} + {{- '' }} + {%- if param_value is string and ('<' in param_value or '&' in param_value or '\n' in param_value) %} + {{- '' }} + {%- else %} + {{- param_value }} + {%- endif %} + {{- '' }} + {%- endfor %} + {%- endif %} + {{- '' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {%- if message.content is string %} + {{- content }} + {%- else %} + {{- message.content | tojson(ensure_ascii=False) }} + {%- endif %} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined %} + {%- if enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- elif enable_thinking is true %} + {{- '\n' }} + {%- endif %} + {%- endif %} +{%- endif %} diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index c38aed8cf..5f71e5da6 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -5593,6 +5593,77 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .expect_content("Hello, world!\nWhat's up?") .run(); } + + // MiniCPM5 - XML tool calls with ... + { + auto tst = peg_tester("models/templates/openbmb-MiniCPM5-1B.jinja", detailed_debug); + + tst.test("Hello, world!\nWhat's up?") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist) + .run(); + + tst.test(R"(print('Hello, World!'))") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ python_tool }) + .expect_tool_calls({ { "python", R"#({"code": "print('Hello, World!')"})#", {} } }) + .run(); + + tst.test(R"()") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ empty_args_tool }) + .expect(simple_assist_msg("", "", "empty_args", "{}")) + .run(); + + tst.test(R"(print('x'))") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .parallel_tool_calls(true) + .tools({ python_tool }) + .expect_tool_calls({ { "python", R"#({"code": "print('x')"})#", {} } }) + .run(); + + // CDATA lets a string value carry characters that would otherwise close the tag. + tst.test(R"(hi ]]>)") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ html_tool }) + .expect_tool_calls({ { "html", R"#({"markup": "hi "})#", {} } }) + .run(); + + tst.test(R"(I'm thinkingprint('hey'))") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ python_tool }) + .expect_reasoning("I'm thinking") + .expect_tool_calls({ { "python", R"#({"code": "print('hey')"})#", {} } }) + .run(); + + tst.test(R"(print('x') +print('y'))") + .enable_thinking(false) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .parallel_tool_calls(true) + .tools({ python_tool }) + .expect_tool_calls({ + { "python", R"#({"code": "print('x')"})#", {} }, + { "python", R"#({"code": "print('y')"})#", {} }, + }) + .run(); + + tst.test(" thinkingHello, world!\nWhat's up?") + .enable_thinking(true) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .messages({ message_user, message_assist_prefill_reasoning }) + .add_generation_prompt(false) + .continue_final_message(COMMON_CHAT_CONTINUATION_REASONING) + .expect_reasoning("I'm thinking") + .expect_content("Hello, world!\nWhat's up?") + .run(); + } } static void test_template_generation_prompt() { @@ -5740,6 +5811,13 @@ static void test_template_generation_prompt() { check(tmpls, continuation_content(), "<｜Assistant｜>I'm thinkingHello, "); check(tmpls, continuation_reasoning(), "<｜Assistant｜>I'm"); } + + { + auto tmpls = read_templates("models/templates/openbmb-MiniCPM5-1B.jinja"); + check(tmpls, basic(), "<|im_start|>assistant\n\n"); + check(tmpls, continuation_content(), "<|im_start|>assistant\n\nI'm thinking\n\n\nHello, "); + check(tmpls, continuation_reasoning(), "<|im_start|>assistant\n\nI'm"); + } } // Test the developer role to system workaround with a simple mock template diff --git a/tests/test-jinja.cpp b/tests/test-jinja.cpp index 81bbcd55a..d8d1892a9 100644 --- a/tests/test-jinja.cpp +++ b/tests/test-jinja.cpp @@ -1584,6 +1584,36 @@ static void test_array_methods(testing & t) { "6" ); + test_template(t, "array|min", + "{{ [tool_calls_count, tool_sep_count]|min }}", + {{"tool_calls_count", 2}, {"tool_sep_count", 1}}, + "1" + ); + + test_template(t, "array|max", + "{{ [tool_calls_count, tool_sep_count]|max }}", + {{"tool_calls_count", 2}, {"tool_sep_count", 1}}, + "2" + ); + + test_template(t, "array|min attribute", + "{{ items|min(attribute='x') }}", + {{"items", json::array({ + json({{"x", 2}}), + json({{"x", 1}}), + })}}, + "{'x': 1}" + ); + + test_template(t, "array|max attribute", + "{{ items|max(attribute='x') }}", + {{"items", json::array({ + json({{"x", 2}}), + json({{"x", 1}}), + })}}, + "{'x': 2}" + ); + // not used by any chat templates // test_template(t, "array.insert()", // "{% set _ = arr.insert(1, 'x') %}{{ arr|join(',') }}", From fa72bc6826a5ff30dda3abd1e2fd87ba91df5762 Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Sun, 28 Jun 2026 20:31:48 +0200 Subject: [PATCH 13/17] dflash: refactor draft model conversion (#25110) * dflash: refactor draft model conversion * apply fix for eagle3 convert --- conversion/llama.py | 6 +++--- conversion/qwen.py | 24 ++++++++++-------------- gguf-py/gguf/constants.py | 1 + gguf-py/gguf/gguf_writer.py | 12 ++++++++++++ gguf-py/gguf/tensor_mapping.py | 5 +++++ 5 files changed, 31 insertions(+), 17 deletions(-) diff --git a/conversion/llama.py b/conversion/llama.py index b43cc994a..315a619c9 100644 --- a/conversion/llama.py +++ b/conversion/llama.py @@ -73,7 +73,7 @@ class LlamaModel(TextModel): target_num_layers = target_config["num_hidden_layers"] target_layers = [2, target_num_layers // 2, target_num_layers - 3] logger.info(f"EAGLE-3: target_layers = {target_layers} (target model has {target_num_layers} layers)") - self.gguf_writer.add_array(f"{self.gguf_writer.arch}.target_layers", target_layers) + self.gguf_writer.add_target_layers(target_layers) # target_hidden_size: prefer eagle3 config, fallback to target config if eagle3_raw_config.get("target_hidden_size") is not None: @@ -83,12 +83,12 @@ class LlamaModel(TextModel): target_hidden_size = target_config["hidden_size"] src = "target model config" logger.info(f"EAGLE-3: target_hidden_size = {target_hidden_size} (from {src})") - self.gguf_writer.add_uint32(f"{self.gguf_writer.arch}.target_hidden_size", target_hidden_size) + self.gguf_writer.add_target_hidden_size(target_hidden_size) # norm_before_residual (RedHat-style eagle3 specific) norm_before_residual = eagle3_raw_config.get("norm_before_residual", False) logger.info(f"EAGLE-3: norm_before_residual = {norm_before_residual}") - self.gguf_writer.add_bool(f"{self.gguf_writer.arch}.norm_before_residual", norm_before_residual) + self.gguf_writer.add_norm_before_residual(norm_before_residual) def set_vocab(self): # eagle3: use tokenizer from target model if provided diff --git a/conversion/qwen.py b/conversion/qwen.py index 81f450e40..0356bd2da 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -643,21 +643,21 @@ class DFlashModel(Qwen3Model): super().set_vocab() self.dir_model = original_dir + mask_token_id = self.hparams.get("dflash_config", {}).get("mask_token_id") + if mask_token_id is not None: + self.gguf_writer.add_mask_token_id(mask_token_id) + def set_gguf_parameters(self): super().set_gguf_parameters() block_size = self.hparams.get("block_size", 16) - self.gguf_writer.add_uint32(f"{self.gguf_writer.arch}.block_size", block_size) + self.gguf_writer.add_block_size(block_size) dflash_config = self.hparams.get("dflash_config", {}) target_layer_ids = dflash_config.get("target_layer_ids", []) if target_layer_ids: extract_layer_ids = [i + 1 for i in target_layer_ids] - self.gguf_writer.add_array(f"{self.gguf_writer.arch}.target_layers", extract_layer_ids) - - mask_token_id = dflash_config.get("mask_token_id", None) - if mask_token_id is not None: - self.gguf_writer.add_mask_token_id(mask_token_id) + 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") @@ -667,13 +667,9 @@ class DFlashModel(Qwen3Model): self.gguf_writer.add_sliding_window(sliding_window) self.gguf_writer.add_sliding_window_pattern(is_swa) - def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - if name == "fc.weight": - yield (name, data_torch) - return - if name == "hidden_norm.weight": - yield (self.format_tensor_name(gguf.MODEL_TENSOR.ENC_OUTPUT_NORM), data_torch) - return + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item if not name.startswith("model."): name = "model." + name - yield from super().modify_tensors(data_torch, name, bid) + return super().filter_tensors((name, gen)) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index bcd10beb0..52e9e54de 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -156,6 +156,7 @@ class Keys: DENSE_FEAT_OUT_SIZE = "{arch}.{dense}_feat_out" TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" + BLOCK_SIZE = "{arch}.block_size" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" class Attention: diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index a06ec88b3..610555f5e 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -940,6 +940,18 @@ class GGUFWriter: def add_sliding_window(self, value: int) -> None: self.add_uint32(Keys.Attention.SLIDING_WINDOW.format(arch=self.arch), value) + def add_block_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.BLOCK_SIZE.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) + + def add_target_hidden_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.TARGET_HIDDEN_SIZE.format(arch=self.arch), value) + + def add_norm_before_residual(self, value: bool) -> None: + self.add_bool(Keys.LLM.NORM_BEFORE_RESIDUAL.format(arch=self.arch), value) + def add_attention_scale(self, value: float) -> None: self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 5f1e28818..9efb36f8a 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1283,6 +1283,11 @@ class TensorNameMap: MODEL_TENSOR.ENC_OUTPUT_NORM: ( "encoder.final_layer_norm", # t5 "layer_norm", # neobert + "model.hidden_norm", # dflash + ), + + MODEL_TENSOR.FC: ( + "model.fc", # dflash ), MODEL_TENSOR.CLS: ( From 7cb8576e7c35789801b0f4e0653ffb41d0c5e046 Mon Sep 17 00:00:00 2001 From: Pascal Date: Sun, 28 Jun 2026 21:06:43 +0200 Subject: [PATCH 14/17] ui: fix stop and reasoning skip in single-model mode (#25084) --- tools/server/server-context.cpp | 2 ++ tools/server/server-models.cpp | 5 +++- tools/server/server-stream.cpp | 7 +++++ tools/ui/src/lib/stores/chat.svelte.ts | 36 ++++++++++++++++++++------ 4 files changed, 41 insertions(+), 9 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index bb1c236cb..06f236b4c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2450,6 +2450,8 @@ private: server_slot * slot = get_slot_by_cmpl_id(task.params.control_cmpl_id); if (slot == nullptr) { + SRV_WRN("control %s on unknown completion id=%s, no live slot\n", + task.params.control_action.c_str(), task.params.control_cmpl_id.c_str()); res->success = false; res->message = "no active completion for this id"; queue_results.send(std::move(res)); diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 0380f98a3..81da00c0e 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -1983,7 +1983,10 @@ void server_models_routes::init_routes() { cli.set_read_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); cli.set_write_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); auto resp = cli.Delete(child_path.c_str()); - (void) resp; // best effort, 404 and network errors are equivalent to no op + (void) resp; // the child logs its own miss when the session is unknown there + } else { + SRV_WRN("router stop for unknown conv_id=%s, no owning child in the conv map\n", + conv_id.c_str()); } // drop the tracking entry, the session is being torn down models.conv_models.forget(conv_id); diff --git a/tools/server/server-stream.cpp b/tools/server/server-stream.cpp index 785c28b3a..c2bba8ec4 100644 --- a/tools/server/server-stream.cpp +++ b/tools/server/server-stream.cpp @@ -218,6 +218,13 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i std::unique_lock lock(map_mu); auto it = sessions.find(conversation_id); if (it == sessions.end()) { + std::string live; + for (const auto & kv : sessions) { + if (!live.empty()) live += ", "; + live += kv.first; + } + SRV_WRN("stop on unknown stream session, conv_id=%s matched nothing, %zu live: [%s]\n", + conversation_id.c_str(), sessions.size(), live.c_str()); return; } s = it->second; diff --git a/tools/ui/src/lib/stores/chat.svelte.ts b/tools/ui/src/lib/stores/chat.svelte.ts index faaaa9755..43f9e72c7 100644 --- a/tools/ui/src/lib/stores/chat.svelte.ts +++ b/tools/ui/src/lib/stores/chat.svelte.ts @@ -154,7 +154,13 @@ class ChatStore { }); if (convId === conversationsStore.activeConversation?.id) this.currentResponse = response; } - private clearChatStreaming(convId: string): void { + private clearChatStreaming(convId: string, messageId?: string): void { + // session aware: a stale generation must not wipe a newer one's streaming state on the + // same conversation, that would drop the frozen stop identity and stop the wrong session + if (messageId !== undefined) { + const cur = this.chatStreamingStates.get(convId); + if (cur && cur.messageId !== messageId) return; + } this.chatStreamingStates.delete(convId); if (convId === conversationsStore.activeConversation?.id) this.currentResponse = ''; } @@ -1055,11 +1061,14 @@ class ChatStore { modelOverride?: string | null, firstUserMessageContent?: string ): Promise { - let effectiveModel = modelOverride; + // the ::model suffix in the stream identity is only for router mode, where it routes to the + // owning child. in single-model mode the identity stays the bare conv id so that attach, stop + // and reattach all agree, regardless of fresh send vs regenerate passing a resolved model + let effectiveModel: string | null | undefined = undefined; - if (isRouterMode() && !effectiveModel) { + if (isRouterMode()) { const conversationModel = this.getConversationModel(allMessages); - effectiveModel = selectedModelName() || conversationModel; + effectiveModel = modelOverride || selectedModelName() || conversationModel; } if (isRouterMode() && effectiveModel) { @@ -1074,6 +1083,9 @@ class ChatStore { let resolvedModel: string | null = null; let modelPersisted = false; const convId = assistantMessage.convId; + // freeze the POST identity from t0 so a stop cancels with the exact session key, + // never a stale or empty model resolved later + this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); const recordModel = (modelName: string | null | undefined, persistImmediately = true): void => { if (!modelName) return; @@ -1103,7 +1115,7 @@ class ChatStore { }; const updateStreamingUI = () => { - this.setChatStreaming(convId, streamedContent, currentMessageId); + this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); const idx = conversationsStore.findMessageIndex(currentMessageId); conversationsStore.updateMessageAtIndex(idx, { content: streamedContent }); }; @@ -1111,7 +1123,7 @@ class ChatStore { const cleanupStreamingState = () => { this.setStreamingActive(false); this.setChatLoading(convId, false); - this.clearChatStreaming(convId); + this.clearChatStreaming(convId, currentMessageId); this.setProcessingState(convId, null); }; @@ -1128,7 +1140,7 @@ class ChatStore { onReasoningChunk: (chunk: string) => { streamedReasoningContent += chunk; // mark streaming state so a stop mid-thinking can persist the partial reasoning - this.setChatStreaming(convId, streamedContent, currentMessageId); + this.setChatStreaming(convId, streamedContent, currentMessageId, effectiveModel); const idx = conversationsStore.findMessageIndex(currentMessageId); conversationsStore.updateMessageAtIndex(idx, { reasoningContent: streamedReasoningContent @@ -1405,7 +1417,7 @@ class ChatStore { // detached drain keeps producing tokens until eos or max_tokens. use the frozen identity // captured when the session started, not the live dropdown const streamStateForStop = this.chatStreamingStates.get(convId); - const modelForStop = streamStateForStop?.model ?? selectedModelName(); + const modelForStop = streamStateForStop?.model; void ChatService.cancelServerStream(convId, modelForStop); this.abortRequest(convId); this.setChatLoading(convId, false); @@ -1846,6 +1858,14 @@ class ChatStore { updateStreamingContent(originalContent + appendedContent); this.setChatReasoning(msg.convId, false); }, + onCompletionId: (id: string) => { + if (!id) return; + // refresh the message id so a later skip targets the live slot after a continue + conversationsStore.updateMessageAtIndex(conversationsStore.findMessageIndex(msg.id), { + completionId: id + }); + DatabaseService.updateMessage(msg.id, { completionId: id }).catch(() => {}); + }, onReasoningChunk: (chunk: string) => { appendedReasoning += chunk; hasReceivedContent = true; From dbdaece23de9ac63f2e7ca9e6bfcdc4fc156a3fa Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Sun, 28 Jun 2026 21:30:03 +0200 Subject: [PATCH 15/17] Revert "ui: fix accessibility for hover-gated interactive elements assisted by claude(in debugging and tests) (#24727)" (#25098) --- .../ChatAttachmentsListItemMcpPrompt.svelte | 2 +- .../ChatMessageUserPending.svelte | 2 +- .../SidebarNavigationConversationItem.svelte | 137 +++++++++++------- 3 files changed, 83 insertions(+), 58 deletions(-) diff --git a/tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte b/tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte index c55dfdec7..636e93f22 100644 --- a/tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte @@ -33,7 +33,7 @@ {#if !readonly && onRemove}
onRemove?.()} />
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte index 5c2913202..4be582b39 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte @@ -56,7 +56,7 @@
diff --git a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte index 2c1b9adf2..b1c2b78f6 100644 --- a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte +++ b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte @@ -39,6 +39,7 @@ depth = 0 }: Props = $props(); + let renderActionsDropdown = $state(false); let dropdownOpen = $state(false); let isLoading = $derived(getAllLoadingChats().includes(conversation.id)); @@ -70,10 +71,26 @@ } } + function handleMouseLeave() { + if (!dropdownOpen) { + renderActionsDropdown = false; + } + } + + function handleMouseOver() { + renderActionsDropdown = true; + } + function handleSelect() { onSelect?.(conversation.id); } + $effect(() => { + if (!dropdownOpen) { + renderActionsDropdown = false; + } + }); + onMount(() => { document.addEventListener('edit-active-conversation', handleGlobalEditEvent as EventListener); @@ -86,19 +103,23 @@ }); -
{ + if (!e.currentTarget.contains(e.relatedTarget as Node | null)) { + handleMouseLeave(); + } + }} > -
{#if depth > 0} @@ -109,7 +130,7 @@ @@ -125,15 +146,18 @@ {#if isLoading} - +
@@ -145,50 +169,52 @@
-
- { - e.stopPropagation(); - handleTogglePin(); - } - }, - { - icon: Pencil, - label: 'Edit', - onclick: handleEdit, - shortcut: ['shift', 'cmd', 'e'] - }, - { - icon: Download, - label: 'Export', - onclick: (e: Event) => { - e.stopPropagation(); - conversationsStore.downloadConversation(conversation.id); + {#if renderActionsDropdown} +
+ { + e.stopPropagation(); + handleTogglePin(); + } }, - shortcut: ['shift', 'cmd', 's'] - }, - { - icon: Trash2, - label: 'Delete', - onclick: handleDelete, - variant: 'destructive', - shortcut: ['shift', 'cmd', 'd'], - separator: true - } - ]} - /> -
-
+ { + icon: Pencil, + label: 'Edit', + onclick: handleEdit, + shortcut: ['shift', 'cmd', 'e'] + }, + { + icon: Download, + label: 'Export', + onclick: (e: Event) => { + e.stopPropagation(); + conversationsStore.downloadConversation(conversation.id); + }, + shortcut: ['shift', 'cmd', 's'] + }, + { + icon: Trash2, + label: 'Delete', + onclick: handleDelete, + variant: 'destructive', + shortcut: ['shift', 'cmd', 'd'], + separator: true + } + ]} + /> +
+ {/if} +