Commit 412ef74e for whisper.cpp
commit 412ef74ec92575a564b0fb09201631452fda46c5
Author: Aman Gupta <amangupta052@gmail.com>
Date: Tue Nov 3 00:13:33 2026 +0800
llama : flash_attn_ext_rows - more efficient unified kv for multi-seq (llama/29510)
* llama : use llm_attn_mask and kv_rows for unified attention
Assisted-by: dsh:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
* cont : optimize `get_seq_cells`
* llama : build the kv_rows lists from per-sequence cell bitmaps
Keep a bitmap over the cells for each sequence (the transpose of `seq`), updated in seq_pos_inc/dec, and build the cell lists by walking its set bits. The lists stay in cell order, cells shared by several sequences go to every list, and nothing has to be rebuilt when cells are removed.
Also leave out of the lists the cells that the SWA hides from every token of the slice (standard and chunked SWA).
* cont : make get_kv_rows_p_min private
* cont : minor metal
* llama : use kv_rows for every ubatch of a unified cache
A ubatch made of equal-length runs with a single sequence each (f.ex. the drafts of several slots) now splits into per-sequence slices.
Any other ubatch (interleaved, uneven or coupled sequences) becomes a single slice over the cells of all its sequences instead of falling back to the dense KQ mask, so the graph keeps the topology of the reserved graph.
* llama : merge get_kv_rows_seq_ids and get_kv_rows_p_min into get_kv_rows_slices
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
index b188c0bf..3ed41814 100644
--- a/ggml/include/ggml.h
+++ b/ggml/include/ggml.h
@@ -2516,6 +2516,24 @@ extern "C" {
float max_bias,
float logit_softcap);
+ // same as ggml_flash_attn_ext, but each q slice attends to its own list of K/V rows
+ // q: [n_embd_k, n_batch, n_head, ne3]
+ // k: [n_embd_k, n_rows, n_head_kv, 1 ] shared by all q slices
+ // v: [n_embd_v, n_rows, n_head_kv, 1 ]
+ // mask: [n_kv, n_batch, ne32, ne3]
+ // kv_rows: [n_kv, ne3] I32, mask column i of slice i3 uses row kv_rows[i, i3] of k and v
+ // negative rows are padding, their mask entries must be -inf
+ GGML_API struct ggml_tensor * ggml_flash_attn_ext_rows(
+ struct ggml_context * ctx,
+ struct ggml_tensor * q,
+ struct ggml_tensor * k,
+ struct ggml_tensor * v,
+ struct ggml_tensor * mask,
+ struct ggml_tensor * kv_rows,
+ float scale,
+ float max_bias,
+ float logit_softcap);
+
GGML_DEPRECATED(GGML_API void ggml_flash_attn_ext_set_prec(
struct ggml_tensor * a,
enum ggml_prec prec),
diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp
index 016d2320..d3a47f83 100644
--- a/ggml/src/ggml-cann/ggml-cann.cpp
+++ b/ggml/src/ggml-cann/ggml-cann.cpp
@@ -2670,6 +2670,10 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
return true;
case GGML_OP_FLASH_ATTN_EXT:
{
+ // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+ if (op->src[5] != nullptr) {
+ return false;
+ }
#ifdef ASCEND_310P
// FA not support on 310p device
return false;
diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp
index d3a156b1..dbdde699 100644
--- a/ggml/src/ggml-cpu/ops.cpp
+++ b/ggml/src/ggml-cpu/ops.cpp
@@ -8844,6 +8844,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
const ggml_tensor * v = dst->src[2];
const ggml_tensor * mask = dst->src[3];
const ggml_tensor * sinks = dst->src[4];
+ const ggml_tensor * kv_rows = dst->src[5];
GGML_TENSOR_LOCALS(int64_t, neq, q, ne)
GGML_TENSOR_LOCALS(size_t, nbq, q, nb)
@@ -8951,6 +8952,8 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
const float * pq = (const float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3));
q_to_vec_dot(pq, Q_q, DK);
+ const int32_t * rows = kv_rows ? (const int32_t *) ((const char *) kv_rows->data + iq3*kv_rows->nb[1]) : nullptr;
+
// online softmax / attention
// loop over n_kv and n_head_kv
// ref: https://arxiv.org/pdf/2112.05682.pdf
@@ -8961,9 +8964,14 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
continue;
}
+ const int64_t ik1 = rows ? rows[ic] : ic;
+ if (ik1 < 0) {
+ continue;
+ }
+
float s; // KQ value
- const char * k_data = (const char *) k->data + ( ic*nbk1 + ik2*nbk2 + ik3*nbk3);
+ const char * k_data = (const char *) k->data + (ik1*nbk1 + ik2*nbk2 + ik3*nbk3);
kq_vec_dot(DK, &s, 0, k_data, 0, Q_q, 0, 1);
s = s*scale; // scale KQ value
@@ -8979,7 +8987,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
float ms = 1.0f; // upon new higher max val, scale VKQ and KQ sum with this value
float vs = 1.0f; // post-softmax KQ value, expf(s - M)
- const char * v_data = ((const char *) v->data + (ic*nbv1 + iv2*nbv2 + iv3*nbv3));
+ const char * v_data = ((const char *) v->data + (ik1*nbv1 + iv2*nbv2 + iv3*nbv3));
if (v->type == GGML_TYPE_F16) {
if (s > M) {
@@ -9146,6 +9154,10 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
static constexpr int Q_TILE_SZ = ggml_fa_tile_config::Q;
static constexpr int KV_TILE_SZ = ggml_fa_tile_config::KV;
+ // with kv_rows the loop runs over the mask columns and reads K/V rows through kv_rows
+ const ggml_tensor * kv_rows = dst->src[5];
+ const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : nek1;
+
int ir = ir0;
while (ir < ir1) {
// q indices for the start of this tile
@@ -9153,6 +9165,10 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
const int iq2 = (ir - iq3*neq2*neq1)/neq1;
const int iq1 = (ir - iq3*neq2*neq1 - iq2*neq1);
+ // padding rows are masked, read row 0 for them
+ const int32_t * rows = kv_rows ? (const int32_t *) ((const char *) kv_rows->data + iq3*kv_rows->nb[1]) : nullptr;
+ auto kv_row = [rows](int64_t i) -> int64_t { return rows ? std::max<int32_t>(rows[i], 0) : i; };
+
// Number of valid rows in this tile:
// - limited by tile size (Q_TILE_SZ)
// - limited by chunk boundary (ir1 - ir)
@@ -9212,8 +9228,8 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
memset(K_f32, 0, DK * KV_TILE_SZ * sizeof(float));
memset(V32, 0, KV_TILE_SZ * DV * sizeof(float));
- for (int64_t ic = 0; ic < nek1; ic += KV_TILE_SZ) {
- const int kv_tile = (int)std::min((int64_t)KV_TILE_SZ, nek1 - ic);
+ for (int64_t ic = 0; ic < n_kv_cols; ic += KV_TILE_SZ) {
+ const int kv_tile = (int)std::min((int64_t)KV_TILE_SZ, n_kv_cols - ic);
// skip the tile entirely if all the masks are -inf
if (mask) {
@@ -9240,7 +9256,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
// Pack K tile transposed: K_f32[dk][kv] so KV_TILE is contiguous (SIMD dim)
// Zero-pad the last tile so the GEMM always operates on KV_TILE_SZ columns
for (int tk = 0; tk < kv_tile; tk++) {
- const char * k_data = (const char *)k->data + (ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3;
+ const char * k_data = (const char *)k->data + kv_row(ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3;
if (kv_type == GGML_TYPE_F16) {
const ggml_fp16_t * k_f16 = (const ggml_fp16_t *)k_data;
for (int64_t dk = 0; dk < DK; dk++) {
@@ -9305,7 +9321,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
// V accumulation: VKQ32 += softmax(KQ) * V
// Pack V tile to contiguous F32, zero-padded
for (int tk = 0; tk < kv_tile; tk++) {
- const char * v_data = (const char *)v->data + (ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3;
+ const char * v_data = (const char *)v->data + kv_row(ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3;
if (kv_type == GGML_TYPE_F16) {
ggml_cpu_fp16_to_fp32((const ggml_fp16_t *)v_data, V32 + tk * DV, DV);
} else {
@@ -9477,8 +9493,12 @@ static void ggml_compute_forward_flash_attn_ext_f16(
// When use_ref is set, force the vec-only reference implementation (no tiling, no KV-chunking)
const bool use_ref = params->use_ref;
+ // with kv_rows the loops run over the mask columns, the split KV path does not support it
+ const ggml_tensor * kv_rows = dst->src[5];
+ const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : nek1;
+
const bool kv_is_f32_or_f16 = (k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16);
- const bool use_split_kv_path = !use_ref && (neq1 == 1 && neq3 == 1) && kv_is_f32_or_f16 && (k->type == v->type) && q->type == GGML_TYPE_F32 && nek1 >= 512;
+ const bool use_split_kv_path = !use_ref && !kv_rows && (neq1 == 1 && neq3 == 1) && kv_is_f32_or_f16 && (k->type == v->type) && q->type == GGML_TYPE_F32 && nek1 >= 512;
if (use_split_kv_path) {
const int64_t chunk_size = (nek1 + nth - 1) / nth;
@@ -9557,7 +9577,7 @@ static void ggml_compute_forward_flash_attn_ext_f16(
if (use_tiled) {
ggml_compute_forward_flash_attn_ext_tiled(params, dst, ir0, ir1);
} else {
- ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, nek1, nullptr, 0);
+ ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, n_kv_cols, nullptr, 0);
}
current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1);
diff --git a/ggml/src/ggml-cpu/spacemit/ime.cpp b/ggml/src/ggml-cpu/spacemit/ime.cpp
index 0683ff88..5f114f58 100644
--- a/ggml/src/ggml-cpu/spacemit/ime.cpp
+++ b/ggml/src/ggml-cpu/spacemit/ime.cpp
@@ -1046,6 +1046,10 @@ class tensor_traits_common : public tensor_traits_base {
return true;
}
case GGML_OP_FLASH_ATTN_EXT:
+ // K/V row lists are handled by the generic CPU path
+ if (op->src[5] != nullptr) {
+ return false;
+ }
forward_flash_attn_ext_f16(params, op);
return true;
case GGML_OP_CONT:
diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh
index 41d94833..beb87794 100644
--- a/ggml/src/ggml-cuda/fattn-common.cuh
+++ b/ggml/src/ggml-cuda/fattn-common.cuh
@@ -25,6 +25,7 @@ typedef void (* fattn_kernel_t)(
const char * __restrict__ mask,
const char * __restrict__ sinks,
const int * __restrict__ KV_max,
+ const int * __restrict__ KV_rows,
float * __restrict__ dst,
float2 * __restrict__ dst_meta,
const float scale,
@@ -991,8 +992,13 @@ void launch_fattn(
const bool V_is_K_view = V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs));
- const ggml_tensor * mask = dst->src[3];
- const ggml_tensor * sinks = dst->src[4];
+ const ggml_tensor * mask = dst->src[3];
+ const ggml_tensor * sinks = dst->src[4];
+ const ggml_tensor * kv_rows = dst->src[5];
+
+ // with kv_rows the kernel iterates over the mask columns and reads K/V rows through kv_rows
+ const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : K->ne[1];
+ GGML_ASSERT(!kv_rows || !use_sparse);
ggml_tensor * KQV = dst;
@@ -1113,8 +1119,8 @@ void launch_fattn(
// Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
// Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
- // multiple sequences of possibly different lengths.
- const bool scan_mask = !use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1);
+ // multiple sequences of possibly different lengths, or kv_rows lists that end with padding.
+ const bool scan_mask = !use_sparse && mask && n_kv_cols % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || kv_rows);
if (scan_mask) {
const int64_t s31 = mask->nb[1] / sizeof(half2);
const int64_t s33 = mask->nb[3] / sizeof(half2);
@@ -1123,7 +1129,7 @@ void launch_fattn(
const dim3 block_dim_KV_max(FATTN_KQ_STRIDE/2, 1, 1);
const int ne_KV_max = blocks_num_KV_max.x*blocks_num_KV_max.y;
- const int iter_k = K->ne[1] / FATTN_KQ_STRIDE;
+ const int iter_k = n_kv_cols / FATTN_KQ_STRIDE;
KV_max.alloc(ne_KV_max);
ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num_KV_max, block_dim_KV_max, 0, main_stream);
@@ -1132,13 +1138,19 @@ void launch_fattn(
CUDA_CHECK(cudaGetLastError());
}
+ // K/V are shared by all slices when read through kv_rows
+ if (kv_rows) {
+ nb13 = 0;
+ nb23 = 0;
+ }
+
const dim3 block_dim(warp_size, nwarps, 1);
int max_blocks_per_sm = 1; // Max. number of active blocks limited by occupancy.
CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, fattn_kernel, block_dim.x * block_dim.y * block_dim.z, nbytes_shared));
GGML_ASSERT(max_blocks_per_sm > 0);
int parallel_blocks = max_blocks_per_sm;
- const int64_t n_kv = use_sparse ? n_kv_max : K->ne[1];
+ const int64_t n_kv = use_sparse ? n_kv_max : n_kv_cols;
const int ntiles_KV = (n_kv + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length.
dim3 blocks_num;
@@ -1255,6 +1267,7 @@ void launch_fattn(
mask ? ((const char *) mask->data) : nullptr,
sinks ? ((const char *) sinks->data) : nullptr,
KV_max.ptr,
+ kv_rows ? (const int *) kv_rows->data : nullptr,
!stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr,
scale, max_bias, m0, m1, n_head_log2, logit_softcap,
Q->ne[0], ne01, Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3],
diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
index 56decfa2..96b596ea 100644
--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh
+++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
@@ -420,6 +420,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
// padded slots gather row 0, the -inf mask removes their contribution
const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : 0;
i_KV = index >= 0 ? index : 0;
+ } else if (indices) {
+ // K/V rows given per mask column, negative rows are padding
+ i_KV = max(indices[k_VKQ_0 + i], 0);
} else {
i_KV = k_VKQ_0 + i;
}
@@ -467,8 +470,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
if constexpr (use_sparse) {
const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1;
src = index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero;
+ } else if (!oob_check || i < i_sup) {
+ const int64_t i_KV = indices ? max(indices[k_VKQ_0 + i], 0) : k_VKQ_0 + i;
+ src = KV + i_KV*stride_KV + k*h2_per_chunk;
} else {
- src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero;
+ src = zero;
}
ggml_cuda_memcpy_1<16>(swizzle<stride_tile>(tile_KV, i*stride_tile + k*h2_per_chunk, i), src);
}
@@ -634,7 +640,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
cp_async_wait_all();
__syncthreads();
flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
- (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr);
+ (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, indices);
} else {
// the sparse mask values are gathered per element, always load them synchronously
constexpr bool use_cp_async = nstages == 1 && !use_sparse;
@@ -991,7 +997,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
(mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
}
flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
- (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, nullptr);
+ (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, indices);
}
}
@@ -1343,7 +1349,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
(mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
}
flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
- (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr);
+ (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, indices);
}
// kb0_start is always < kb0_stop so the last iter can be executed unconditionally.
@@ -1802,6 +1808,7 @@ static __global__ void flash_attn_ext_f16(
const char * mask_ptr,
const char * sinks_ptr,
const int * KV_max_ptr,
+ const int * KV_rows_ptr,
float * dst_ptr,
float2 * dst_meta_ptr,
const float scale,
@@ -1924,7 +1931,8 @@ static __global__ void flash_attn_ext_f16(
const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
- const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr;
+ const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 :
+ KV_rows_ptr ? KV_rows_ptr + int64_t(sequence)*ne11 : nullptr;
const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
@@ -1973,7 +1981,8 @@ static __global__ void flash_attn_ext_f16(
const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
- const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr;
+ const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 :
+ KV_rows_ptr ? KV_rows_ptr + int64_t(sequence)*ne11 : nullptr;
const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
@@ -1989,7 +1998,7 @@ static __global__ void flash_attn_ext_f16(
(Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
#else
- GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+ GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
max_bias, m0, m1, n_head_log2, logit_softcap,
ne00, ne01, ne02, ne03,
nb01, nb02, nb03,
diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh
index bd4a7710..ba4628fe 100644
--- a/ggml/src/ggml-cuda/fattn-tile.cuh
+++ b/ggml/src/ggml-cuda/fattn-tile.cuh
@@ -797,6 +797,7 @@ static __global__ void flash_attn_tile(
const char * mask_ptr,
const char * sinks_ptr,
const int * KV_max_ptr,
+ const int * KV_rows_ptr,
float * dst_ptr,
float2 * dst_meta_ptr,
const float scale,
@@ -821,6 +822,7 @@ static __global__ void flash_attn_tile(
const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr;
float * GGML_CUDA_RESTRICT dst = dst_ptr;
float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
+ GGML_UNUSED(KV_rows_ptr);
// Skip unused kernel variants for faster compilation:
@@ -1132,7 +1134,7 @@ static __global__ void flash_attn_tile(
}
}
#else
- GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+ GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
max_bias, m0, m1, n_head_log2, logit_softcap,
ne00, ne01, ne02, ne03,
nb01, nb02, nb03,
diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh
index 74f28204..9df15128 100644
--- a/ggml/src/ggml-cuda/fattn-vec.cuh
+++ b/ggml/src/ggml-cuda/fattn-vec.cuh
@@ -25,6 +25,7 @@ static __global__ void flash_attn_ext_vec(
const char * mask_ptr,
const char * sinks_ptr,
const int * KV_max_ptr,
+ const int * KV_rows_ptr,
float * dst_ptr,
float2 * dst_meta_ptr,
const float scale,
@@ -50,6 +51,7 @@ static __global__ void flash_attn_ext_vec(
const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr;
float * GGML_CUDA_RESTRICT dst = dst_ptr;
float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
+ GGML_UNUSED(KV_rows_ptr);
// Skip unused kernel variants for faster compilation:
if (use_logit_softcap && !(D == 128 || D == 256)) {
@@ -514,7 +516,7 @@ static __global__ void flash_attn_ext_vec(
dst_meta[((sequence*int(ne01.z) + ic0 + tid)*ne02 + head)*gridDim.y + blockIdx.y] = make_float2(KQ_max[tid], KQ_sum[tid]);
}
#else
- GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+ GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
max_bias, m0, m1, n_head_log2, logit_softcap,
ne00, ne01, ne02, ne03,
nb01, nb02, nb03,
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index f1ebd4e8..5d7ebf96 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -161,7 +161,7 @@ bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_
// the dense kernel handles up to 64/ncols2 queries per K/V pass, the single-query gather has to beat that
const int64_t n_gather = (ncols1 == 1 ? std::min<int64_t>(Q->ne[1], 64/ncols2) : ncols1) * (int64_t) n_kv_max;
- return ((GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc)) || amd_mfma_available(cc)) &&
+ return ((GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc)) || amd_mfma_available(cc)) && dst->src[5] == nullptr &&
mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f &&
mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 &&
K->ne[1] >= std::max<int64_t>(4096, GGML_CUDA_CC_IS_NVIDIA(cc) ? 2*n_gather : 2LL*n_kv_max);
@@ -216,12 +216,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con
const ggml_tensor * V = dst->src[2];
const ggml_tensor * mask = dst->src[3];
+ // with kv_rows the kernel iterates over the mask columns
+ const int64_t n_kv = dst->src[5] ? dst->src[5]->ne[0] : K->ne[1];
+
float max_bias = 0.0f;
memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float));
// Edge cases like no mask, ALiBi, unpadded K/V, or misaligned addresses for large data transfers
// are put into the template specialization without GQA optimizations.
- bool use_gqa_opt = mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;
+ bool use_gqa_opt = mask && max_bias == 0.0f && n_kv % FATTN_KQ_STRIDE == 0;
for (const ggml_tensor * t : {Q, K, V, mask}) {
if (t == nullptr || ggml_is_quantized(t->type)) {
continue;
@@ -567,6 +570,10 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
const ggml_tensor * K = dst->src[1];
const ggml_tensor * V = dst->src[2];
const ggml_tensor * mask = dst->src[3];
+ const ggml_tensor * kv_rows = dst->src[5];
+
+ // with kv_rows the kernel iterates over the mask columns
+ const int64_t n_kv = kv_rows ? kv_rows->ne[0] : K->ne[1];
const int gqa_ratio = Q->ne[2] / K->ne[2];
GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
@@ -576,7 +583,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
// The effective batch size for the kernel can be increased by gqa_ratio.
// The kernel versions without this optimization are also used for ALiBi, if there is no mask, or if the KV cache is not padded,
- bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;
+ bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && n_kv % FATTN_KQ_STRIDE == 0;
for (const ggml_tensor * t : {Q, K, V, mask}) {
if (t == nullptr || ggml_is_quantized(t->type)) {
continue;
@@ -648,6 +655,12 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
return BEST_FATTN_KERNEL_NONE;
}
+ // only the MMA kernel reads K/V through kv_rows
+ if (kv_rows) {
+ return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && mask && Q->ne[0] != 40 && Q->ne[0] != 72 ?
+ BEST_FATTN_KERNEL_MMA_F16 : BEST_FATTN_KERNEL_NONE;
+ }
+
// For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes:
// 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path.
const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0;
diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp
index 35a9afab..cdd947c5 100644
--- a/ggml/src/ggml-et/ggml-et.cpp
+++ b/ggml/src/ggml-et/ggml-et.cpp
@@ -1273,7 +1273,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
case GGML_OP_FLASH_ATTN_EXT:
if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
(op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && op->src[2] &&
- (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr &&
+ (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr && op->src[5] == nullptr &&
ggml_is_contiguous_rows(op) && ggml_is_contiguous_rows(op->src[0])) {
float max_bias = 0.0f;
float logit_softcap = 0.0f;
diff --git a/ggml/src/ggml-hexagon/htp-kparams.cpp b/ggml/src/ggml-hexagon/htp-kparams.cpp
index c8c1947e..03b5a3b5 100644
--- a/ggml/src/ggml-hexagon/htp-kparams.cpp
+++ b/ggml/src/ggml-hexagon/htp-kparams.cpp
@@ -348,6 +348,11 @@ static bool htp_kparams_solve_flash_attn(
return false;
}
+ // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+ if (op->src[5] != nullptr) {
+ return false;
+ }
+
memset(kparams, 0, sizeof(*kparams));
const struct ggml_tensor * q = op->src[0];
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index 548b9896..463d57c7 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -1987,6 +1987,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(
return res;
}
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_rows(
+ ggml_metal_library_t lib,
+ const ggml_tensor * op) {
+ assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+
+ const char * base = "kernel_flash_attn_ext_rows";
+
+ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base);
+ if (!res.pipeline) {
+ res = ggml_metal_library_compile_pipeline(lib, base, base, nullptr);
+ }
+
+ GGML_UNUSED(op);
+
+ return res;
+}
+
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
ggml_metal_library_t lib,
const ggml_tensor * op) {
@@ -2016,6 +2033,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
bool has_scap,
bool has_kvpad,
bool has_sparse,
+ bool has_rows,
int32_t nqpsg,
int32_t ne,
int32_t nsg,
@@ -2045,7 +2063,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
dv,
qne_suffix);
- snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
+ snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_rows=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
base,
has_mask,
has_sinks,
@@ -2053,6 +2071,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
has_scap,
has_kvpad,
has_sparse,
+ has_rows,
ns10,
ns20,
nsg, nwg);
@@ -2067,6 +2086,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
ggml_metal_cv_set_bool(cv, has_scap, FC_FLASH_ATTN_EXT_VEC + 3);
ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT_VEC + 4);
ggml_metal_cv_set_bool(cv, has_sparse, FC_FLASH_ATTN_EXT_VEC + 5);
+ ggml_metal_cv_set_bool(cv, has_rows, FC_FLASH_ATTN_EXT_VEC + 6);
ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20);
ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21);
diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h
index 1858e7fc..e096d611 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.h
+++ b/ggml/src/ggml-metal/ggml-metal-device.h
@@ -219,6 +219,10 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
int32_t ns10,
int32_t ns20);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_rows(
+ ggml_metal_library_t lib,
+ const struct ggml_tensor * op);
+
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
ggml_metal_library_t lib,
const struct ggml_tensor * op);
@@ -232,6 +236,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
bool has_scap,
bool has_kvpad,
bool has_sparse,
+ bool has_rows,
int32_t nqpsg,
int32_t ne,
int32_t nsg,
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index 19ed68f9..68a9addb 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -1727,6 +1727,39 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
case GGML_OP_ROLL:
return ggml_is_contiguous(op->src[0]);
case GGML_OP_FLASH_ATTN_EXT:
+ // K/V row lists (ggml_flash_attn_ext_rows): one list of cache rows per q slice
+ if (op->src[5] != NULL) {
+ const struct ggml_tensor * kv_rows = op->src[5];
+
+ if (kv_rows->type != GGML_TYPE_I32 ||
+ !ggml_is_contiguous(kv_rows) ||
+ kv_rows->ne[2] != 1 || kv_rows->ne[3] != 1 ||
+ kv_rows->ne[1] != op->src[0]->ne[3] ||
+ !op->src[3] || op->src[1]->ne[3] != 1 || op->src[2]->ne[3] != 1) {
+ return false;
+ }
+
+ // the non-vec kernels attend over a contiguous F16 copy of the referenced rows:
+ // - F16 K/V is gathered directly
+ // - quantized K/V is first dequantized with kernel_flash_attn_ext_kv_*_f16
+ // - other types keep attending to the whole cache (dense mask)
+ switch (op->src[1]->type) {
+ case GGML_TYPE_F16:
+ case GGML_TYPE_Q4_0:
+ case GGML_TYPE_Q4_1:
+ case GGML_TYPE_Q5_0:
+ case GGML_TYPE_Q5_1:
+ case GGML_TYPE_Q8_0:
+ break;
+ default:
+ return false;
+ }
+
+ // the row gather copies F16 values in groups of 4
+ if (op->src[2]->ne[0] % 4 != 0) {
+ return false;
+ }
+ }
// for new head sizes, add checks here
if (op->src[0]->ne[0] != 32 &&
op->src[0]->ne[0] != 40 &&
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 6e566084..11c15fef 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -385,6 +385,19 @@ typedef struct {
int32_t nblocks;
} ggml_metal_kargs_flash_attn_ext_kv_f16;
+// gather the K/V cache rows referenced by the kv_rows lists into contiguous F16 buffers
+typedef struct {
+ int32_t ne11; // number of rows per list (n_kv_rows)
+ int32_t ne_12_2; // number of K/V heads
+ int32_t ne10; // DK
+ int32_t ne20; // DV, set to 0 to skip gathering V (V is a view of K)
+ uint64_t v_off; // byte offset of the gathered V region in dst
+ uint64_t nb11; // K source strides (ne13 = 1: the cache is shared by all streams)
+ uint64_t nb12;
+ uint64_t nb21; // V source strides (ne23 = 1)
+ uint64_t nb22;
+} ggml_metal_kargs_flash_attn_ext_rows;
+
typedef struct {
int32_t ne11;
int32_t ne_12_2; // assume K and V are same shape
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 35e3dada..cace34bd 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -3261,7 +3261,10 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool
return false;
}
- return op->src[1]->ne[1] % OP_FLASH_ATTN_EXT_TENSOR_NCPSG == 0;
+ // with kv_rows, the attention runs over the compacted row lists (padded to 256), not over the cache
+ const int64_t ne11_attn = op->src[5] ? op->src[5]->ne[0] : op->src[1]->ne[1];
+
+ return ne11_attn % OP_FLASH_ATTN_EXT_TENSOR_NCPSG == 0;
}
// returns the n_kv_max hint if the sparse path is available for this op, or 0 otherwise
@@ -3270,6 +3273,11 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool
static int ggml_metal_op_flash_attn_ext_n_kv_max_sparse(const ggml_tensor * op) {
assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+ // K/V row lists are handled by the dense kernels, not by the sparse path
+ if (op->src[5] != nullptr) {
+ return 0;
+ }
+
int32_t n_kv_max = 0;
memcpy(&n_kv_max, ((const int32_t *) op->op_params) + 4, sizeof(n_kv_max));
@@ -3513,6 +3521,32 @@ size_t ggml_metal_op_flash_attn_ext_extra_idx(const ggml_tensor * op) {
return GGML_PAD(sizeof(int32_t)*(size_t) n_kv_max_padded*ne31*ne32*ne33, 16);
}
+// size of the contiguous F16 buffers holding the K/V rows referenced by the kv_rows lists
+// the attention over these buffers is a standard dense op (no per-tile gather in the main kernel)
+size_t ggml_metal_op_flash_attn_ext_extra_rows(const ggml_tensor * op) {
+ assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+
+ const ggml_tensor * kv_rows = op->src[5];
+
+ if (!kv_rows) {
+ return 0;
+ }
+
+ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
+ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne);
+
+ const int64_t n_rows = kv_rows->ne[0];
+ const int64_t n_slices = kv_rows->ne[1];
+
+ // note: when V is a view of K, the gathered V is read from the gathered K, so only K is gathered
+ const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
+
+ const size_t k_size = GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne10*n_rows*ne12*n_slices, 16);
+ const size_t v_size = v_is_view_of_k ? 0 : GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne20*n_rows*ne22*n_slices, 16);
+
+ return k_size + v_size;
+}
+
int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
@@ -3545,6 +3579,22 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] >= op->src[0]->ne[1] &&
"the Flash-Attention Metal kernel requires the mask to be at least n_queries big");
+ // K/V row lists (ggml_flash_attn_ext_rows): src[5] holds one list of cache rows per q slice
+ // the attention loops then run over the mask columns and read K/V rows through the lists
+ const ggml_tensor * kv_rows = op->src[5];
+ const bool has_rows = kv_rows != nullptr;
+
+ GGML_ASSERT(!has_rows || (kv_rows->type == GGML_TYPE_I32 && ggml_is_contiguous(kv_rows)));
+ GGML_ASSERT(!has_rows || (kv_rows->ne[1] == ne03 && kv_rows->ne[2] == 1 && kv_rows->ne[3] == 1));
+ GGML_ASSERT(!has_rows || op->src[3]);
+
+ // the loop bound over the mask columns / K/V row lists
+ const int64_t ne11_attn = has_rows ? kv_rows->ne[0] : ne11;
+
+ // the row lists are padded by the KV cache to a multiple of 256, so the pad path never triggers
+ GGML_ASSERT(!has_rows || ne11_attn % OP_FLASH_ATTN_EXT_NCPSG == 0);
+ GGML_ASSERT(!has_rows || ne11_attn % OP_FLASH_ATTN_EXT_VEC_NCPSG == 0);
+
float scale;
float max_bias;
float logit_softcap;
@@ -3559,6 +3609,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
const bool has_mask = op->src[3] != NULL;
const bool has_sinks = op->src[4] != NULL;
+
+ GGML_ASSERT(!has_rows || has_mask);
const bool has_bias = max_bias != 0.0f;
const bool has_scap = logit_softcap != 0.0f;
@@ -3595,12 +3647,25 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
const bool use_sparse = n_kv_max_sparse > 0;
const int n_kv_max_padded = use_sparse ? GGML_PAD(n_kv_max_sparse, OP_FLASH_ATTN_EXT_VEC_NCPSG) : 0;
+ const bool use_vec = ggml_metal_op_flash_attn_ext_use_vec(op);
+
+ // with kv_rows, the non-vec path gathers the referenced rows into a contiguous F16 buffer and
+ // attends over it with the F16 kernel - quantized KV is first dequantized with the F16 pass
+ const bool rows_gather = has_rows && !use_vec;
+
// the vec kernels dequantize the KV inline; no need for the F16 dequant pass in the sparse path
- const bool use_kv_f16 = !use_sparse && ggml_metal_op_flash_attn_ext_use_kv_f16(op);
+ const bool dequant_kv = !use_sparse && (rows_gather ?
+ ggml_is_quantized(op->src[1]->type) : ggml_metal_op_flash_attn_ext_use_kv_f16(op));
+
+ const bool use_kv_f16 = !use_sparse && (rows_gather ? true : dequant_kv);
ggml_metal_buffer_id bid_idx = bid_kv_f16;
bid_idx.offs += ggml_metal_op_flash_attn_ext_extra_kv_f16(op);
+ // contiguous K/V buffers holding the rows referenced by the kv_rows lists
+ ggml_metal_buffer_id bid_rows = bid_idx;
+ bid_rows.offs += ggml_metal_op_flash_attn_ext_extra_idx(op);
+
ggml_metal_buffer_id bid_k = bid_src1;
ggml_metal_buffer_id bid_v = bid_src2;
@@ -3613,7 +3678,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
uint64_t nb22_attn = nb22;
uint64_t nb23_attn = nb23;
- if (use_kv_f16) {
+ if (dequant_kv) {
assert(ggml_metal_op_flash_attn_ext_extra_kv_f16(op) != 0);
const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
@@ -3701,6 +3766,70 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
}
}
+ // with kv_rows, gather the cache rows referenced by the lists into contiguous F16 buffers
+ // the attention then runs as a standard dense op over the gathered rows (tensor or half8x8 kernel)
+ if (rows_gather) {
+ assert(ggml_metal_op_flash_attn_ext_extra_rows(op) != 0);
+
+ const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
+
+ GGML_ASSERT(ne20 % 4 == 0);
+
+ const size_t k_size = GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne10*ne11_attn*ne12*ne03, 16);
+
+ ggml_metal_kargs_flash_attn_ext_rows args0 = {
+ /*.ne11 =*/ (int32_t) ne11_attn,
+ /*.ne_12_2 =*/ ne12,
+ /*.ne10 =*/ ne10,
+ /*.ne20 =*/ v_is_view_of_k ? 0 : ne20, // the gathered V is a view of the gathered K
+ /*.v_off =*/ k_size,
+ /*.nb11 =*/ nb11_attn,
+ /*.nb12 =*/ nb12_attn,
+ /*.nb21 =*/ nb21_attn,
+ /*.nb22 =*/ nb22_attn,
+ };
+
+ auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_rows(lib, op);
+ const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline0), 64);
+
+ ggml_metal_encoder_set_pipeline(enc, pipeline0);
+ ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0);
+ ggml_metal_encoder_set_buffer (enc, bid_k, 1);
+ ggml_metal_encoder_set_buffer (enc, bid_v, 2);
+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(kv_rows), 3);
+ ggml_metal_encoder_set_buffer (enc, bid_rows, 4);
+
+ ggml_metal_encoder_dispatch_threadgroups(enc, ne11_attn, ne03, 1, nth, 1, 1);
+
+ ggml_metal_op_concurrency_reset(ctx);
+
+ // contiguous F16 layout of the gathered rows: [DK, n_rows, n_head_kv, n_streams]
+ bid_k = bid_rows;
+
+ nb10_attn = sizeof(ggml_fp16_t);
+ nb11_attn = nb10_attn*ne10;
+ nb12_attn = nb11_attn*ne11_attn;
+ nb13_attn = nb12_attn*ne12;
+
+ // if V is a view of K, the gathered V is read from the gathered K with K's strides
+ if (v_is_view_of_k) {
+ bid_v = bid_k;
+
+ nb20_attn = nb10_attn;
+ nb21_attn = nb11_attn;
+ nb22_attn = nb12_attn;
+ nb23_attn = nb13_attn;
+ } else {
+ bid_v = bid_rows;
+ bid_v.offs += k_size;
+
+ nb20_attn = sizeof(ggml_fp16_t);
+ nb21_attn = nb20_attn*ne20;
+ nb22_attn = nb21_attn*ne11_attn;
+ nb23_attn = nb22_attn*ne22;
+ }
+ }
+
if (!use_sparse && ggml_metal_op_flash_attn_ext_use_tensor(op, props_dev->has_tensor)) {
// tensor API kernel
const int nqptg = ne00 >= 512 ? OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE : OP_FLASH_ATTN_EXT_TENSOR_NQPSG; // queries per threadgroup
@@ -3746,9 +3875,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
/*.nb01 =*/ nb01,
/*.nb02 =*/ nb02,
/*.nb03 =*/ nb03,
- /*.ne11 =*/ ne11,
+ /*.ne11 =*/ (int32_t) ne11_attn,
/*.ne_12_2 =*/ ne12,
- /*.ne_12_3 =*/ ne13,
+ /*.ne_12_3 =*/ has_rows ? ne03 : ne13, // the gathered K/V holds one copy per stream
/*.ns10 =*/ ns10,
/*.nb11 =*/ nb11_attn,
/*.nb12 =*/ nb12_attn,
@@ -3806,7 +3935,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
bool need_sync = false;
- const bool has_kvpad = ne11 % ncpsg != 0;
+ const bool has_kvpad = ne11_attn % ncpsg != 0;
if (has_kvpad) {
assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0);
@@ -3906,9 +4035,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
/*.nb01 =*/ nb01,
/*.nb02 =*/ nb02,
/*.nb03 =*/ nb03,
- /*.ne11 =*/ ne11,
+ /*.ne11 =*/ (int32_t) ne11_attn,
/*.ne_12_2 =*/ ne12,
- /*.ne_12_3 =*/ ne13,
+ /*.ne_12_3 =*/ has_rows ? ne03 : ne13, // the gathered K/V holds one copy per stream
/*.ns10 =*/ ns10,
/*.nb11 =*/ nb11_attn,
/*.nb12 =*/ nb12_attn,
@@ -3974,7 +4103,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
bool need_sync = false;
- const bool has_kvpad = !use_sparse && ne11 % ncpsg != 0;
+ const bool has_kvpad = !use_sparse && ne11_attn % ncpsg != 0;
if (use_sparse) {
assert(ggml_metal_op_flash_attn_ext_extra_idx(op) != 0);
@@ -4057,13 +4186,13 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
// queries: Q*NSG*PAD2(ne00, 128)
// attn + mask: NSG*4*Q*C
// results: 2*NSG*Q*PAD2(ne20, 128)
- // sparse idx: NSG*C ints (only when use_sparse)
+ // sparse idx: NSG*C ints (when use_sparse, or when gathering K/V rows)
const int64_t dk_pad = GGML_PAD(ne00, 128);
const int64_t dv_pad = GGML_PAD(ne20, 128);
auto fa_vec_smem = [&](int64_t nsg, int32_t nqptg) -> size_t {
const size_t smem_half = (size_t) (dk_pad + 4*ncpsg + 2*dv_pad)*nqptg*nsg;
- return GGML_PAD(smem_half*sizeof(ggml_fp16_t) + (use_sparse ? (size_t) nsg*ncpsg*sizeof(int) : 0), 16);
+ return GGML_PAD(smem_half*sizeof(ggml_fp16_t) + ((use_sparse || has_rows) ? (size_t) nsg*ncpsg*sizeof(int) : 0), 16);
};
int64_t nsg = 1;
@@ -4094,7 +4223,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
} else {
nwg = 32;
nsg = 1;
- while (2*nwg*nsg*ncpsg < ne11 && nsg < 4) {
+ while (2*nwg*nsg*ncpsg < ne11_attn && nsg < 4) {
nsg *= 2;
}
}
@@ -4115,7 +4244,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
/*.nb01 =*/ nb01,
/*.nb02 =*/ nb02,
/*.nb03 =*/ nb03,
- /*.ne11 =*/ use_sparse ? n_kv_max_padded : ne11,
+ /*.ne11 =*/ use_sparse ? n_kv_max_padded : (int32_t) ne11_attn,
/*.ne_12_2 =*/ ne12,
/*.ne_12_3 =*/ ne13,
/*.ns10 =*/ ns10,
@@ -4144,7 +4273,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
/*.n_kv_max_padded =*/ n_kv_max_padded,
};
- auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, use_sparse, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
+ auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, use_sparse, has_rows, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
@@ -4155,7 +4284,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_set_buffer (enc, bid_v, 3);
ggml_metal_encoder_set_buffer (enc, bid_src3, 4);
ggml_metal_encoder_set_buffer (enc, bid_src4, 5);
- ggml_metal_encoder_set_buffer (enc, use_sparse ? bid_idx : bid_src0, 8);
+ ggml_metal_encoder_set_buffer (enc, use_sparse ? bid_idx : (has_rows ? ggml_metal_get_buffer_id(kv_rows) : bid_src0), 8);
const size_t smem = fa_vec_smem(nsg, nqptg);
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h
index 0a77d8d6..47889fa8 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.h
+++ b/ggml/src/ggml-metal/ggml-metal-ops.h
@@ -45,6 +45,7 @@ size_t ggml_metal_op_flash_attn_ext_extra_blk(const struct ggml_tensor * op);
size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op);
size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op);
size_t ggml_metal_op_flash_attn_ext_extra_idx(const struct ggml_tensor * op);
+size_t ggml_metal_op_flash_attn_ext_extra_rows(const struct ggml_tensor * op);
size_t ggml_metal_op_argmax_extra(const struct ggml_tensor * op);
diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp
index 1699a713..9bed6421 100644
--- a/ggml/src/ggml-metal/ggml-metal.cpp
+++ b/ggml/src/ggml-metal/ggml-metal.cpp
@@ -300,6 +300,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor);
res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor);
res += ggml_metal_op_flash_attn_ext_extra_idx(tensor);
+ res += ggml_metal_op_flash_attn_ext_extra_rows(tensor);
} break;
case GGML_OP_CUMSUM:
case GGML_OP_ARGSORT:
diff --git a/ggml/src/ggml-metal/kernels/fa_aux.metal b/ggml/src/ggml-metal/kernels/fa_aux.metal
index ebee1bec..6fe47470 100644
--- a/ggml/src/ggml-metal/kernels/fa_aux.metal
+++ b/ggml/src/ggml-metal/kernels/fa_aux.metal
@@ -48,6 +48,80 @@ template [[host_name("kernel_flash_attn_ext_kv_q5_0_f16")]] kernel kernel_flash_
template [[host_name("kernel_flash_attn_ext_kv_q5_1_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q5_1, 32, dequantize_q5_1>;
template [[host_name("kernel_flash_attn_ext_kv_q8_0_f16")]] kernel kernel_flash_attn_ext_kv_f16_t kernel_flash_attn_ext_kv_f16<block_q8_0, 32, dequantize_q8_0>;
+// gather the K/V cache rows referenced by the kv_rows lists into contiguous F16 buffers
+// - one threadgroup per (row, stream); the attention is then run as a standard dense op over the
+// gathered rows, which avoids the per-tile K/V gather in the main kernels
+// - negative rows are padding: filled with zeros, their mask entries are -INF so they contribute nothing
+kernel void kernel_flash_attn_ext_rows(
+ constant ggml_metal_kargs_flash_attn_ext_rows & args,
+ device const char * k_src,
+ device const char * v_src,
+ device const int * rows,
+ device char * dst,
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ ushort tiitg[[thread_index_in_threadgroup]],
+ ushort3 ntg[[threads_per_threadgroup]]) {
+ const int64_t i1 = tgpig[0]; // row index
+ const int64_t i3 = tgpig[1]; // stream
+
+ const int64_t row = rows[i1 + i3*args.ne11];
+
+ // gathered layout: [DK, n_rows, n_head_kv, n_streams] and [DV, n_rows, n_head_kv, n_streams] F16
+ device half4 * k_dst = (device half4 *) (dst + i3*args.ne_12_2*args.ne11*args.ne10*sizeof(half));
+ device half4 * v_dst = (device half4 *) (dst + args.v_off + i3*args.ne_12_2*args.ne11*args.ne20*sizeof(half));
+
+ // note: the K/V cache is shared by all streams (ne13 = ne23 = 1), the row list selects the rows
+ // each thread copies up to 4 half4 values (32 bytes) to keep the memory transactions large
+
+ // K: all heads of this row
+ {
+ const int64_t per_head = args.ne10/4;
+
+ for (int64_t i = 4*tiitg; i < args.ne_12_2*per_head; i += 4*ntg.x) {
+ const int64_t i2 = i/per_head;
+ const int64_t i0 = i%per_head;
+ const int n = (int) min((int64_t) 4, per_head - i0);
+
+ device half4 * dst4 = k_dst + (i2*args.ne11 + i1)*per_head + i0;
+
+ if (row >= 0) {
+ device const half4 * src4 = (device const half4 *) (k_src + row*args.nb11 + i2*args.nb12) + i0;
+ for (int j = 0; j < n; ++j) {
+ dst4[j] = src4[j];
+ }
+ } else {
+ for (int j = 0; j < n; ++j) {
+ dst4[j] = half4(0.0h);
+ }
+ }
+ }
+ }
+
+ // V: all heads of this row
+ {
+ const int64_t per_head = args.ne20/4;
+
+ for (int64_t i = 4*tiitg; i < args.ne_12_2*per_head; i += 4*ntg.x) {
+ const int64_t i2 = i/per_head;
+ const int64_t i0 = i%per_head;
+ const int n = (int) min((int64_t) 4, per_head - i0);
+
+ device half4 * dst4 = v_dst + (i2*args.ne11 + i1)*per_head + i0;
+
+ if (row >= 0) {
+ device const half4 * src4 = (device const half4 *) (v_src + row*args.nb21 + i2*args.nb22) + i0;
+ for (int j = 0; j < n; ++j) {
+ dst4[j] = src4[j];
+ }
+ } else {
+ for (int j = 0; j < n; ++j) {
+ dst4[j] = half4(0.0h);
+ }
+ }
+ }
+ }
+}
+
constant bool FC_flash_attn_ext_pad_has_mask [[function_constant(FC_FLASH_ATTN_EXT_PAD + 0)]];
constant int32_t FC_flash_attn_ext_pad_ncpsg [[function_constant(FC_FLASH_ATTN_EXT_PAD + 25)]];
diff --git a/ggml/src/ggml-metal/kernels/fa_vec_common.metal b/ggml/src/ggml-metal/kernels/fa_vec_common.metal
index 46270ba4..c81e4d09 100644
--- a/ggml/src/ggml-metal/kernels/fa_vec_common.metal
+++ b/ggml/src/ggml-metal/kernels/fa_vec_common.metal
@@ -13,6 +13,7 @@ constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EX
constant int32_t FC_flash_attn_ext_vec_nsg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]];
constant int32_t FC_flash_attn_ext_vec_nwg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 23)]];
constant bool FC_flash_attn_ext_vec_has_sparse [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]];
+constant bool FC_flash_attn_ext_vec_has_rows [[function_constant(FC_FLASH_ATTN_EXT_VEC + 6)]];
template<
typename q4_t, // query types in shared memory
typename k4_t, // key types in shared memory
@@ -90,7 +91,8 @@ kernel void kernel_flash_attn_ext_vec(
threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + SMEM_Q + SMEM_S + 2*sgitg*Q*PV); // scratch buffer for the results
// sparse indices for the current block
- threadgroup int * spidx = FC_flash_attn_ext_vec_has_sparse
+ // note: with kv_rows, the buffer holds the row list of the current slice instead
+ threadgroup int * spidx = (FC_flash_attn_ext_vec_has_sparse || FC_flash_attn_ext_vec_has_rows)
? (threadgroup int *) (shmem_f16 + SMEM) + sgitg*C
: nullptr;
@@ -165,6 +167,12 @@ kernel void kernel_flash_attn_ext_vec(
((int64_t)(iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded + (iq1%args.ne31)*args.n_kv_max_padded;
}
+ // K/V row list of the current slice (ggml_flash_attn_ext_rows)
+ // with kv_rows, the loop below runs over the mask columns (args.ne11) and the K/V cache rows
+ // are read indirectly through the list. negative entries are padding, their mask is -INF
+ device const int * rows_list = FC_flash_attn_ext_vec_has_rows ?
+ (device const int *) idx + (int64_t)(iq3%args.ne33)*args.ne11 : nullptr;
+
float slope = 1.0f;
// ALiBi
@@ -222,12 +230,12 @@ kernel void kernel_flash_attn_ext_vec(
ic = 0;
}
- // load the sparse KV indices for the current block into shared memory
- if (FC_flash_attn_ext_vec_has_sparse) {
+ // load the KV row indices for the current block into shared memory
+ if (FC_flash_attn_ext_vec_has_sparse || FC_flash_attn_ext_vec_has_rows) {
FOR_UNROLL (short ii = 0; ii < C/NW; ++ii) {
const short i = ii*NW + tiisg;
- spidx[i] = pidx[ic + i];
+ spidx[i] = FC_flash_attn_ext_vec_has_rows ? rows_list[ic + i] : pidx[ic + i];
}
simdgroup_barrier(mem_flags::mem_threadgroup);
}
@@ -242,6 +250,15 @@ kernel void kernel_flash_attn_ext_vec(
sm[qq*C + tiisg] = -MAXHALF;
}
}
+ } else if (FC_flash_attn_ext_vec_has_rows) {
+ // the mask is compact: column ic maps to cache row spidx[tiisg]
+ FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
+ if ((iq1*Q + qq) < args.ne01) {
+ sm[qq*C + tiisg] = pm[qq][ic + tiisg];
+ } else {
+ sm[qq*C + tiisg] = -MAXHALF;
+ }
+ }
} else {
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
if ((iq1*Q + qq) < args.ne01) {
@@ -276,7 +293,7 @@ kernel void kernel_flash_attn_ext_vec(
{
device const k4_t * pk4 = nullptr;
- if (!FC_flash_attn_ext_vec_has_sparse) {
+ if (!FC_flash_attn_ext_vec_has_sparse && !FC_flash_attn_ext_vec_has_rows) {
pk4 = (device const k4_t *) (k + ic*args.nb11);
pk4 += ty*NS10/4 + tx;
@@ -291,7 +308,7 @@ kernel void kernel_flash_attn_ext_vec(
// each simdgroup processes Q queries and NE (NW/NL) cache elements
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
- if (FC_flash_attn_ext_vec_has_sparse) {
+ if (FC_flash_attn_ext_vec_has_sparse || FC_flash_attn_ext_vec_has_rows) {
// the KV rows are gathered from the index list; -1 entries are padding
const int i11 = spidx[NE*cc + ty];
if (i11 >= 0) {
@@ -434,7 +451,7 @@ kernel void kernel_flash_attn_ext_vec(
}
}
- if (FC_flash_attn_ext_vec_has_sparse) {
+ if (FC_flash_attn_ext_vec_has_sparse || FC_flash_attn_ext_vec_has_rows) {
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
// the KV rows are gathered from the index list; -1 entries are padding
const int i11 = spidx[NE*cc + ty];
diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp
index 97956bed..9999aac9 100644
--- a/ggml/src/ggml-opencl/ggml-opencl.cpp
+++ b/ggml/src/ggml-opencl/ggml-opencl.cpp
@@ -11849,6 +11849,10 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
case GGML_OP_MEAN:
return op->src[0]->type == GGML_TYPE_F32;
case GGML_OP_FLASH_ATTN_EXT: {
+ // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+ if (op->src[5] != nullptr) {
+ return false;
+ }
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
if (use_fa_bin_kernels_prefill(backend_ctx, op->src[0], op->src[1], op->src[2])) {
return true;
diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp
index 20741b98..88634557 100644
--- a/ggml/src/ggml-openvino/ggml-openvino.cpp
+++ b/ggml/src/ggml-openvino/ggml-openvino.cpp
@@ -1520,6 +1520,11 @@ static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) {
return {false, "FLASH_ATTN_EXT gemma3n pattern on GPU is not supported"};
}
+ // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+ if (op->src[5] != nullptr) {
+ return {false, "FLASH_ATTN_EXT with K/V row lists is not supported"};
+ }
+
if (op->src[4] != nullptr) {
return {false, "FLASH_ATTN_EXT with sinks is not supported"};
}
diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp
index fd4786ec..9ee16423 100644
--- a/ggml/src/ggml-sycl/ggml-sycl.cpp
+++ b/ggml/src/ggml-sycl/ggml-sycl.cpp
@@ -7653,7 +7653,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
case GGML_OP_SOLVE_TRI:
return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K;
case GGML_OP_FLASH_ATTN_EXT:
- return ggml_sycl_flash_attn_ext_supported(device, op);
+ // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+ return op->src[5] == nullptr && ggml_sycl_flash_attn_ext_supported(device, op);
default:
return false;
}
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 2b6561d4..ba857871 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -15682,6 +15682,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
}
case GGML_OP_FLASH_ATTN_EXT:
{
+ // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+ if (op->src[5] != nullptr) {
+ return false;
+ }
bool coopmat2 = device->coopmat2;
uint32_t HSK = op->src[1]->ne[0];
uint32_t HSV = op->src[2]->ne[0];
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index fac6cb60..44ff0b10 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -4598,7 +4598,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
{
// conservative support checks for whether the more resource-intensive shader paths
// can be used, to avoid cases where flash_attn is assigned to the CPU later on
- supports_op = src0->type == GGML_TYPE_F32 &&
+ // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+ supports_op = op->src[5] == nullptr && src0->type == GGML_TYPE_F32 &&
(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16 ||
src1->type == GGML_TYPE_Q4_0 || src1->type == GGML_TYPE_Q8_0) &&
(src2->type == GGML_TYPE_F32 || src2->type == GGML_TYPE_F16 ||
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index a324ec95..3d60c44e 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -5557,20 +5557,30 @@ struct ggml_tensor * ggml_arange(
// ggml_flash_attn_ext
-struct ggml_tensor * ggml_flash_attn_ext(
+static struct ggml_tensor * ggml_flash_attn_ext_impl(
struct ggml_context * ctx,
struct ggml_tensor * q,
struct ggml_tensor * k,
struct ggml_tensor * v,
struct ggml_tensor * mask,
+ struct ggml_tensor * kv_rows,
float scale,
float max_bias,
float logit_softcap) {
GGML_ASSERT(ggml_can_mul_mat(k, q));
// TODO: check if vT can be multiplied by (k*qT)
- GGML_ASSERT(q->ne[3] == k->ne[3]);
- GGML_ASSERT(q->ne[3] == v->ne[3]);
+ if (kv_rows) {
+ // k and v are shared by all q slices, kv_rows picks the rows of each slice
+ GGML_ASSERT(mask && mask->ne[0] == kv_rows->ne[0]);
+ GGML_ASSERT(kv_rows->type == GGML_TYPE_I32 && ggml_is_contiguous(kv_rows));
+ GGML_ASSERT(kv_rows->ne[1] == q->ne[3] && kv_rows->ne[2] == 1 && kv_rows->ne[3] == 1);
+ GGML_ASSERT(k->ne[3] == 1);
+ GGML_ASSERT(v->ne[3] == 1);
+ } else {
+ GGML_ASSERT(q->ne[3] == k->ne[3]);
+ GGML_ASSERT(q->ne[3] == v->ne[3]);
+ }
if (mask) {
GGML_ASSERT(mask->type == GGML_TYPE_F16);
@@ -5597,10 +5607,36 @@ struct ggml_tensor * ggml_flash_attn_ext(
result->src[1] = k;
result->src[2] = v;
result->src[3] = mask;
+ result->src[5] = kv_rows;
return result;
}
+struct ggml_tensor * ggml_flash_attn_ext(
+ struct ggml_context * ctx,
+ struct ggml_tensor * q,
+ struct ggml_tensor * k,
+ struct ggml_tensor * v,
+ struct ggml_tensor * mask,
+ float scale,
+ float max_bias,
+ float logit_softcap) {
+ return ggml_flash_attn_ext_impl(ctx, q, k, v, mask, NULL, scale, max_bias, logit_softcap);
+}
+
+struct ggml_tensor * ggml_flash_attn_ext_rows(
+ struct ggml_context * ctx,
+ struct ggml_tensor * q,
+ struct ggml_tensor * k,
+ struct ggml_tensor * v,
+ struct ggml_tensor * mask,
+ struct ggml_tensor * kv_rows,
+ float scale,
+ float max_bias,
+ float logit_softcap) {
+ return ggml_flash_attn_ext_impl(ctx, q, k, v, mask, kv_rows, scale, max_bias, logit_softcap);
+}
+
void ggml_flash_attn_ext_set_prec(
struct ggml_tensor * a,