Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-18 00:35:04 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'upstream' into concedo_experimental

Browse Source 
# Conflicts:
#	.ecrc
#	.github/actions/get-tag-name/action.yml
#	.github/actions/windows-setup-cuda/action.yml
#	.github/workflows/build-cuda-ubuntu.yml
#	.github/workflows/build-cuda-windows.yml
#	.github/workflows/build-self-hosted.yml
#	.github/workflows/release.yml
#	.github/workflows/server-self-hosted.yml
#	CONTRIBUTING.md
#	app/CMakeLists.txt
#	ci/run.sh
#	docs/backend/OPENVINO.md
#	docs/ops.md
#	docs/ops/CUDA.csv
#	examples/eval-callback/CMakeLists.txt
#	ggml/src/ggml-hexagon/htp/cpy-ops.c
#	ggml/src/ggml-opencl/ggml-opencl.cpp
#	ggml/src/ggml-opencl/kernels/moe_reorder_b.cl
#	ggml/src/ggml-opencl/kernels/ssm_scan.cl
#	ggml/src/ggml-openvino/ggml-decoder.cpp
#	ggml/src/ggml-openvino/ggml-decoder.h
#	ggml/src/ggml-openvino/ggml-openvino-extra.cpp
#	ggml/src/ggml-openvino/ggml-openvino-extra.h
#	ggml/src/ggml-openvino/ggml-openvino.cpp
#	ggml/src/ggml-openvino/ggml-quants.cpp
#	ggml/src/ggml-openvino/ggml-quants.h
#	ggml/src/ggml-openvino/openvino/frontend.cpp
#	ggml/src/ggml-openvino/openvino/node_context.h
#	ggml/src/ggml-openvino/openvino/op/add.cpp
#	ggml/src/ggml-openvino/openvino/op/diag.cpp
#	ggml/src/ggml-openvino/openvino/op/div.cpp
#	ggml/src/ggml-openvino/openvino/op/glu_geglu_quick.cpp
#	ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp
#	ggml/src/ggml-openvino/openvino/op/mul_mat_id.cpp
#	ggml/src/ggml-openvino/openvino/op/mulmat.cpp
#	ggml/src/ggml-openvino/openvino/op/norm.cpp
#	ggml/src/ggml-openvino/openvino/op/pad.cpp
#	ggml/src/ggml-openvino/openvino/op/permute.cpp
#	ggml/src/ggml-openvino/openvino/op/rope.cpp
#	ggml/src/ggml-openvino/openvino/op/set_rows.cpp
#	ggml/src/ggml-openvino/openvino/op/transpose.cpp
#	ggml/src/ggml-openvino/openvino/op/unary_softplus.cpp
#	ggml/src/ggml-openvino/openvino/op_table.cpp
#	ggml/src/ggml-openvino/openvino/op_table.h
#	ggml/src/ggml-openvino/openvino/pass/squeeze_matmul.cpp
#	ggml/src/ggml-openvino/openvino/translate_session.cpp
#	ggml/src/ggml-openvino/utils.cpp
#	ggml/src/ggml-openvino/utils.h
#	scripts/hip/gcn-cdna-vgpr-check.py
#	tests/CMakeLists.txt
#	tests/test-backend-ops.cpp
#	tools/server/CMakeLists.txt
#	tools/tuning/CMakeLists.txt
This commit is contained in:
Concedo
2026-09-16 15:32:56 +08:00
parent 4929a7326c 930e2fa599
commit 802810260d
32 changed files with 695 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/test-cmake/.gitignore
+1
View File
@@ -1,3 +1,4 @@
llama-build-install
install
build
build-subdir
examples/test-cmake/CMakeLists.txt
+14 -5
View File
@@ -3,11 +3,20 @@ project(llama-simple)
set(CMAKE_CXX_STANDARD 17)
find_package(llama 0.1.0 REQUIRED)
option(LLAMA_TEST_USE_SUBDIR "Use add_subdirectory instead of find_package" OFF)
if(LLAMA_TEST_USE_SUBDIR)
add_subdirectory(../../ llama.cpp)
else()
find_package(llama 0.1.0 REQUIRED)
endif()
add_executable(test-cmake test-cmake.cpp)
target_link_libraries(test-cmake PRIVATE llama)
target_compile_definitions(test-cmake PRIVATE
LLAMA_BUILD_NUMBER=${LLAMA_BUILD_NUMBER}
LLAMA_BUILD_COMMIT="${LLAMA_BUILD_COMMIT}"
)
if(DEFINED LLAMA_BUILD_NUMBER)
target_compile_definitions(test-cmake PRIVATE
LLAMA_BUILD_NUMBER=${LLAMA_BUILD_NUMBER}
LLAMA_BUILD_COMMIT="${LLAMA_BUILD_COMMIT}"
)
endif()
examples/test-cmake/README.md
+14 -5
View File
@@ -5,17 +5,18 @@ enable troubleshooting issues and exploration. The idea is that this can be used
after making changes to llama.cpp installation cmake configuration and then
verify it locally.
### Usage
The following will configure, build, and install llama.cpp
### find_package
The following will configure, build, and install llama.cpp, and the build a
project that uses find_package to use the installation.
Configuring/build/install:
```console
./build-install.sh
```
The above command will create a directory named `install` in the current directory
which will have the follwing files in its lib directory:
which will have the following files in its lib directory:
```console
(venv) $ ls install/lib/
$ ls install/lib/
cmake libggml.so libllama-common.so.0 libllama.so.0.1.0 llama.cpp
libggml-base.so libggml.so.0 libllama-common.so.0.1.0 libmtmd.so pkgconfig
libggml-base.so.0 libggml.so.0.19.0 libllama.so libmtmd.so.0
@@ -24,7 +25,7 @@ libggml-base.so.0.19.0 libllama-common.so libllama.so.0 libmtmd.so
Build/run this project using the installation created above:
```console
(venv) $ ./build.sh
$ ./build.sh
-- Configuring done (0.0s)
-- Generating done (0.0s)
-- Build files have been written to: /path/to/llama.cpp/examples/test-cmake/build
@@ -34,3 +35,11 @@ Build/run this project using the installation created above:
load_backend: loaded CPU backend from /path/to/llama.cpp/examples/test-cmake/install/lib/llama.cpp/libggml-cpu-alderlake.so
[test-cmake] Backend initialized.
```
### add_subdirectory
The following will use add_subdirectory to include llama.cpp in a cmake project
and is intended to simulate projects that build llama.cpp in this way.
```console
$ USE_SUBDIR=ON ./build.sh
```
examples/test-cmake/build.sh
+14 -3
View File
@@ -2,6 +2,17 @@
set -e
cmake -S . -B build -DCMAKE_PREFIX_PATH="${PWD}/install"
cmake --build build
LD_LIBRARY_PATH="${PWD}/install/lib/llama.cpp:${PWD}/install/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" ./build/test-cmake
if [ "${USE_SUBDIR:-OFF}" = "ON" ]; then
BUILD_DIR="build-subdir"
CMAKE_ARGS="-DLLAMA_TEST_USE_SUBDIR=ON -DLLAMA_BUILD_COMMON=ON -DLLAMA_BUILD_TOOLS=ON -DLLAMA_BUILD_SERVER=ON-DLLAMA_BUILD_TESTS=ON"
LIB_PATH="${PWD}/${BUILD_DIR}/bin"
else
BUILD_DIR="build"
CMAKE_ARGS="-DCMAKE_PREFIX_PATH=${PWD}/install"
LIB_PATH="${PWD}/install/lib/llama.cpp"
fi
cmake --fresh -S . -B "${BUILD_DIR}" ${CMAKE_ARGS}
cmake --build "${BUILD_DIR}" -j 8
LD_LIBRARY_PATH="${LIB_PATH}:${PWD}/install/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" "./${BUILD_DIR}/test-cmake"
examples/test-cmake/test-cmake.cpp
+4
View File
@@ -2,8 +2,12 @@
#include <cstdio>
int main(void) {
#ifdef LLAMA_BUILD_NUMBER
printf("[test-cmake] llama.cpp version: %s, build: %d (%s)\n",
llama_version(), LLAMA_BUILD_NUMBER, LLAMA_BUILD_COMMIT);
#else
printf("[test-cmake] llama.cpp version: %s\n", llama_version());
#endif
printf("[test-cmake] ggml version: %s, commit: %s\n", ggml_version(), ggml_commit());
printf("[test-cmake] Initializing backend...\n");
llama_backend_init();
ggml/include/ggml-rpc.h
+1 -1
View File
@@ -6,7 +6,7 @@
extern "C" {
#endif
#define RPC_PROTO_MAJOR_VERSION 6
#define RPC_PROTO_MAJOR_VERSION 7
#define RPC_PROTO_MINOR_VERSION 0
#define RPC_PROTO_PATCH_VERSION 0
ggml/src/ggml-cuda/allreduce.cu
+32 -26
View File
@@ -1,6 +1,6 @@
#include "allreduce.cuh"
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
#if !defined(GGML_USE_MUSA)
#include "convert.cuh"
#include "ggml-impl.h"
@@ -11,11 +11,12 @@
#include <limits>
// ---------------------------------------------------------------------------
// CUDA AllReduce for tensor-parallel inference across two GPUs.
// AllReduce for tensor-parallel inference across two GPUs (CUDA or
// ROCm/HIP).
//
// Provides an in-place sum reduction over matching tensors on two CUDA
// devices in the same process. Used by the tensor-split path alongside
// NCCL; targets setups without NVLink, where data is exchanged between the
// Provides an in-place sum reduction over matching tensors on two GPUs
// in the same process. Used by the tensor-split path alongside NCCL;
// targets setups without NVLink/xGMI, where data is exchanged between the
// GPUs by staging it through pinned host memory over PCIe.
//
// Two reduction strategies are selected per call by tensor size:
@@ -161,11 +162,14 @@ static __global__ void ggml_cuda_ar_kernel(
__threadfence_system(); // make our signal visible system-wide
while (ggml_cuda_ar_signal_get(other_slot) != token) {
#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
#ifdef GGML_USE_HIP
// Equals ~100ns at 2500 MHz (sleeps for n * [1,64] clock cycles)
__builtin_amdgcn_s_sleep(4);
#elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
__nanosleep(100);
#else
NO_DEVICE_CODE;
#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
#endif // GGML_USE_HIP
}
}
@@ -280,7 +284,7 @@ struct ggml_cuda_ar_host_mapping {
}
rc = cudaHostGetDevicePointer(reinterpret_cast<void **>(&dev), host, 0);
if (rc != cudaSuccess) {
cudaFreeHost(host);
CUDA_CHECK(cudaFreeHost(host));
host = nullptr;
dev = nullptr;
}
@@ -289,7 +293,7 @@ struct ggml_cuda_ar_host_mapping {
void free() {
if (host) {
cudaFreeHost(host);
CUDA_CHECK(cudaFreeHost(host));
host = nullptr;
dev = nullptr;
}
@@ -401,7 +405,8 @@ ggml_cuda_ar_pipeline * ggml_cuda_ar_pipeline_init(const int * devices, size_t n
return nullptr;
}
// The chunked kernel uses __nanosleep, which is sm70+ (Volta+).
// The chunked kernel uses __nanosleep (NVIDIA, sm70+) or
// __builtin_amdgcn_s_sleep (AMD).
for (size_t i = 0; i < n_devices; ++i) {
const int cc = ggml_cuda_info().devices[devices[i]].cc;
if (cc < GGML_CUDA_CC_VOLTA) {
@@ -543,7 +548,7 @@ void ggml_cuda_ar_pipeline_free(ggml_cuda_ar_pipeline * p) {
for (int i = 0; i < p->n_devices; ++i) {
if (p->streams[i]) {
ggml_cuda_set_device(p->devices[i]);
cudaStreamSynchronize(p->streams[i]);
CUDA_CHECK(cudaStreamSynchronize(p->streams[i]));
}
}
@@ -552,28 +557,28 @@ void ggml_cuda_ar_pipeline_free(ggml_cuda_ar_pipeline * p) {
p->host_large[i].free();
if (p->dev_tmp[i]) {
ggml_cuda_set_device(p->devices[i]);
cudaFree(p->dev_tmp[i]);
CUDA_CHECK(cudaFree(p->dev_tmp[i]));
}
ggml_cuda_set_device(p->devices[i]);
for (int s = 0; s < GGML_CUDA_AR_POOL_SIZE; ++s) {
if (p->ev_pool[i][s].app) { cudaEventDestroy(p->ev_pool[i][s].app); }
if (p->ev_pool[i][s].app) { CUDA_CHECK(cudaEventDestroy(p->ev_pool[i][s].app)); }
for (int c = 0; c < GGML_CUDA_AR_COPY_MAX_CHUNKS; ++c) {
if (p->ev_pool[i][s].cpy[c]) { cudaEventDestroy(p->ev_pool[i][s].cpy[c]); }
if (p->ev_pool[i][s].cpy[c]) { CUDA_CHECK(cudaEventDestroy(p->ev_pool[i][s].cpy[c])); }
}
if (p->ev_pool[i][s].h2d) { cudaEventDestroy(p->ev_pool[i][s].h2d); }
if (p->ev_pool[i][s].ker) { cudaEventDestroy(p->ev_pool[i][s].ker); }
if (p->ev_pool[i][s].h2d) { CUDA_CHECK(cudaEventDestroy(p->ev_pool[i][s].h2d)); }
if (p->ev_pool[i][s].ker) { CUDA_CHECK(cudaEventDestroy(p->ev_pool[i][s].ker)); }
}
if (p->host_large_read_done[i]) {
ggml_cuda_set_device(p->devices[i]);
cudaEventDestroy(p->host_large_read_done[i]);
CUDA_CHECK(cudaEventDestroy(p->host_large_read_done[i]));
}
if (p->dev_tmp_kernel_done[i]) {
ggml_cuda_set_device(p->devices[i]);
cudaEventDestroy(p->dev_tmp_kernel_done[i]);
CUDA_CHECK(cudaEventDestroy(p->dev_tmp_kernel_done[i]));
}
if (p->streams[i]) {
ggml_cuda_set_device(p->devices[i]);
cudaStreamDestroy(p->streams[i]);
CUDA_CHECK(cudaStreamDestroy(p->streams[i]));
}
}
p->arrival.free();
@@ -952,13 +957,14 @@ bool ggml_cuda_ar_allreduce(
return ok;
}
#else // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
#else // defined(GGML_USE_MUSA)
// HIP and MUSA lack the host-mapped pinned-memory APIs (cudaHostAllocPortable
// / cudaHostAllocMapped / cudaHostGetDevicePointer) and __nanosleep that this
// implementation relies on, so the internal AllReduce is a CUDA-only feature.
// The dispatcher in ggml-cuda.cu treats a nullptr pipeline as "init failed"
// and silently falls back to the meta backend's generic AllReduce.
// MUSA lacks the host-mapped pinned-memory APIs (cudaHostAllocPortable
// / cudaHostAllocMapped / cudaHostGetDevicePointer) and a device-side
// sleep intrinsic that this implementation relies on, so the internal
// AllReduce is unavailable there. The dispatcher in ggml-cuda.cu treats
// a nullptr pipeline as "init failed" and silently falls back to the meta
// backend's generic AllReduce.
ggml_cuda_ar_pipeline * ggml_cuda_ar_pipeline_init(const int *, size_t) {
return nullptr;
}
@@ -968,4 +974,4 @@ bool ggml_cuda_ar_allreduce(ggml_cuda_ar_pipeline *, ggml_backend_t *, ggml_tens
return false;
}
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
#endif // !defined(GGML_USE_MUSA)
ggml/src/ggml-cuda/allreduce.cuh
+1 -1
View File
@@ -9,7 +9,7 @@
struct ggml_cuda_ar_pipeline;
// Allocate a pipeline for n_devices GPUs.
// devices[] holds the CUDA device IDs in rank order.
// devices[] holds the GPU device IDs in rank order.
// Returns nullptr on allocation failure.
ggml_cuda_ar_pipeline * ggml_cuda_ar_pipeline_init(
const int * devices, size_t n_devices);
ggml/src/ggml-cuda/cpy.cu
+8
View File
@@ -589,6 +589,14 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg
ggml_cpy_scalar_cuda<int32_t, int32_t>
(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
}
} else if (src0->type == GGML_TYPE_I16 && src1->type == GGML_TYPE_I16) {
if (can_be_transposed) {
ggml_cpy_scalar_cuda<int16_t, int16_t, true>
(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
} else {
ggml_cpy_scalar_cuda<int16_t, int16_t>
(src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream);
}
} else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_I32) {
if (contiguous_srcs) {
ggml_cpy_scalar_contiguous_cuda<float, int32_t>
ggml/src/ggml-cuda/fattn-mma-f16.cuh
+5 -3
View File
@@ -181,7 +181,7 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 8, 128, 1, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 16, 256, 2, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 32, 256, 2, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 64, 256, 4, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 64, 256, 3, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 8, 256, 2, 64, 40, 40, 40, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 16, 256, 2, 64, 40, 40, 40, 1, true);
@@ -1141,7 +1141,7 @@ template<int DV, int ncols> struct mma_tile_sizes {
using T_C_KQ = tile<16, 16, float>; // column-major
using T_A_VKQ = tile<16, 8, half2>; // row-major
using T_B_VKQ = tile<16, 8, half2>; // column-major
using T_C_VKQ = tile<16, 8, half2>; // column-major
using T_C_VKQ = tile<16, 16, float>; // column-major
};
#else // Volta
template<int DV, int ncols> struct mma_tile_sizes {
@@ -1227,7 +1227,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
T_C_VKQ VKQ_C[cols_per_warp == 8 ? DV/T_C_VKQ::I : DV/(2*T_C_VKQ::J)];
#elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
T_C_VKQ VKQ_C[DV % 32 != 0 ? DV/T_C_VKQ::J : DV/(2*T_C_VKQ::J)];
#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
#elif defined(AMD_MFMA_AVAILABLE)
T_C_VKQ VKQ_C[ DV/T_C_VKQ::J];
#elif defined(AMD_WMMA_AVAILABLE)
T_C_VKQ VKQ_C[ DV/(2*T_C_VKQ::J)];
#else // Volta
T_C_VKQ VKQ_C[ DV/(2*T_C_VKQ::J)];
ggml/src/ggml-cuda/ggml-cuda.cu
+3 -4
View File
@@ -5311,10 +5311,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
return false;
} break;
case GGML_OP_DUP:
{
ggml_type src0_type = op->src[0]->type;
return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16;
} break;
return true;
case GGML_OP_ARGMAX:
case GGML_OP_COUNT_EQUAL:
{
@@ -5482,7 +5479,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
return true;
#endif
case GGML_OP_SUM_ROWS:
return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && ggml_is_contiguous_rows(op->src[0]);
case GGML_OP_MEAN:
return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && ggml_is_contiguous_rows(op->src[0]);
case GGML_OP_GROUP_NORM:
return ggml_is_contiguous(op->src[0]);
case GGML_OP_PAD:
ggml/src/ggml-cuda/mean.cu
+14 -7
View File
@@ -18,7 +18,7 @@ void ggml_cuda_op_mean(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
GGML_ASSERT(src0->type == GGML_TYPE_F32);
GGML_ASSERT(dst->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_is_contiguous(src0));
GGML_ASSERT(ggml_is_contiguous_rows(src0));
const int64_t ncols = src0->ne[0];
const int64_t nrows = ggml_nrows(src0);
@@ -65,13 +65,20 @@ void ggml_cuda_op_mean(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
// Heuristic for block size selection to optimize occupancy.
// See discussion in: https://github.com/ggml-org/llama.cpp/pull/15132
dim3 block_dims;
if ((nrows / nsm) < 2) {
const dim3 block_dims(512, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
ggml_cuda_kernel_launch(reduce_rows_f32</*norm=*/true>, launch_params, src0_d, dst_d, ncols);
block_dims = dim3(512, 1, 1);
} else {
const dim3 block_dims(ncols < 1024 ? 32 : 128, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
ggml_cuda_kernel_launch(reduce_rows_f32</*norm=*/true>, launch_params, src0_d, dst_d, ncols);
block_dims = dim3(ncols < 1024 ? 32 : 128, 1, 1);
}
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
if (ggml_is_contiguous(src0)) {
ggml_cuda_kernel_launch(reduce_rows_f32</*norm=*/true>, launch_params, src0_d, dst_d, ncols);
return;
}
const char * src0_d_bytes = (const char *) src0->data;
ggml_cuda_kernel_launch(reduce_rows_f32_strided</*norm=*/true>, launch_params, src0_d_bytes, dst_d, ncols,
src0->ne[1], src0->ne[2], src0->nb[1], src0->nb[2], src0->nb[3]);
}
ggml/src/ggml-cuda/reduce_rows.cuh
+36 -7
View File
@@ -1,11 +1,6 @@
#include "common.cuh"
// Row reduction kernel template - compute sum (norm=false) or mean (norm=true)
template <bool norm>
static __global__ void reduce_rows_f32(const float * x_ptr, float * dst_ptr, const int ncols) {
const float * GGML_CUDA_RESTRICT x = x_ptr;
float * GGML_CUDA_RESTRICT dst = dst_ptr;
const int row = blockIdx.x;
static __device__ __forceinline__ float reduce_row_f32(const float * x, const int ncols) {
const int col = threadIdx.x;
float sum = 0.0f;
@@ -17,7 +12,7 @@ static __global__ void reduce_rows_f32(const float * x_ptr, float * dst_ptr, con
for (int i = col; i < ncols;) {
for (int j = 0; j < num_unroll; ++j) {
if (i < ncols) {
temp[j] = x[row * ncols + i];
temp[j] = x[i];
} else {
temp[j] = 0;
}
@@ -35,6 +30,40 @@ static __global__ void reduce_rows_f32(const float * x_ptr, float * dst_ptr, con
__shared__ float shared_vals[32];
sum = block_reduce<block_reduce_method::SUM>(sum, shared_vals);
return sum;
}
// Row reduction kernel template - compute sum (norm=false) or mean (norm=true)
template <bool norm>
static __global__ void reduce_rows_f32(const float * x_ptr, float * dst_ptr, const int ncols) {
float * GGML_CUDA_RESTRICT dst = dst_ptr;
const int64_t row = blockIdx.x;
const int col = threadIdx.x;
const float * GGML_CUDA_RESTRICT x = x_ptr + row*ncols;
const float sum = reduce_row_f32(x, ncols);
if (col != 0) {
return;
}
dst[row] = norm ? sum / ncols : sum;
}
template <bool norm>
static __global__ void reduce_rows_f32_strided(const char * x_ptr, float * dst_ptr, const int ncols,
const int64_t ne1, const int64_t ne2, const int64_t nb1, const int64_t nb2, const int64_t nb3) {
float * GGML_CUDA_RESTRICT dst = dst_ptr;
const int64_t row = blockIdx.x;
const int col = threadIdx.x;
const int64_t i1 = row % ne1;
const int64_t i2 = (row / ne1) % ne2;
const int64_t i3 = row / (ne1 * ne2);
const float * GGML_CUDA_RESTRICT x = (const float *) (x_ptr + i1*nb1 + i2*nb2 + i3*nb3);
const float sum = reduce_row_f32(x, ncols);
if (col != 0) {
return;
}
ggml/src/ggml-cuda/sumrows.cu
+13 -7
View File
@@ -24,24 +24,30 @@ void ggml_cuda_op_sum_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
GGML_ASSERT(src0->type == GGML_TYPE_F32);
GGML_ASSERT( dst->type == GGML_TYPE_F32);
GGML_ASSERT(ggml_is_contiguous(src0));
GGML_ASSERT(ggml_is_contiguous_rows(src0));
const int64_t ncols = src0->ne[0];
const int64_t nrows = ggml_nrows(src0);
if (ggml_is_contiguous(src0)) {
sum_rows_f32_cuda(src0_d, dst_d, ncols, nrows, stream);
return;
}
const dim3 block_nums(nrows, 1, 1);
const int id = ggml_cuda_get_device();
const int nsm = ggml_cuda_info().devices[id].nsm;
dim3 block_dims;
if ((nrows / nsm) < 2) {
// Increase num threads to 512 for small nrows to better hide the latency
const dim3 block_dims(512, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
ggml_cuda_kernel_launch(reduce_rows_f32</*norm=*/false>, launch_params, src0_d, dst_d, ncols);
block_dims = dim3(512, 1, 1);
} else {
// Enough active SMs to hide latency, use smaller blocks to allow better scheduling
const dim3 block_dims(ncols < 1024 ? 32 : 128, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
ggml_cuda_kernel_launch(reduce_rows_f32</*norm=*/false>, launch_params, src0_d, dst_d, ncols);
block_dims = dim3(ncols < 1024 ? 32 : 128, 1, 1);
}
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream);
const char * src0_d_bytes = (const char *) src0->data;
ggml_cuda_kernel_launch(reduce_rows_f32_strided</*norm=*/false>, launch_params, src0_d_bytes, dst_d, ncols,
src0->ne[1], src0->ne[2], src0->nb[1], src0->nb[2], src0->nb[3]);
}
ggml/src/ggml-cuda/vendors/hip.h
Vendored
+4
View File
@@ -73,6 +73,10 @@
#define cudaGetDeviceProperties hipGetDeviceProperties
#define cudaGetErrorString hipGetErrorString
#define cudaGetLastError hipGetLastError
#define cudaHostAlloc hipHostMalloc
#define cudaHostAllocPortable hipHostMallocPortable
#define cudaHostAllocMapped hipHostMallocMapped
#define cudaHostGetDevicePointer hipHostGetDevicePointer
#define cudaHostRegister hipHostRegister
#define cudaHostRegisterPortable hipHostRegisterPortable
#define cudaHostRegisterReadOnly hipHostRegisterReadOnly
ggml/src/ggml-metal/ggml-metal-ops.cpp
+1
View File
@@ -2923,6 +2923,7 @@ static int ggml_metal_op_flash_attn_ext_n_kv_max_sparse(const ggml_tensor * op)
const bool dk_dv_ok = (dk == 32 && dv == 32) ||
(dk == 64 && dv == 64) ||
(dk == 96 && dv == 96) ||
(dk == 96 && dv == 64) ||
(dk == 128 && dv == 128) ||
(dk == 192 && dv == 128) ||
(dk == 192 && dv == 192) ||
ggml/src/ggml-metal/ggml-metal-tuning.cpp
+3
View File
@@ -34,6 +34,9 @@ int fa_vec_baseline_ne(int dk, int dv) {
if (dk == 96 && dv == 96) {
return 4;
}
if (dk == 96 && dv == 64) {
return 4;
}
if (dk == 128 && dv == 128) {
return 1;
}
ggml/src/ggml-metal/kernels/fa.metal
+31
View File
@@ -930,6 +930,7 @@ template [[host_name("kernel_flash_attn_ext_f32_dk64_dv64" )]] kernel flash_at
template [[host_name("kernel_flash_attn_ext_f32_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 72, 72>;
template [[host_name("kernel_flash_attn_ext_f32_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 80, 80>;
template [[host_name("kernel_flash_attn_ext_f32_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 96, 96>;
template [[host_name("kernel_flash_attn_ext_f32_dk96_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 96, 64>;
template [[host_name("kernel_flash_attn_ext_f32_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 112, 112>;
template [[host_name("kernel_flash_attn_ext_f32_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 128, 128>;
template [[host_name("kernel_flash_attn_ext_f32_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_F32, float4x4, 1, dequantize_f32, float4x4, 1, dequantize_f32, 192, 192>;
@@ -946,6 +947,7 @@ template [[host_name("kernel_flash_attn_ext_f16_dk64_dv64" )]] kernel flash_at
template [[host_name("kernel_flash_attn_ext_f16_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 72, 72>;
template [[host_name("kernel_flash_attn_ext_f16_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 80, 80>;
template [[host_name("kernel_flash_attn_ext_f16_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 96, 96>;
template [[host_name("kernel_flash_attn_ext_f16_dk96_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 96, 64>;
template [[host_name("kernel_flash_attn_ext_f16_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 112, 112>;
template [[host_name("kernel_flash_attn_ext_f16_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 128, 128>;
template [[host_name("kernel_flash_attn_ext_f16_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, half4x4, 1, dequantize_f16, half4x4, 1, dequantize_f16, 192, 192>;
@@ -963,6 +965,7 @@ template [[host_name("kernel_flash_attn_ext_bf16_dk64_dv64" )]] kernel flash_at
template [[host_name("kernel_flash_attn_ext_bf16_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 72, 72>;
template [[host_name("kernel_flash_attn_ext_bf16_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 80, 80>;
template [[host_name("kernel_flash_attn_ext_bf16_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 96, 96>;
template [[host_name("kernel_flash_attn_ext_bf16_dk96_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 96, 64>;
template [[host_name("kernel_flash_attn_ext_bf16_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 112, 112>;
template [[host_name("kernel_flash_attn_ext_bf16_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 128, 128>;
template [[host_name("kernel_flash_attn_ext_bf16_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES_BF, bfloat4x4, 1, dequantize_bf16, bfloat4x4, 1, dequantize_bf16, 192, 192>;
@@ -980,6 +983,7 @@ template [[host_name("kernel_flash_attn_ext_q4_0_dk64_dv64" )]] kernel flash_at
template [[host_name("kernel_flash_attn_ext_q4_0_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 72, 72>;
template [[host_name("kernel_flash_attn_ext_q4_0_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 80, 80>;
template [[host_name("kernel_flash_attn_ext_q4_0_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 96, 96>;
template [[host_name("kernel_flash_attn_ext_q4_0_dk96_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 96, 64>;
template [[host_name("kernel_flash_attn_ext_q4_0_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 112, 112>;
template [[host_name("kernel_flash_attn_ext_q4_0_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 128, 128>;
template [[host_name("kernel_flash_attn_ext_q4_0_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_0, 2, dequantize_q4_0, block_q4_0, 2, dequantize_q4_0, 192, 192>;
@@ -996,6 +1000,7 @@ template [[host_name("kernel_flash_attn_ext_q4_1_dk64_dv64" )]] kernel flash_at
template [[host_name("kernel_flash_attn_ext_q4_1_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 72, 72>;
template [[host_name("kernel_flash_attn_ext_q4_1_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 80, 80>;
template [[host_name("kernel_flash_attn_ext_q4_1_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 96, 96>;
template [[host_name("kernel_flash_attn_ext_q4_1_dk96_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 96, 64>;
template [[host_name("kernel_flash_attn_ext_q4_1_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 112, 112>;
template [[host_name("kernel_flash_attn_ext_q4_1_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 128, 128>;
template [[host_name("kernel_flash_attn_ext_q4_1_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q4_1, 2, dequantize_q4_1, block_q4_1, 2, dequantize_q4_1, 192, 192>;
@@ -1012,6 +1017,7 @@ template [[host_name("kernel_flash_attn_ext_q5_0_dk64_dv64" )]] kernel flash_at
template [[host_name("kernel_flash_attn_ext_q5_0_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 72, 72>;
template [[host_name("kernel_flash_attn_ext_q5_0_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 80, 80>;
template [[host_name("kernel_flash_attn_ext_q5_0_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 96, 96>;
template [[host_name("kernel_flash_attn_ext_q5_0_dk96_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 96, 64>;
template [[host_name("kernel_flash_attn_ext_q5_0_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 112, 112>;
template [[host_name("kernel_flash_attn_ext_q5_0_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 128, 128>;
template [[host_name("kernel_flash_attn_ext_q5_0_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_0, 2, dequantize_q5_0, block_q5_0, 2, dequantize_q5_0, 192, 192>;
@@ -1028,6 +1034,7 @@ template [[host_name("kernel_flash_attn_ext_q5_1_dk64_dv64" )]] kernel flash_at
template [[host_name("kernel_flash_attn_ext_q5_1_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 72, 72>;
template [[host_name("kernel_flash_attn_ext_q5_1_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 80, 80>;
template [[host_name("kernel_flash_attn_ext_q5_1_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 96, 96>;
template [[host_name("kernel_flash_attn_ext_q5_1_dk96_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 96, 64>;
template [[host_name("kernel_flash_attn_ext_q5_1_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 112, 112>;
template [[host_name("kernel_flash_attn_ext_q5_1_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 128, 128>;
template [[host_name("kernel_flash_attn_ext_q5_1_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q5_1, 2, dequantize_q5_1, block_q5_1, 2, dequantize_q5_1, 192, 192>;
@@ -1044,6 +1051,7 @@ template [[host_name("kernel_flash_attn_ext_q8_0_dk64_dv64" )]] kernel flash_at
template [[host_name("kernel_flash_attn_ext_q8_0_dk72_dv72" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 72, 72>;
template [[host_name("kernel_flash_attn_ext_q8_0_dk80_dv80" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 80, 80>;
template [[host_name("kernel_flash_attn_ext_q8_0_dk96_dv96" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 96, 96>;
template [[host_name("kernel_flash_attn_ext_q8_0_dk96_dv64" )]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 96, 64>;
template [[host_name("kernel_flash_attn_ext_q8_0_dk112_dv112")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 112, 112>;
template [[host_name("kernel_flash_attn_ext_q8_0_dk128_dv128")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 128, 128>;
template [[host_name("kernel_flash_attn_ext_q8_0_dk192_dv192")]] kernel flash_attn_ext_t kernel_flash_attn_ext<FA_TYPES, block_q8_0, 2, dequantize_q8_0, block_q8_0, 2, dequantize_q8_0, 192, 192>;
@@ -1905,6 +1913,29 @@ template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv96")]] kernel flas
template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv96_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 96, 96, 4, 2>;
template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv96_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 96, 96, 4, 4>;
template [[host_name("kernel_flash_attn_ext_vec_f32_dk96_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 96, 64, 4>;
template [[host_name("kernel_flash_attn_ext_vec_f16_dk96_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 96, 64, 4>;
template [[host_name("kernel_flash_attn_ext_vec_f16_dk96_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 96, 64, 4, 2>;
template [[host_name("kernel_flash_attn_ext_vec_f16_dk96_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 96, 64, 4, 4>;
#if defined(GGML_METAL_HAS_BF16)
template [[host_name("kernel_flash_attn_ext_vec_bf16_dk96_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, bfloat4, 1, dequantize_bf16_t4, bfloat4, 1, dequantize_bf16_t4, 96, 64, 4>;
#endif
template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk96_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 96, 64, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk96_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 96, 64, 4, 2>;
template [[host_name("kernel_flash_attn_ext_vec_q4_0_dk96_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_0, 8, dequantize_q4_0_t4, block_q4_0, 8, dequantize_q4_0_t4, 96, 64, 4, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk96_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 96, 64, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk96_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 96, 64, 4, 2>;
template [[host_name("kernel_flash_attn_ext_vec_q4_1_dk96_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q4_1, 8, dequantize_q4_1_t4, block_q4_1, 8, dequantize_q4_1_t4, 96, 64, 4, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk96_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 96, 64, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk96_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 96, 64, 4, 2>;
template [[host_name("kernel_flash_attn_ext_vec_q5_0_dk96_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_0, 8, dequantize_q5_0_t4, block_q5_0, 8, dequantize_q5_0_t4, 96, 64, 4, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk96_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 96, 64, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk96_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 96, 64, 4, 2>;
template [[host_name("kernel_flash_attn_ext_vec_q5_1_dk96_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q5_1, 8, dequantize_q5_1_t4, block_q5_1, 8, dequantize_q5_1_t4, 96, 64, 4, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv64")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 96, 64, 4>;
template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv64_q2_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 96, 64, 4, 2>;
template [[host_name("kernel_flash_attn_ext_vec_q8_0_dk96_dv64_q4_ne4")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, block_q8_0, 8, dequantize_q8_0_t4, block_q8_0, 8, dequantize_q8_0_t4, 96, 64, 4, 4>;
template [[host_name("kernel_flash_attn_ext_vec_f32_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES_F32, float4, 1, dequantize_f32_t4, float4, 1, dequantize_f32_t4, 128, 128, 1>;
template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 1>;
template [[host_name("kernel_flash_attn_ext_vec_f16_dk128_dv128_q1_ne2")]] kernel flash_attn_ext_vec_t kernel_flash_attn_ext_vec<FA_TYPES, half4, 1, dequantize_f16_t4, half4, 1, dequantize_f16_t4, 128, 128, 2, 1>;
ggml/src/ggml-rpc/ggml-rpc.cpp
+44 -25
View File
@@ -700,10 +700,31 @@ static void ggml_backend_rpc_buffer_memset_tensor(
ctx->dispatcher->send(RPC_CMD_MEMSET_TENSOR, request, sizeof(*request));
}
// input serialization format: | rpc_tensor | cache_flag (1 byte) | offset (8 bytes) | data (size bytes)
static std::shared_ptr<uint8_t> serialize_set_tensor(const rpc_tensor & rpc_tensor, uint8_t cache_flag, uint64_t offset, const void * data, size_t size, size_t & input_size) {
input_size = sizeof(rpc_tensor) + sizeof(cache_flag) + sizeof(offset) + size;
uint8_t * input = new uint8_t[input_size]();
uint8_t * p = input;
memcpy(p, &rpc_tensor, sizeof(rpc_tensor)); p += sizeof(rpc_tensor);
memcpy(p, &cache_flag, sizeof(cache_flag)); p += sizeof(cache_flag);
memcpy(p, &offset, sizeof(offset)); p += sizeof(offset);
memcpy(p, data, size);
return std::shared_ptr<uint8_t>(input, std::default_delete<uint8_t[]>());
}
// the hash cache is meant for weights, so that a model reload can skip re-sending them.
// compute-buffer inputs (the activations ggml_backend_sched copies between backends) must not
// take this path, otherwise with `rpc-server -c` every ubatch above the threshold is written
// to the cache directory and later served from there.
static bool rpc_use_hash_cache(const ggml_tensor * tensor, size_t size) {
return size > HASH_THRESHOLD && tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS;
}
static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context;
rpc_tensor rpc_tensor = serialize_tensor(tensor);
if (size > HASH_THRESHOLD) {
uint8_t cache_flag = 0;
if (rpc_use_hash_cache(tensor, size)) {
auto request = std::make_shared<rpc_msg_set_tensor_hash_req>();
request->tensor = rpc_tensor;
request->offset = offset;
@@ -714,15 +735,12 @@ static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggm
// the server has the same data, no need to send it
return;
}
// the server has no cache entry for this tensor - ask it to save one
cache_flag = 1;
}
// input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes)
size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size;
uint8_t * input = new uint8_t[input_size]();
memcpy(input, &rpc_tensor, sizeof(rpc_tensor));
memcpy(input + sizeof(rpc_tensor), &offset, sizeof(offset));
memcpy(input + sizeof(rpc_tensor) + sizeof(offset), data, size);
std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
ctx->dispatcher->send(RPC_CMD_SET_TENSOR, input_ptr, input_size);
size_t input_size;
auto input = serialize_set_tensor(rpc_tensor, cache_flag, offset, data, size, input_size);
ctx->dispatcher->send(RPC_CMD_SET_TENSOR, input, input_size);
}
static void ggml_backend_rpc_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
@@ -930,7 +948,8 @@ static void ggml_backend_rpc_free(ggml_backend_t backend) {
static void ggml_backend_rpc_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
ggml_backend_rpc_context * ctx = (ggml_backend_rpc_context *)backend->context;
rpc_tensor rpc_tensor = serialize_tensor(tensor);
if (size > HASH_THRESHOLD) {
uint8_t cache_flag = 0;
if (rpc_use_hash_cache(tensor, size)) {
auto request = std::make_shared<rpc_msg_set_tensor_hash_req>();
request->tensor = rpc_tensor;
request->offset = offset;
@@ -942,15 +961,12 @@ static void ggml_backend_rpc_set_tensor_async(ggml_backend_t backend, ggml_tenso
// the server has the same data, no need to send it
return;
}
// the server has no cache entry for this tensor - ask it to save one
cache_flag = 1;
}
// input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes)
size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size;
uint8_t * input = new uint8_t[input_size]();
memcpy(input, &rpc_tensor, sizeof(rpc_tensor));
memcpy(input + sizeof(rpc_tensor), &offset, sizeof(offset));
memcpy(input + sizeof(rpc_tensor) + sizeof(offset), data, size);
std::shared_ptr<uint8_t> input_ptr(input, std::default_delete<uint8_t[]>());
ctx->dispatcher->send_async(RPC_CMD_SET_TENSOR, input_ptr, input_size);
size_t input_size;
auto input = serialize_set_tensor(rpc_tensor, cache_flag, offset, data, size, input_size);
ctx->dispatcher->send_async(RPC_CMD_SET_TENSOR, input, input_size);
}
static void ggml_backend_rpc_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) {
@@ -1401,14 +1417,17 @@ ggml_tensor * rpc_server::deserialize_tensor(struct ggml_context * ctx, const rp
bool rpc_server::set_tensor(const std::vector<uint8_t> & input) {
// serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes) |
if (input.size() < sizeof(rpc_tensor) + sizeof(uint64_t)) {
// serialization format: | rpc_tensor | cache_flag (1 byte) | offset (8 bytes) | data (size bytes) |
uint8_t cache_flag;
uint64_t offset;
const size_t header_size = sizeof(rpc_tensor) + sizeof(cache_flag) + sizeof(offset);
if (input.size() < header_size) {
return false;
}
const rpc_tensor * in_tensor = (const rpc_tensor *)input.data();
uint64_t offset;
memcpy(&offset, input.data() + sizeof(rpc_tensor), sizeof(offset));
const size_t size = input.size() - sizeof(rpc_tensor) - sizeof(offset);
memcpy(&cache_flag, input.data() + sizeof(rpc_tensor), sizeof(cache_flag));
memcpy(&offset, input.data() + sizeof(rpc_tensor) + sizeof(cache_flag), sizeof(offset));
const size_t size = input.size() - header_size;
struct ggml_init_params params {
/*.mem_size =*/ ggml_tensor_overhead(),
@@ -1437,8 +1456,8 @@ bool rpc_server::set_tensor(const std::vector<uint8_t> & input) {
}
}
const void * data = input.data() + sizeof(rpc_tensor) + sizeof(offset);
if (cache_dir && size > HASH_THRESHOLD) {
const void * data = input.data() + header_size;
if (cache_dir && cache_flag) {
uint64_t hash = fnv_hash((const uint8_t*)data, size);
char hash_str[17];
snprintf(hash_str, sizeof(hash_str), "%016" PRIx64, hash);
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+118 -24
View File
@@ -1177,6 +1177,10 @@ struct vk_device_struct {
std::map<std::pair<uint32_t, uint32_t>, vk_pipeline> pipeline_fa_mask_opt;
vk_pipeline pipeline_fa_sparse_compact;
vk_pipeline pipeline_fa_sparse_compact_subgroup;
bool fa_sparse_compact_use_subgroups;
vk_pipeline pipeline_flash_attn_split_k_reduce;
vk_pipeline pipeline_count_experts;
@@ -2202,6 +2206,16 @@ struct vk_op_flash_attn_mask_opt_push_constants {
uint32_t nbd3;
};
struct vk_op_flash_attn_sparse_compact_push_constants {
uint32_t KV;
uint32_t nem1;
uint32_t nem2;
uint32_t nbm1;
uint32_t nbm2;
uint32_t nbm3;
uint32_t n_kv_max;
};
// Allow pre-recording command buffers
struct vk_staging_memcpy {
vk_staging_memcpy(void * _dst, const void * _src, size_t _n) : dst(_dst), src(_src), n(_n) {}
@@ -4125,14 +4139,15 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_
}
static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool aligned, bool f32acc,
bool use_mask, bool use_mask_opt, bool use_logit_softcap, ggml_type k_type, ggml_type v_type) {
bool use_mask, bool use_mask_opt, bool use_logit_softcap, bool use_sparse, ggml_type k_type, ggml_type v_type) {
const bool old_amd_windows = device->vendor_id == VK_VENDOR_ID_AMD && device->driver_id == vk::DriverId::eAmdProprietary &&
(device->architecture == AMD_GCN || device->architecture == AMD_RDNA1 || device->architecture == AMD_RDNA2);
uint32_t flags = (use_mask_opt ? 1 : 0) |
(use_mask ? 2 : 0) |
(use_logit_softcap ? 4 : 0) |
(old_amd_windows ? 8 : 0);
(old_amd_windows ? 8 : 0) |
(use_sparse ? 16 : 0);
const uint32_t subgroup_size = params.disable_subgroups ? 0 : params.subgroup_size;
@@ -4752,7 +4767,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
}
name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16";
}
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7,
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 8,
sizeof(vk_flash_attn_push_constants), {Br, 1, 1},
get_fa_spec_constants(fa.first), aligned ? Bc : 1, true,
!fa_ds, !fa_ds ? fa_sgs : 0);
@@ -4788,7 +4803,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
else { spv_data = flash_attn_f32_f16_f16acc_cm1_data; spv_size = flash_attn_f32_f16_f16acc_cm1_len; }
name = aligned ? "flash_attn_f32_f16_aligned_cm1" : "flash_attn_f32_f16_cm1";
}
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7,
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 8,
sizeof(vk_flash_attn_push_constants), {Br, 1, 1},
get_fa_spec_constants(fa.first), aligned ? Bc : 1, true,
!fa_ds, !fa_ds ? fa_sgs : 0);
@@ -4825,7 +4840,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
if (f32acc) { spv_data = flash_attn_f32_f16_cm2_data; spv_size = flash_attn_f32_f16_cm2_len; name = "flash_attn_f32_f16_f32acc_cm2"; }
else { spv_data = flash_attn_f32_f16_f16acc_cm2_data; spv_size = flash_attn_f32_f16_f16acc_cm2_len; name = "flash_attn_f32_f16_f16acc_cm2"; }
}
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 7,
ggml_vk_create_pipeline(device, fa.second, name, spv_size, spv_data, "main", 8,
sizeof(vk_flash_attn_push_constants), {Br, 1, 1},
get_fa_spec_constants(fa.first), aligned ? Bc : 1, true, false, 0);
}
@@ -5789,6 +5804,22 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, it.second, "fa_mask_opt", fa_mask_opt_len, fa_mask_opt_data, "main", 2, sizeof(vk_op_flash_attn_mask_opt_push_constants), {1, 1, 1}, {128, 128 / device->subgroup_size, BrBc.first, BrBc.second}, 1, true, true, device->subgroup_size);
}
{
// Large workgroup so the per-row KV scan parallelizes; capped to device limits.
const uint32_t compact_max = std::min({1024u, device->properties.limits.maxComputeWorkGroupInvocations, device->properties.limits.maxComputeWorkGroupSize[0]});
// Fast ballot prefix-sum path when the device supports full subgroups; otherwise
// a shared-memory prefix-sum fallback. Both emit a deterministic ascending list.
device->fa_sparse_compact_use_subgroups = device->subgroup_ballot && device->subgroup_require_full_support;
if (device->fa_sparse_compact_use_subgroups) {
const uint32_t compact_wg = std::max(device->subgroup_size, (compact_max / device->subgroup_size) * device->subgroup_size);
const uint32_t compact_num_sg = compact_wg / device->subgroup_size;
ggml_vk_create_pipeline(device, device->pipeline_fa_sparse_compact_subgroup, "fa_sparse_compact_subgroup", fa_sparse_compact_subgroup_len, fa_sparse_compact_subgroup_data, "main", 2, sizeof(vk_op_flash_attn_sparse_compact_push_constants), {1, 1, 1}, {compact_wg, compact_num_sg}, 1, true, true, device->subgroup_size);
} else {
ggml_vk_create_pipeline(device, device->pipeline_fa_sparse_compact, "fa_sparse_compact", fa_sparse_compact_len, fa_sparse_compact_data, "main", 2, sizeof(vk_op_flash_attn_sparse_compact_push_constants), {1, 1, 1}, {compact_max}, 1, true);
}
}
if (device->subgroup_clustered && device->subgroup_require_full_support) {
ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4, "quantize_q8_1_x4", quantize_q8_1_x4_subgroup_len, quantize_q8_1_x4_subgroup_data, "main", 2, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1, true, true);
} else {
@@ -11310,6 +11341,30 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc);
float scale = 1.0f;
float max_bias = 0.0f;
float logit_softcap = 0.0f;
memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float));
memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float));
memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
if (logit_softcap != 0) {
scale /= logit_softcap;
}
// Sparse mask hint (op_params[4]): compact the <= n_kv_max finite positions and gather only those.
const int32_t n_kv_max = mask ? ggml_get_op_params_i32(dst, 4) : 0;
static const bool disable_sparse = getenv("GGML_VK_FA_SPARSE_DISABLE") != nullptr;
// cm2 dense is fast, so it needs a larger reduction to win.
const int64_t min_ratio = tuning_params.path == FA_COOPMAT2 ? 4 : 2;
const bool use_sparse = !disable_sparse && n_kv_max > 0 && mask &&
max_bias == 0.0f && logit_softcap == 0.0f &&
k_type_eff == GGML_TYPE_F16 && v_type_eff == GGML_TYPE_F16 &&
nem0 == KV &&
(int64_t)KV >= std::max<int64_t>(4096, min_ratio * (int64_t)n_kv_max) &&
(gqa_ratio > 1 || (tuning_params.path == FA_SCALAR && N == 1));
const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type));
uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type));
uint32_t v_stride = (uint32_t)(nbv1 / ggml_type_size(v->type));
@@ -11332,7 +11387,6 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t));
nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t));
}
const uint32_t alignment = tuning_params.block_cols;
bool aligned = (KV % alignment) == 0 &&
// the "aligned" shader variant will forcibly align strides, for performance
@@ -11343,23 +11397,11 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
aligned = false;
}
float scale = 1.0f;
float max_bias = 0.0f;
float logit_softcap = 0.0f;
memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float));
memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float));
memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
if (logit_softcap != 0) {
scale /= logit_softcap;
}
// Only use mask opt when the mask is fairly large. This hasn't been tuned extensively.
bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16
bool use_mask_opt = mask && !use_sparse && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16
&& (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256);
vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc,
mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff);
mask != nullptr, use_mask_opt, logit_softcap != 0, use_sparse, k_type_eff, v_type_eff);
vk_pipeline pipeline = nullptr;
@@ -11394,7 +11436,19 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
const uint32_t Tr = CEIL_DIV(N, Br);
// Try to use split_k when KV is large enough to be worth the overhead.
if (gqa_ratio > 1 && workgroups_x <= Br) {
// Sparse: split_kv carries n_kv_max, split_k partitions its blocks for occupancy.
if (use_sparse) {
split_kv = (uint32_t)n_kv_max;
const uint32_t total_blocks = CEIL_DIV((uint32_t)n_kv_max, Bc);
const uint32_t base_wgs = (gqa_ratio > 1 ? workgroups_x : Tr) * workgroups_y * workgroups_z;
if (base_wgs < shader_core_count * 2) {
split_k = shader_core_count * 2 / base_wgs;
}
split_k = std::max(1u, std::min(split_k, total_blocks));
// Match the shader's per-split block count so no split is empty.
const uint32_t per_blocks = CEIL_DIV(total_blocks, split_k);
split_k = CEIL_DIV(total_blocks, per_blocks);
} else if (gqa_ratio > 1 && workgroups_x <= Br) {
split_k = shader_core_count * 2 / (workgroups_x * workgroups_y * workgroups_z);
} else if (gqa_ratio <= 1) {
uint32_t total_wgs_no_split = Tr * workgroups_y * workgroups_z;
@@ -11403,7 +11457,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
}
}
if (split_k > 1) {
if (!use_sparse && split_k > 1) {
// Try to evenly split KV into split_k chunks, but it needs to be a multiple
// of "align", so recompute split_k based on that.
split_kv = ROUNDUP_POW2(std::max(1u, KV / split_k), alignment);
@@ -11450,6 +11504,24 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
}
}
// Sparse index scratch reuses prealloc_y (mutually exclusive with mask opt).
const uint64_t sparse_idx_size = use_sparse
? sizeof(int32_t) * (uint64_t)n_kv_max * nem1 * nem2 * nem3
: 0;
vk_pipeline sparse_compact_pipeline = ctx->device->fa_sparse_compact_use_subgroups
? ctx->device->pipeline_fa_sparse_compact_subgroup
: ctx->device->pipeline_fa_sparse_compact;
if (use_sparse) {
ggml_pipeline_request_descriptor_sets(ctx, sparse_compact_pipeline, 1);
if (ctx->prealloc_size_y < sparse_idx_size) {
ctx->prealloc_size_y = sparse_idx_size;
ggml_vk_preallocate_buffers(ctx, subctx);
}
if (ctx->prealloc_y_need_sync) {
ggml_vk_sync_buffers(ctx, subctx);
}
}
const uint32_t n_head_kv = neq2;
const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv));
const float m0 = powf(2.0f, -(max_bias ) / n_head_log2);
@@ -11462,6 +11534,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
vk_subbuffer mask_buf = mask ? ggml_vk_tensor_subbuffer(ctx, mask) : q_buf;
vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf;
vk_subbuffer mask_opt_buf = use_mask_opt ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf;
vk_subbuffer sparse_buf = use_sparse ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf;
if (use_dequant_kv) {
const uint64_t fp = sizeof(ggml_fp16_t);
@@ -11513,6 +11586,24 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
ggml_vk_sync_buffers(ctx, subctx);
}
if (use_sparse)
{
const vk_op_flash_attn_sparse_compact_push_constants sc_pc = {
KV,
nem1,
nem2,
(uint32_t)(mask->nb[1] / sizeof(ggml_fp16_t)),
(uint32_t)(mask->nb[2] / sizeof(ggml_fp16_t)),
(uint32_t)(mask->nb[3] / sizeof(ggml_fp16_t)),
(uint32_t)n_kv_max,
};
ggml_vk_dispatch_pipeline(ctx, subctx, sparse_compact_pipeline,
{ mask_buf, sparse_buf }, sc_pc,
{ nem1, nem2, nem3 });
ggml_vk_sync_buffers(ctx, subctx);
}
const vk_flash_attn_push_constants pc = { N, KV,
(uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3,
(uint32_t)neq2, (uint32_t)neq3,
@@ -11545,7 +11636,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
vk_subbuffer split_k_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0);
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
{q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, mask_opt_buf},
{q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, mask_opt_buf, sparse_buf},
pc, { dispatch_x, workgroups_y, workgroups_z });
ggml_vk_sync_buffers(ctx, subctx);
@@ -11560,13 +11651,16 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
workgroups_x *= pipeline->wg_denoms[0];
}
ggml_vk_dispatch_pipeline(ctx, subctx, pipeline,
{q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf},
{q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf, sparse_buf},
pc, { workgroups_x, workgroups_y, workgroups_z });
}
if (use_dequant_kv) {
ctx->prealloc_x_need_sync = true;
}
if (use_mask_opt || use_sparse) {
ctx->prealloc_y_need_sync = true;
}
}
static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, uint32_t K, uint32_t NPQ) {
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
+28 -18
View File
@@ -218,12 +218,14 @@ void main() {
uint32_t c = (idx + tid) % Bc;
uint32_t r = (idx + tid) / Bc;
if (idx + tid < Bc * Br) {
if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r < p.nem1)) {
FLOAT_TYPE m = FLOAT_TYPE(data_m[m_offset + (i * Br + r) * m_stride + (j * Bc + c)]);
uint32_t kcol;
bool kv_active = fa_kv_index(j * Bc + c, kcol);
if (kv_active && (!nem1_bounds_check || i * Br + r < p.nem1)) {
FLOAT_TYPE m = FLOAT_TYPE(data_m[m_offset + (i * Br + r) * m_stride + kcol]);
masksh[c * masksh_stride + r] = m;
max_mask = max(max_mask, float(m));
} else {
masksh[c * masksh_stride + r] = FLOAT_TYPE(0);
masksh[c * masksh_stride + r] = USE_SPARSE ? FLOAT_TYPE(NEG_FLT_MAX_OVER_2) : FLOAT_TYPE(0);
}
}
}
@@ -258,14 +260,15 @@ void main() {
uint32_t c = (idx + tid) / (HSK / 4);
if (idx + gl_WorkGroupSize.x <= Bc * HSK / 4 || c < Bc) {
FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0);
if (!KV_bounds_check || j * Bc + c < KV) {
uint32_t kcol;
if (fa_kv_index(j * Bc + c, kcol)) {
if (USE_DECODE_K) {
uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d;
uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * d;
uint ib = coord / BLOCK_SIZE_K;
uint iqs = (coord % BLOCK_SIZE_K);
K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K);
} else {
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]);
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d]);
}
}
@@ -305,7 +308,9 @@ void main() {
}
[[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
uint32_t kcol;
bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, kcol);
if (!kv_active) {
continue;
}
@@ -313,12 +318,12 @@ void main() {
if (SHMEM_STAGING != 0) {
K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)];
} else if (USE_DECODE_K) {
uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid);
uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid);
uint ib = coord / BLOCK_SIZE_K;
uint iqs = (coord % BLOCK_SIZE_K);
K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K);
} else {
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]);
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * D_split + d_tid]);
}
[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
Sf[r][c] = dot_product(Q_cache[r], K_Tf, Sf[r][c]);
@@ -327,7 +332,9 @@ void main() {
}
} else {
[[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
uint32_t kcol;
bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, kcol);
if (!kv_active) {
continue;
}
@@ -336,12 +343,12 @@ void main() {
if (SHMEM_STAGING != 0) {
K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)];
} else if (USE_DECODE_K) {
uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid);
uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid);
uint ib = coord / BLOCK_SIZE_K;
uint iqs = (coord % BLOCK_SIZE_K);
K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K);
} else {
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]);
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * D_split + d_tid]);
}
[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
Sf[r][c] = dot_product(Qf[tile_row(r) * qf_stride + d * D_split + d_tid], K_Tf, Sf[r][c]);
@@ -489,14 +496,15 @@ void main() {
uint32_t c = (idx + tid) / (HSV / 4);
if (idx + gl_WorkGroupSize.x <= Bc * HSV / 4 || c < Bc) {
FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0);
if (!KV_bounds_check || j * Bc + c < KV) {
uint32_t vcol;
if (fa_kv_index(j * Bc + c, vcol)) {
if (USE_DECODE_V) {
uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d;
uint coord = vcol * v_stride * BLOCK_SIZE_V + 4 * d;
uint ib = coord / BLOCK_SIZE_V;
uint iqs = (coord % BLOCK_SIZE_V);
V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V);
} else {
V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]);
V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + vcol * v_stride / 4 + d]);
}
}
@@ -507,7 +515,9 @@ void main() {
}
[[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
uint32_t vcol;
bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, vcol);
if (!kv_active) {
continue;
}
@@ -522,12 +532,12 @@ void main() {
if (SHMEM_STAGING != 0) {
Vf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)];
} else if (USE_DECODE_V) {
uint coord = (j * Bc + c * cols_per_iter + col_tid) * v_stride * BLOCK_SIZE_V + 4 * (d * D_split + d_tid);
uint coord = vcol * v_stride * BLOCK_SIZE_V + 4 * (d * D_split + d_tid);
uint ib = coord / BLOCK_SIZE_V;
uint iqs = (coord % BLOCK_SIZE_V);
Vf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V);
} else {
Vf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * v_stride / 4 + d * D_split + d_tid]);
Vf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + vcol * v_stride / 4 + d * D_split + d_tid]);
}
[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
Of[r][d] += FLOAT_TYPEV4(Pf[r] * Vf);
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
+32 -1
View File
@@ -24,6 +24,8 @@ const bool USE_MASK_OPT = (Flags & 1) != 0;
const bool MASK_ENABLE = (Flags & 2) != 0;
const bool LOGIT_SOFTCAP = (Flags & 4) != 0;
const bool OLD_AMD_WINDOWS = (Flags & 8) != 0;
// Sparse: gather binding-7 indices instead of scanning [0,KV); p.split_kv = n_kv_max.
const bool USE_SPARSE = (Flags & 16) != 0;
// Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths
const uint32_t HSK_pad = (HSK + 15) & ~15;
@@ -82,6 +84,8 @@ layout (binding = 5) writeonly buffer OV4 {D_TYPEV4 data_ov4[];};
layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
layout (binding = 7) readonly buffer SP {int32_t data_sparse[];};
#define MASK_OPT_ALL_NEG_INF 1
#define MASK_OPT_ALL_ZERO 2
@@ -144,7 +148,7 @@ ACC_TYPE perElemOpGetSink(const in uint32_t r, const in uint32_t c, const in ACC
uint32_t i, N, KV, split_k_index, Tr, start_j, end_j,
gqa_iq1, iq2, iq3, rk2, rk3, rv2, rv3, ik2, ik3, iv2, iv3,
q_stride, k_stride, v_stride, m_stride;
q_stride, k_stride, v_stride, m_stride, sparse_base;
void init_indices()
{
@@ -208,6 +212,33 @@ void init_indices()
// that prevents the compiler from folding the "&" through the select
// and breaking the alignment detection.
m_stride = (p.gqa_ratio > 1) ? (p.gqa_ratio >> 16) : KV;
// Sparse: the tile shares one mask row (gqa heads, or Br==1). split_k
// partitions the n_kv_max blocks.
if (USE_SPARSE) {
uint32_t qrow = (p.gqa_ratio > 1) ? gqa_iq1 : (i * Br);
sparse_base = (((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 + qrow) * p.split_kv;
uint32_t total_blocks = CEIL_DIV(p.split_kv, Bc);
uint32_t per_blocks = CEIL_DIV(total_blocks, p.k_num);
start_j = min(split_k_index * per_blocks, total_blocks);
end_j = min((split_k_index + 1) * per_blocks, total_blocks);
}
}
// Resolve a linear KV slot to a real column; false for inactive (sparse padding/-1, or dense OOB).
bool fa_kv_index(uint lin, out uint kv_col) {
if (USE_SPARSE) {
if (lin >= p.split_kv) {
kv_col = 0;
return false;
}
int idx = data_sparse[sparse_base + lin];
kv_col = idx >= 0 ? uint(idx) : 0;
return idx >= 0;
}
kv_col = lin;
return !KV_bounds_check || lin < KV;
}
// Bias applied to softmax to stay in fp16 range.
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp
+32 -16
View File
@@ -176,9 +176,16 @@ void main() {
uint32_t c = (idx + tid) / (Br / 4);
uint32_t r = (idx + tid) % (Br / 4);
if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) {
if ((!KV_bounds_check || j * Bc + c < KV)) {
uint32_t kcol;
bool kv_active = fa_kv_index(j * Bc + c, kcol);
if (kv_active) {
f16vec4 m;
if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) {
if (USE_SPARSE) {
// sparse is gqa-gated (m_stride == 0): all four rows share the value
FLOAT_TYPE mv = FLOAT_TYPE(data_m[m_offset + kcol]);
m = f16vec4(mv);
max_mask = max(max_mask, float(mv));
} else if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) {
m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)],
data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)],
data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)],
@@ -206,6 +213,8 @@ void main() {
m = f16vec4(0.0);
}
mask_cache[idx / WorkGroupSize] = m;
} else if (USE_SPARSE) {
mask_cache[idx / WorkGroupSize] = f16vec4(NEG_FLT_MAX_OVER_2);
}
}
}
@@ -231,17 +240,19 @@ void main() {
uint32_t c = (idx + tid) / (HSK_pad / 4);
if (idx + gl_WorkGroupSize.x <= Bc * HSK_pad / 4 || c < Bc) {
FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0);
if ((!KV_bounds_check || j * Bc + c < KV) && (HSK == HSK_pad || d < HSK / 4)) {
uint32_t kcol;
bool kv_active = fa_kv_index(j * Bc + c, kcol);
if (kv_active && (HSK == HSK_pad || d < HSK / 4)) {
#if !defined(BFLOAT16)
if (USE_DECODE_K) {
uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d;
uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * d;
uint ib = coord / BLOCK_SIZE_K;
uint iqs = (coord % BLOCK_SIZE_K);
K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K);
} else
#endif
{
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]);
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d]);
}
}
@@ -266,7 +277,7 @@ void main() {
if (SHMEM_STAGING == 0) {
// For quants we always need to dequant into kvsh; for f16/bf16 we can load
// directly from global memory when alignment / bounds allow it.
const bool stage_k = USE_DECODE_K || KV_bounds_check || d * 16 + 16 > HSK;
const bool stage_k = USE_DECODE_K || KV_bounds_check || USE_SPARSE || d * 16 + 16 > HSK;
if (stage_k) {
barrier();
[[unroll]] for (uint32_t idx = 0; idx < Bc * MatBr / 4; idx += gl_WorkGroupSize.x) {
@@ -274,17 +285,19 @@ void main() {
uint32_t row = (idx + tid) / (MatBr / 4);
if (idx + tid < Bc * MatBr / 4) {
FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0);
if ((!KV_bounds_check || j * Bc + row < KV) && (HSK == HSK_pad || d * 16 + col_vec * 4 < HSK)) {
uint32_t kcol;
bool kv_active = fa_kv_index(j * Bc + row, kcol);
if (kv_active && (HSK == HSK_pad || d * 16 + col_vec * 4 < HSK)) {
#if !defined(BFLOAT16)
if (USE_DECODE_K) {
uint coord = (j * Bc + row) * k_stride * BLOCK_SIZE_K + d * 16 + col_vec * 4;
uint coord = kcol * k_stride * BLOCK_SIZE_K + d * 16 + col_vec * 4;
uint ib = coord / BLOCK_SIZE_K;
uint iqs = (coord % BLOCK_SIZE_K);
K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K);
} else
#endif
{
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + row) * k_stride / 4 + d * 16 / 4 + col_vec]);
K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * 16 / 4 + col_vec]);
}
}
@@ -401,17 +414,19 @@ void main() {
uint32_t c = (idx + tid) / (HSV_pad / 4);
if (idx + gl_WorkGroupSize.x <= Bc * HSV_pad / 4 || c < Bc) {
FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0);
if ((!KV_bounds_check || j * Bc + c < KV) && (HSV == HSV_pad || d < HSV / 4)) {
uint32_t v_row;
bool kv_active = fa_kv_index(j * Bc + c, v_row);
if (kv_active && (HSV == HSV_pad || d < HSV / 4)) {
#if !defined(BFLOAT16)
if (USE_DECODE_V) {
uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d;
uint coord = v_row * v_stride * BLOCK_SIZE_V + 4 * d;
uint ib = coord / BLOCK_SIZE_V;
uint iqs = (coord % BLOCK_SIZE_V);
V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V);
} else
#endif
{
V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]);
V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + v_row * v_stride / 4 + d]);
}
}
@@ -441,21 +456,22 @@ void main() {
if (SHMEM_STAGING == 0) {
// For quants we always preload via kvsh. For f16/bf16 we only preload when
// alignment / bounds force it (otherwise we coopMatLoad direct from data_vv4).
const bool stage_v = USE_DECODE_V || KV_bounds_check;
const bool stage_v = USE_DECODE_V || KV_bounds_check || USE_SPARSE;
if (stage_v) {
[[unroll]] for (uint32_t i = 0; i < v_loads_per_thread; ++i) {
const uint idx = i * gl_WorkGroupSize.x + tid;
const uint row = idx / v_cols;
const uint col = idx % v_cols;
const uint v_row = j * Bc + row;
uint32_t v_row;
bool kv_active = fa_kv_index(j * Bc + row, v_row);
const uint v_col = hsv_tile * MatBc * row_split + col * 4;
const uint coord = v_row * v_stride * BLOCK_SIZE_V + v_col;
const uint ib = coord / BLOCK_SIZE_V;
const uint iqs = coord % BLOCK_SIZE_V;
if (!KV_bounds_check || (v_row < KV && v_col < HSV)) {
if (USE_SPARSE ? (kv_active && v_col < HSV) : (!KV_bounds_check || (v_row < KV && v_col < HSV))) {
#if !defined(BFLOAT16)
if (USE_DECODE_V) {
kvsh[row * vsh_stride + col] = dequantize4(ib, iqs, v_offset, BINDING_IDX_V);
@@ -479,7 +495,7 @@ void main() {
coopMatLoad(KMat, Psh, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor);
if (SHMEM_STAGING == 0) {
if (!USE_DECODE_V && !KV_bounds_check) {
if (!USE_DECODE_V && !KV_bounds_check && !USE_SPARSE) {
// F16/BF16 values can be loaded directly from global memory
const uint v_tile_row = j * Bc + bc_chunk * MatBc;
const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4;
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
+87 -8
View File
@@ -29,6 +29,12 @@
#include "dequant_funcs_cm2.glsl"
#endif
#ifdef GL_NV_cooperative_matrix_decode_vector
#define FA_GATHER_BS 4u
#else
#define FA_GATHER_BS 1u
#endif
// buffer_reference stride = sizeof(struct) = FaBlockBytesK/V.
layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_K {
uint8_t raw[FaBlockBytesK];
@@ -107,6 +113,67 @@ layout (binding = 1) readonly buffer K {uint8_t data_k[];};
layout (binding = 2) readonly buffer V {uint8_t data_v[];};
layout (binding = 3) readonly buffer M {uint8_t data_m[];};
// f16 aliases for the sparse gather callbacks.
layout (binding = 1) readonly buffer KF16 {float16_t data_kf16[];};
layout (binding = 2) readonly buffer VF16 {float16_t data_vf16[];};
layout (binding = 3) readonly buffer MF16 {float16_t data_mf16[];};
#ifdef GL_NV_cooperative_matrix_decode_vector
layout (binding = 1) readonly buffer KF16V4 {f16vec4 data_kf16v4[];};
layout (binding = 2) readonly buffer VF16V4 {f16vec4 data_vf16v4[];};
#endif
// K/V/mask f16-element offsets for the current head/batch, set in main().
uint32_t g_k_off_elem, g_v_off_elem, g_m_off_elem;
#if !defined(BFLOAT16)
// blockCoords are in block units: KV slot = blockCoords[0],
// head dim = blockCoords[1]*FA_GATHER_BS + coordInBlock[1].
float16_t faGatherK(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return float16_t(0); }
const int r = data_sparse[sparse_base + blockCoords[0]];
return r < 0 ? float16_t(0) : data_kf16[g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
}
float16_t faGatherV(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return float16_t(0); }
const int r = data_sparse[sparse_base + blockCoords[0]];
return r < 0 ? float16_t(0) : data_vf16[g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
}
#ifdef GL_NV_cooperative_matrix_decode_vector
f16vec4 faGatherKVector(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return f16vec4(0); }
const int r = data_sparse[sparse_base + blockCoords[0]];
if (r < 0) { return f16vec4(0); }
const uint32_t o = g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
return data_kf16v4[o / 4];
}
f16vec4 faGatherVVector(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return f16vec4(0); }
const int r = data_sparse[sparse_base + blockCoords[0]];
if (r < 0) { return f16vec4(0); }
const uint32_t o = g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
return data_vf16v4[o / 4];
}
#define FAGATHERK , faGatherK, faGatherKVector
#define FAGATHERV , faGatherV, faGatherVVector
#else
#define FAGATHERK , faGatherK
#define FAGATHERV , faGatherV
#endif
#endif
// Add gathered mask to S (slope==1 since sparse requires max_bias==0). col = slot in block jblk.
ACC_TYPE faAddSparseMask(const uint32_t row, const uint32_t col, const ACC_TYPE elem, const uint32_t jblk) {
const float NEG = uintBitsToFloat(0xFEFFFFFF);
const uint32_t kvslot = jblk * Bc + col;
if (kvslot >= p.split_kv) { return ACC_TYPE(NEG); }
const int r = data_sparse[sparse_base + kvslot];
return r < 0 ? ACC_TYPE(NEG) : elem + ACC_TYPE(data_mf16[g_m_off_elem + row * m_stride + uint(r)]);
}
ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) {
return max(x, y);
}
@@ -185,14 +252,16 @@ void main() {
tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0);
const uint bs_k = fa_block_elems(FaTypeK);
const uint bs_v = fa_block_elems(FaTypeV);
const uint bs_k = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeK);
const uint bs_v = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeV);
tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, bs_k);
tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, bs_v);
// Sparse iterates n_kv_max (in split_kv); the decode callbacks remap each slot.
const uint32_t KV_iter = USE_SPARSE ? p.split_kv : KV;
tensorLayoutQ = setTensorLayoutDimensionNV(tensorLayoutQ, N, HSK);
tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, HSK);
tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, HSV);
tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV_iter, HSK);
tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV_iter, HSV);
// hint to the compiler that strides are aligned for the aligned variant of the shader
if (Clamp != gl_CooperativeMatrixClampModeConstantNV)
@@ -250,6 +319,10 @@ void main() {
mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride;
}
g_k_off_elem = (ik2*p.nb12 + ik3*p.nb13) / 2;
g_v_off_elem = (iv2*p.nb22 + iv3*p.nb23) / 2;
g_m_off_elem = m_offset / 2;
uint32_t mask_opt = 0;
uint32_t mask_opt_idx = ~0;
@@ -257,7 +330,7 @@ void main() {
for (uint32_t j = start_j; j < end_j; ++j) {
coopmat<float16_t, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> mv = coopmat<float16_t, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(0);
if (MASK_ENABLE) {
if (MASK_ENABLE && !USE_SPARSE) {
if (USE_MASK_OPT && mask_opt_idx != j / 16) {
mask_opt_idx = j / 16;
@@ -315,7 +388,9 @@ void main() {
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
#else
const bool k_use_decode = (bs_k > 1u);
if (k_use_decode) {
if (USE_SPARSE) {
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FAGATHERK);
} else if (k_use_decode) {
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FADECODEK);
} else {
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
@@ -330,7 +405,9 @@ void main() {
}
}
if (MASK_ENABLE) {
if (MASK_ENABLE && USE_SPARSE) {
coopMatPerElementNV(S, S, faAddSparseMask, j);
} else if (MASK_ENABLE) {
S += slopeMat*coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(mv);
}
@@ -385,7 +462,9 @@ void main() {
coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad));
#else
const bool v_use_decode = (bs_v > 1u);
if (v_use_decode) {
if (USE_SPARSE) {
coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FAGATHERV);
} else if (v_use_decode) {
coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FADECODEV);
} else {
coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad));
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp
+102
View File
@@ -0,0 +1,102 @@
#version 450
#extension GL_EXT_control_flow_attributes : enable
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
#ifdef USE_SUBGROUPS
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_ballot : require
#endif
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
layout(constant_id = 0) const uint BLOCK_SIZE = 128;
layout(constant_id = 1) const uint NUM_SUBGROUPS = 1;
layout (binding = 0) readonly buffer M {float16_t data_m[];};
layout (binding = 1) writeonly buffer I {int32_t data_i[];};
layout (push_constant) uniform parameter {
uint KV;
uint nem1;
uint nem2;
uint nbm1;
uint nbm2;
uint nbm3;
uint n_kv_max;
} p;
#ifdef USE_SUBGROUPS
shared uvec4 ballots_sh[NUM_SUBGROUPS];
#else
shared uint scan[BLOCK_SIZE];
#endif
// One workgroup per mask row: compact the finite-mask KV positions into a
// per-row index list of length n_kv_max, -1 padded. Emitted in ascending KV
// order so the downstream attention accumulation is deterministic.
void main() {
const uint i1 = gl_WorkGroupID.x;
const uint i2 = gl_WorkGroupID.y;
const uint i3 = gl_WorkGroupID.z;
const uint tid = gl_LocalInvocationIndex;
const uint m_base = i3 * p.nbm3 + i2 * p.nbm2 + i1 * p.nbm1;
const uint out_base = ((i3 * p.nem2 + i2) * p.nem1 + i1) * p.n_kv_max;
uint base = 0;
for (uint chunk = 0; chunk < p.KV; chunk += BLOCK_SIZE) {
const uint k = chunk + tid;
bool selected = false;
if (k < p.KV) {
const float v = float(data_m[m_base + k]);
selected = !isinf(v) && !isnan(v);
}
#ifdef USE_SUBGROUPS
const uvec4 ballot = subgroupBallot(selected);
if (subgroupElect()) {
ballots_sh[gl_SubgroupID] = ballot;
}
barrier();
uint subgroup_base = 0;
uint total = 0;
[[unroll]] for (uint s = 0; s < gl_NumSubgroups; ++s) {
if (s == gl_SubgroupID) {
subgroup_base = total;
}
total += subgroupBallotBitCount(ballots_sh[s]);
}
barrier();
const uint slot = base + subgroup_base + subgroupBallotExclusiveBitCount(ballot);
#else
// Hillis-Steele inclusive prefix sum over the workgroup.
scan[tid] = selected ? 1u : 0u;
barrier();
for (uint off = 1; off < BLOCK_SIZE; off <<= 1) {
uint add = 0;
if (tid >= off) {
add = scan[tid - off];
}
barrier();
scan[tid] += add;
barrier();
}
const uint inclusive = scan[tid];
const uint total = scan[BLOCK_SIZE - 1];
const uint slot = base + inclusive - 1u;
#endif
if (selected && slot < p.n_kv_max) {
data_i[out_base + slot] = int32_t(k);
}
base += total;
barrier();
}
for (uint s = min(base, p.n_kv_max) + tid; s < p.n_kv_max; s += BLOCK_SIZE) {
data_i[out_base + s] = int32_t(-1);
}
}
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+2
View File
@@ -948,6 +948,8 @@ void process_shaders() {
string_to_spv("fa_split_k_reduce", "flash_attn_split_k_reduce.comp", {});
string_to_spv("fa_mask_opt", "flash_attn_mask_opt.comp", {});
string_to_spv("fa_sparse_compact", "flash_attn_sparse_compact.comp", {});
string_to_spv("fa_sparse_compact_subgroup", "flash_attn_sparse_compact.comp", {{"USE_SUBGROUPS", "1"}});
string_to_spv("quantize_q8_1", "quantize_q8_1.comp", {});
string_to_spv("quantize_q8_1_subgroup", "quantize_q8_1.comp", {{"USE_SUBGROUPS", "1"}});
src/models/dflash.cpp
+15 -15
View File
@@ -240,21 +240,6 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
}
}
std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const llm_graph_params & params) const {
switch (params.gtype) {
case LLM_GRAPH_TYPE_ENCODER:
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
if (hparams.dsv4_hc_mult > 0) {
return std::make_unique<graph_dsv4>(*this, params);
}
return std::make_unique<graph<false>>(*this, params);
default:
GGML_ABORT("invalid graph type");
};
}
template <>
ggml_tensor * llama_model_dflash::graph<true>::build_inp_embd_enc() const {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
@@ -999,3 +984,18 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
build_dspark_markov_head(*this, model, inp_tokens);
}
}
std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const llm_graph_params & params) const {
switch (params.gtype) {
case LLM_GRAPH_TYPE_ENCODER:
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
if (hparams.dsv4_hc_mult > 0) {
return std::make_unique<graph_dsv4>(*this, params);
}
return std::make_unique<graph<false>>(*this, params);
default:
GGML_ABORT("invalid graph type");
};
}
src/models/eagle3.cpp
+12 -12
View File
@@ -100,18 +100,6 @@ void llama_model_eagle3::load_arch_tensors(llama_model_loader &) {
}
}
std::unique_ptr<llm_graph_context> llama_model_eagle3::build_arch_graph(const llm_graph_params & params) const {
switch (params.gtype) {
case LLM_GRAPH_TYPE_ENCODER:
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
return std::make_unique<graph<false>>(*this, params);
default:
GGML_ABORT("invalid graph type");
};
}
template <>
ggml_tensor * llama_model_eagle3::graph<true>::build_inp_embd_enc() const {
ggml_tensor * cur = nullptr;
@@ -336,3 +324,15 @@ llama_model_eagle3::graph<false>::graph(const llama_model & model, const llm_gra
ggml_build_forward_expand(gf, cur);
}
std::unique_ptr<llm_graph_context> llama_model_eagle3::build_arch_graph(const llm_graph_params & params) const {
switch (params.gtype) {
case LLM_GRAPH_TYPE_ENCODER:
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
return std::make_unique<graph<false>>(*this, params);
default:
GGML_ABORT("invalid graph type");
};
}
src/models/t5.cpp
+12 -12
View File
@@ -106,18 +106,6 @@ void llama_model_t5::load_arch_tensors(llama_model_loader &) {
}
}
std::unique_ptr<llm_graph_context> llama_model_t5::build_arch_graph(const llm_graph_params & params) const {
switch (params.gtype) {
case LLM_GRAPH_TYPE_ENCODER:
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
return std::make_unique<graph<false>>(*this, params);
default:
GGML_ABORT("invalid graph type");
};
}
template <>
llama_model_t5::graph<false>::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();
@@ -368,3 +356,15 @@ llama_model_t5::graph<true>::graph(const llama_model & model, const llm_graph_pa
ggml_build_forward_expand(gf, cur);
}
std::unique_ptr<llm_graph_context> llama_model_t5::build_arch_graph(const llm_graph_params & params) const {
switch (params.gtype) {
case LLM_GRAPH_TYPE_ENCODER:
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
return std::make_unique<graph<false>>(*this, params);
default:
GGML_ABORT("invalid graph type");
};
}
tools/ui/src/lib/hooks/use-tools-panel.svelte.ts
+1 -1
View File
@@ -47,7 +47,7 @@ export function useToolsPanel(): UseToolsPanelReturn {
if (toolsStore.toolGroups.length > 0) return null;
// Tools endpoint is unreachable (404) — server started without --tools
// Tools endpoint unreachable (403) — server started without tools
if (toolsStore.isToolsEndpointUnreachable) {
return `The KoboldCpp MCP bridge may accessible on the same URL with /mcp at the end but must be added manually.`;
}
tools/ui/src/lib/stores/agentic/index.svelte.ts
+8 -2
View File
@@ -315,8 +315,14 @@ class AgenticStore {
// Clear any pending permissions/continue requests for this conversation when starting a new flow
this.gates.clear(conversationId);
// Ensure server tools are fetched before checking if agentic is enabled
if (toolsStore.serverTools.length === 0 && !toolsStore.loading) {
// Ensure server tools are fetched before checking if agentic is enabled.
// A disabled /tools endpoint stays disabled for the life of the server,
// so the tools panel is the only place that probes it again.
if (
toolsStore.serverTools.length === 0 &&
!toolsStore.loading &&
!toolsStore.isToolsEndpointUnreachable
) {
await toolsStore.fetchServerTools();
}
tools/ui/src/lib/stores/tools.svelte.ts
+3 -6
View File
@@ -32,7 +32,7 @@ import { mcpStore } from '$lib/stores/mcp/index.svelte';
import { modelsStore } from '$lib/stores/models/index.svelte';
import { settingsStore } from '$lib/stores/settings/index.svelte';
import type { OpenAIToolDefinition, ToolEntry, ToolGroup } from '$lib/types';
import { buildSandboxToolDefinition } from '$lib/utils';
import { ApiError, buildSandboxToolDefinition } from '$lib/utils';
import { SvelteMap, SvelteSet } from 'svelte/reactivity';
/** Stable selection identity for a tool, shared by the disabled set and the permission store */
@@ -246,13 +246,10 @@ class ToolsStore {
toolInfos.filter((info) => info.uses_cwd).map((info) => info.tool)
);
} catch (err) {
const errorMessage = err instanceof Error ? err.message : String(err);
this._error = errorMessage;
this._error = err instanceof Error ? err.message : String(err);
// 403 from /tools means the server was started without --tools
// TODO: check status code instead of relying on message
if (errorMessage.includes('this feature is disabled')) {
if (err instanceof ApiError && err.status === 403) {
this._toolsEndpointUnreachable = true;
console.info('[ToolsStore] Server tools are disabled on the server');
} else {