Commit 1b6935b22 for llama.cpp
commit 1b6935b227b9f9aa89deea8016c0e208f547446f
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 (#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 b188c0bf9..3ed418140 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 016d2320d..d3a47f839 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 d3a156b12..dbdde6997 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 0683ff881..5f114f587 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 41d948335..beb877941 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 56decfa2b..96b596ea6 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 bd4a7710f..ba4628fec 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 74f282047..9df15128f 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 f1ebd4e8b..5d7ebf966 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 35a9afabf..cdd947c52 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 c8c1947e6..03b5a3b50 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 548b9896a..463d57c7b 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 1858e7fc5..e096d611b 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 19ed68f9a..68a9addb4 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 6e566084a..11c15feff 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 35e3dada2..cace34bd4 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 0a77d8d69..47889fa84 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 1699a7131..9bed6421b 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 ebee1becd..6fe47470f 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 46270ba4b..c81e4d09b 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 97956bed1..9999aac90 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 20741b98e..88634557c 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 fd4786ec5..9ee16423a 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 2b6561d43..ba8578716 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 fac6cb60b..44ff0b10c 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 a324ec951..3d60c44e5 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,
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 833c38fd7..066c6ddc9 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -129,6 +129,12 @@ static const llm_fused_op_probe llm_fused_op_flash_attn_probe = {
/*.n_tokens_per_seq =*/ 1,
};
+static const llm_fused_op_probe llm_fused_op_flash_attn_kv_rows_probe = {
+ /*.op =*/ LLM_FUSED_OP_FLASH_ATTN_KV_ROWS,
+ /*.name =*/ "Flash Attention over K/V rows",
+ /*.n_tokens_per_seq =*/ 1,
+};
+
static const llm_fused_op_probe llm_fused_op_gdn_ar_probe = {
/*.op =*/ LLM_FUSED_OP_GDN_AR,
/*.name =*/ "fused Gated Delta Net (autoregressive)",
@@ -370,6 +376,19 @@ llama_context::llama_context(
cparams.kv_unified = params.kv_unified;
cparams.moe_cache_size = params.moe_cache_size;
+ // resolved after flash attention, see resolve_fused_ops()
+ {
+ cparams.fused_kv_rows = false;
+
+ const char * LLAMA_KV_ROWS_DISABLE = getenv("LLAMA_KV_ROWS_DISABLE");
+ if (!LLAMA_KV_ROWS_DISABLE) {
+ cparams.auto_fkvr = cparams.kv_unified;
+ } else {
+ LLAMA_LOG_INFO("%s: LLAMA_KV_ROWS_DISABLE set, disabling kv_rows\n", __func__);
+ cparams.auto_fkvr = false;
+ }
+ }
+
// initialized later
cparams.pipeline_parallel = false;
cparams.training = false;
@@ -662,6 +681,13 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3
cparams.auto_fa = false;
}
+ // the fallback attends to the whole unified cache with a dense mask
+ if (cparams.auto_fkvr) {
+ cparams.fused_kv_rows = cparams.flash_attn;
+ resolve(llm_fused_op_flash_attn_kv_rows_probe, cparams.fused_kv_rows);
+ cparams.auto_fkvr = false;
+ }
+
if (cparams.auto_fgdn) {
LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", func);
resolve(llm_fused_op_gdn_ar_probe, cparams.fused_gdn_ar);
diff --git a/src/llama-cparams.h b/src/llama-cparams.h
index 0c260f7b6..c02e8bcac 100644
--- a/src/llama-cparams.h
+++ b/src/llama-cparams.h
@@ -50,6 +50,8 @@ struct llama_cparams {
bool auto_fhc;
bool fused_top_k_mask; // use fused top-k sparse attention mask
bool auto_ftkm;
+ bool fused_kv_rows; // unified cache: attend to the cells of each sequence only (ggml_flash_attn_ext_rows)
+ bool auto_fkvr;
bool no_perf;
bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
bool op_offload;
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 43eb46efb..725bdb048 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -65,6 +65,70 @@ static bool can_reuse_kq_mask(
return res;
}
+static bool can_reuse_attn_mask(
+ const llm_attn_mask & attn_mask,
+ const llama_kv_cache_context * mctx,
+ const llama_ubatch & ubatch,
+ const llama_cparams & cparams) {
+ const uint32_t n_kv_rows = cparams.fused_kv_rows ? mctx->get_n_kv_rows() : 0;
+ const uint32_t n_slices = llama_kv_cache::get_n_kv_slices(ubatch);
+ const auto n_tokens = ubatch.n_tokens;
+
+ // the current ubatch uses rows only when rows are enabled and it can split into streams
+ const bool use_rows = n_kv_rows > 0 && n_slices > 0;
+
+ // the previous mask and the current ubatch must agree on the layout
+ if (use_rows != (attn_mask.rows != nullptr)) {
+ return false;
+ }
+
+ if (!use_rows) {
+ return can_reuse_kq_mask(attn_mask.raw, mctx, ubatch, cparams);
+ }
+
+ bool res = true;
+ res &= attn_mask.n_kv == mctx->get_n_kv();
+ res &= attn_mask.rows->ne[0] == n_kv_rows;
+ res &= attn_mask.rows->ne[1] == n_slices;
+ res &= attn_mask.raw->ne[0] == n_kv_rows;
+ res &= attn_mask.raw->ne[1] == n_tokens/n_slices;
+ res &= attn_mask.raw->ne[3] == n_slices;
+
+ return res;
+}
+
+// KQ mask over the cell lists of kv_rows, or over the whole cache if kv_rows is not used
+static llm_attn_mask build_attn_inp_attn_mask(
+ ggml_context * ctx,
+ const llama_kv_cache_context * mctx,
+ const llama_ubatch & ubatch,
+ const llama_cparams & cparams) {
+ llm_attn_mask attn_mask;
+
+ attn_mask.n_kv = mctx->get_n_kv();
+
+ const uint32_t n_kv_rows = cparams.fused_kv_rows ? mctx->get_n_kv_rows() : 0;
+
+ if (n_kv_rows == 0) {
+ attn_mask.rows = nullptr;
+ attn_mask.raw = build_attn_inp_kq_mask(ctx, mctx, ubatch, cparams);
+ } else {
+ const uint32_t n_slices = llama_kv_cache::get_n_kv_slices(ubatch);
+
+ attn_mask.rows = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_kv_rows, n_slices);
+ ggml_set_input(attn_mask.rows);
+ ggml_set_name(attn_mask.rows, "attn_inp_kv_rows");
+
+ attn_mask.raw = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv_rows, ubatch.n_tokens/n_slices, 1, n_slices);
+ ggml_set_input(attn_mask.raw);
+ ggml_set_name(attn_mask.raw, "attn_inp_kq_mask");
+ }
+
+ attn_mask.cnv = attn_mask.raw;
+
+ return attn_mask;
+}
+
// impl
void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) {
@@ -476,7 +540,7 @@ static void print_mask(const T * data, int64_t n_tokens, int64_t n_kv, int64_t n
}
void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) {
- const int64_t n_kv = self_kq_mask->ne[0];
+ const int64_t n_kv = self_attn_mask.raw->ne[0];
const int64_t n_tokens = ubatch->n_tokens;
const auto fill_mask = [&](auto * data, int64_t ne, int n_swa, llama_swa_type swa_type, bool causal) {
@@ -521,21 +585,21 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) {
const bool causal = cparams.causal_attn;
const bool causal_swa = cparams.causal_attn || hparams.non_causal_type == LLAMA_NON_CAUSAL_TYPE_DENSE_ONLY;
- GGML_ASSERT(self_kq_mask);
- GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));
- if (self_kq_mask->type == GGML_TYPE_F16) {
- fill_mask((ggml_fp16_t *) self_kq_mask->data, ggml_nelements(self_kq_mask), 0, LLAMA_SWA_TYPE_NONE, causal);
+ GGML_ASSERT(self_attn_mask.raw);
+ GGML_ASSERT(ggml_backend_buffer_is_host(self_attn_mask.raw->buffer));
+ if (self_attn_mask.raw->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) self_attn_mask.raw->data, ggml_nelements(self_attn_mask.raw), 0, LLAMA_SWA_TYPE_NONE, causal);
} else {
- fill_mask((float *) self_kq_mask->data, ggml_nelements(self_kq_mask), 0, LLAMA_SWA_TYPE_NONE, causal);
+ fill_mask((float *) self_attn_mask.raw->data, ggml_nelements(self_attn_mask.raw), 0, LLAMA_SWA_TYPE_NONE, causal);
}
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
- GGML_ASSERT(self_kq_mask_swa);
- GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer));
- if (self_kq_mask_swa->type == GGML_TYPE_F16) {
- fill_mask((ggml_fp16_t *) self_kq_mask_swa->data, ggml_nelements(self_kq_mask_swa), hparams.n_swa, hparams.swa_type, causal_swa);
+ GGML_ASSERT(self_attn_mask_swa.raw);
+ GGML_ASSERT(ggml_backend_buffer_is_host(self_attn_mask_swa.raw->buffer));
+ if (self_attn_mask_swa.raw->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) self_attn_mask_swa.raw->data, ggml_nelements(self_attn_mask_swa.raw), hparams.n_swa, hparams.swa_type, causal_swa);
} else {
- fill_mask((float *) self_kq_mask_swa->data, ggml_nelements(self_kq_mask_swa), hparams.n_swa, hparams.swa_type, causal_swa);
+ fill_mask((float *) self_attn_mask_swa.raw->data, ggml_nelements(self_attn_mask_swa.raw), hparams.n_swa, hparams.swa_type, causal_swa);
}
}
}
@@ -551,8 +615,8 @@ void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) {
// the mask is left unallocated when the graph only stores K/V without attending
// (e.g. DFlash's KV-injection pass)
- if (self_kq_mask && self_kq_mask->buffer) {
- mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
+ if (self_attn_mask.raw && self_attn_mask.raw->buffer) {
+ mctx->set_input_attn_mask(self_attn_mask, ubatch, cparams.causal_attn);
}
}
@@ -566,7 +630,7 @@ bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) {
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
//res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
- res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(self_attn_mask, mctx, params.ubatch, params.cparams);
return res;
}
@@ -574,7 +638,7 @@ bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) {
void llm_graph_input_attn_k::set_input(const llama_ubatch * ubatch) {
mctx->set_input_k_idxs(self_k_idxs, ubatch);
- mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
+ mctx->set_input_attn_mask(self_attn_mask, ubatch, cparams.causal_attn);
}
bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) {
@@ -588,7 +652,7 @@ bool llm_graph_input_attn_k::can_reuse_impl(const llm_graph_params & params) {
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
- res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(self_attn_mask, mctx, params.ubatch, params.cparams);
return res;
}
@@ -622,7 +686,7 @@ bool llm_graph_input_attn_kv_msa::can_reuse(const llm_graph_params & params) {
res &= self_k_idxs_idx->ne[0] == params.ubatch.n_tokens;
}
- res &= can_reuse_kq_mask(self_kq_mask, this->mctx, params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(self_attn_mask, this->mctx, params.ubatch, params.cparams);
return res;
}
@@ -631,13 +695,13 @@ void llm_graph_input_attn_k_dsa::set_input(const llama_ubatch * ubatch) {
mctx->get_mla()->set_input_k_idxs(self_k_idxs_mla, ubatch);
// leave unallocated when the top-k layers use the indexer mask
- if (self_kq_mask_mla->buffer) {
- mctx->get_mla()->set_input_kq_mask(self_kq_mask_mla, ubatch, cparams.causal_attn);
+ if (self_attn_mask_mla.raw->buffer) {
+ mctx->get_mla()->set_input_attn_mask(self_attn_mask_mla, ubatch, cparams.causal_attn);
}
mctx->get_lid()->set_input_k_idxs(self_k_idxs_lid, ubatch);
- mctx->get_lid()->set_input_kq_mask(self_kq_mask_lid, ubatch, cparams.causal_attn);
+ mctx->get_lid()->set_input_attn_mask(self_attn_mask_lid, ubatch, cparams.causal_attn);
}
bool llm_graph_input_attn_k_dsa::can_reuse(const llm_graph_params & params) {
@@ -652,8 +716,8 @@ bool llm_graph_input_attn_k_dsa::can_reuse_impl(const llm_graph_params & params)
res &= self_k_idxs_mla->ne[0] == params.ubatch.n_tokens;
res &= self_k_idxs_lid->ne[0] == params.ubatch.n_tokens;
- res &= can_reuse_kq_mask(self_kq_mask_mla, mctx->get_mla(), params.ubatch, params.cparams);
- res &= can_reuse_kq_mask(self_kq_mask_lid, mctx->get_lid(), params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(self_attn_mask_mla, mctx->get_mla(), params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(self_attn_mask_lid, mctx->get_lid(), params.ubatch, params.cparams);
return res;
}
@@ -687,9 +751,7 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {
}
// the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live
- if (self_kq_mask && self_kq_mask->buffer) {
- mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
- }
+ mctx->get_base()->set_input_attn_mask(self_attn_mask, ubatch, cparams.causal_attn);
// swa tensors may not be allocated if there are no SWA attention layers
if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
@@ -699,9 +761,7 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {
}
}
- if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
- mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
- }
+ mctx->get_swa()->set_input_attn_mask(self_attn_mask_swa, ubatch, cparams.causal_attn);
}
bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
@@ -717,8 +777,8 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
//res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
}
- if (self_kq_mask && self_kq_mask->buffer) {
- res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams);
+ if (self_attn_mask.raw && self_attn_mask.raw->buffer) {
+ res &= can_reuse_attn_mask(self_attn_mask, mctx->get_base(), params.ubatch, params.cparams);
}
// swa tensors may not be allocated if there are no SWA attention layers
@@ -727,8 +787,8 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
//res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
}
- if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
- res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
+ if (self_attn_mask_swa.raw && self_attn_mask_swa.raw->buffer) {
+ res &= can_reuse_attn_mask(self_attn_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
}
return res;
@@ -741,18 +801,14 @@ void llm_graph_input_attn_k_iswa::set_input(const llama_ubatch * ubatch) {
}
// the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live
- if (self_kq_mask && self_kq_mask->buffer) {
- mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
- }
+ mctx->get_base()->set_input_attn_mask(self_attn_mask, ubatch, cparams.causal_attn);
// swa tensors may not be allocated if there are no SWA attention layers
if (self_k_idxs_swa && self_k_idxs_swa->buffer) {
mctx->get_swa()->set_input_k_idxs(self_k_idxs_swa, ubatch);
}
- if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
- mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
- }
+ mctx->get_swa()->set_input_attn_mask(self_attn_mask_swa, ubatch, cparams.causal_attn);
}
bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) {
@@ -767,8 +823,8 @@ bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) {
res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
}
- if (self_kq_mask && self_kq_mask->buffer) {
- res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams);
+ if (self_attn_mask.raw && self_attn_mask.raw->buffer) {
+ res &= can_reuse_attn_mask(self_attn_mask, mctx->get_base(), params.ubatch, params.cparams);
}
// swa tensors may not be allocated if there are no SWA attention layers
@@ -776,8 +832,8 @@ bool llm_graph_input_attn_k_iswa::can_reuse(const llm_graph_params & params) {
res &= self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;
}
- if (self_kq_mask_swa && self_kq_mask_swa->buffer) {
- res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
+ if (self_attn_mask_swa.raw && self_attn_mask_swa.raw->buffer) {
+ res &= can_reuse_attn_mask(self_attn_mask_swa, mctx->get_swa(), params.ubatch, params.cparams);
}
return res;
@@ -1029,8 +1085,8 @@ void llm_graph_input_dsv4_raw::set_input(const llama_ubatch * ubatch) {
mctx->set_input_k_idxs(self_k_idxs);
}
- if (self_kq_mask && self_kq_mask->buffer) {
- mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
+ if (self_attn_mask.raw && self_attn_mask.raw->buffer) {
+ mctx->set_input_attn_mask(self_attn_mask, ubatch, cparams.causal_attn);
}
}
@@ -1067,8 +1123,8 @@ bool llm_graph_input_dsv4::can_reuse(const llm_graph_params & params) {
if (inp_raw->self_k_idxs && inp_raw->self_k_idxs->buffer) {
res &= inp_raw->self_k_idxs->ne[0] == raw_ctx->get_n_write();
}
- if (inp_raw->self_kq_mask && inp_raw->self_kq_mask->buffer) {
- res &= dsv4_can_reuse_raw_kq_mask(inp_raw->self_kq_mask, raw_ctx, params.ubatch, n_stream);
+ if (inp_raw->self_attn_mask.raw && inp_raw->self_attn_mask.raw->buffer) {
+ res &= dsv4_can_reuse_raw_kq_mask(inp_raw->self_attn_mask.raw, raw_ctx, params.ubatch, n_stream);
}
res &= dsv4_can_reuse_comp_input(inp_csa, plan_csa, params.ubatch.n_tokens, n_stream);
@@ -1079,12 +1135,12 @@ bool llm_graph_input_dsv4::can_reuse(const llm_graph_params & params) {
}
void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) {
- GGML_ASSERT(cross_kq_mask);
+ GGML_ASSERT(cross_attn_mask.raw);
- const int64_t n_enc = cross_kq_mask->ne[0];
+ const int64_t n_enc = cross_attn_mask.raw->ne[0];
const int64_t n_tokens = ubatch->n_tokens;
- GGML_ASSERT(ggml_backend_buffer_is_host(cross_kq_mask->buffer));
+ GGML_ASSERT(ggml_backend_buffer_is_host(cross_attn_mask.raw->buffer));
GGML_ASSERT(!ubatch->equal_seqs()); // TODO: use ubatch->n_seqs instead of failing
const auto fill_mask = [&](auto * data) {
@@ -1107,10 +1163,10 @@ void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) {
}
};
- if (cross_kq_mask->type == GGML_TYPE_F16) {
- fill_mask((ggml_fp16_t *) cross_kq_mask->data);
+ if (cross_attn_mask.raw->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) cross_attn_mask.raw->data);
} else {
- fill_mask((float *) cross_kq_mask->data);
+ fill_mask((float *) cross_attn_mask.raw->data);
}
}
@@ -1118,10 +1174,8 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {
mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);
- // leave the mask unallocated if all attention layers have their own (e.g. qwen4exp QSA)
- if (inp_attn->self_kq_mask->buffer) {
- mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
- }
+ // kv_rows and the mask stay unallocated if all attention layers have their own (e.g. qwen4exp QSA)
+ mctx->get_attn()->set_input_attn_mask(inp_attn->self_attn_mask, ubatch, cparams.causal_attn);
const int64_t n_rs = mctx->get_recr()->get_n_rs();
@@ -1146,7 +1200,7 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {
res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens;
//res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
- res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(inp_attn->self_attn_mask, mctx->get_attn(), params.ubatch, params.cparams);
res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();
@@ -1165,8 +1219,8 @@ void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) {
mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
// leave the mask unallocated if all attention layers have their own (e.g. glm5-next DSA)
- if (inp_attn->self_kq_mask->buffer) {
- mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
+ if (inp_attn->self_attn_mask.raw->buffer) {
+ mctx->get_attn()->set_input_attn_mask(inp_attn->self_attn_mask, ubatch, cparams.causal_attn);
}
const int64_t n_rs = mctx->get_recr()->get_n_rs();
@@ -1191,7 +1245,7 @@ bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) {
res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens;
- res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(inp_attn->self_attn_mask, mctx->get_attn(), params.ubatch, params.cparams);
res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();
@@ -1212,8 +1266,8 @@ void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) {
attn_ctx->get_base()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);
}
- if (inp_attn->self_kq_mask && inp_attn->self_kq_mask->buffer) {
- attn_ctx->get_base()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn);
+ if (inp_attn->self_attn_mask.raw && inp_attn->self_attn_mask.raw->buffer) {
+ attn_ctx->get_base()->set_input_attn_mask(inp_attn->self_attn_mask, ubatch, cparams.causal_attn);
}
// swa tensors may not be allocated if there are no SWA attention layers
@@ -1222,8 +1276,8 @@ void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) {
attn_ctx->get_swa()->set_input_v_idxs(inp_attn->self_v_idxs_swa, ubatch);
}
- if (inp_attn->self_kq_mask_swa && inp_attn->self_kq_mask_swa->buffer) {
- attn_ctx->get_swa()->set_input_kq_mask(inp_attn->self_kq_mask_swa, ubatch, cparams.causal_attn);
+ if (inp_attn->self_attn_mask_swa.raw && inp_attn->self_attn_mask_swa.raw->buffer) {
+ attn_ctx->get_swa()->set_input_attn_mask(inp_attn->self_attn_mask_swa, ubatch, cparams.causal_attn);
}
const int64_t n_rs = mctx->get_recr()->get_n_rs();
@@ -1254,7 +1308,7 @@ bool llm_graph_input_mem_hybrid_iswa::can_reuse(const llm_graph_params & params)
//res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
}
- res &= can_reuse_kq_mask(inp_attn->self_kq_mask, attn_ctx->get_base(), params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(inp_attn->self_attn_mask, attn_ctx->get_base(), params.ubatch, params.cparams);
// swa tensors may not be allocated if there are no SWA attention layers
if (inp_attn->self_k_idxs_swa && inp_attn->self_k_idxs_swa->buffer) {
@@ -1262,7 +1316,7 @@ bool llm_graph_input_mem_hybrid_iswa::can_reuse(const llm_graph_params & params)
//res &= inp_attn->self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there
}
- res &= can_reuse_kq_mask(inp_attn->self_kq_mask_swa, attn_ctx->get_swa(), params.ubatch, params.cparams);
+ res &= can_reuse_attn_mask(inp_attn->self_attn_mask_swa, attn_ctx->get_swa(), params.ubatch, params.cparams);
res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs();
@@ -2897,6 +2951,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(
ggml_tensor * v,
ggml_tensor * kq_b,
ggml_tensor * kq_mask,
+ ggml_tensor * kv_rows,
ggml_tensor * sinks,
ggml_tensor * v_mla,
int64_t n_kv_max,
@@ -2905,7 +2960,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(
const bool v_trans = v->nb[1] > v->nb[2];
// split the batch into streams if needed
- const auto n_stream = k->ne[3];
+ const auto n_stream = kv_rows ? kv_rows->ne[1] : k->ne[3];
// the stream dim steps over one stream's worth of dim-2 (tokens): (ne[2]/n_stream) rows of stride nb[2].
// q->nb[3]/n_stream only equals that for a contiguous q; nope-only MLA (glm5-next) passes a permuted
@@ -2919,6 +2974,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(
ggml_tensor * cur;
const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr;
+ GGML_ASSERT(use_flash_attn || kv_rows == nullptr);
if (use_flash_attn) {
GGML_ASSERT(kq_b == nullptr && "Flash attention does not support KQ bias yet");
@@ -2935,9 +2991,14 @@ ggml_tensor * llm_graph_context::build_attn_mha(
v = ggml_cast(ctx0, v, GGML_TYPE_F16);
}
- cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,
- hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
- res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il});
+ if (kv_rows) {
+ cur = ggml_flash_attn_ext_rows(ctx0, q, k, v, kq_mask, kv_rows, kq_scale, hparams.f_max_alibi_bias,
+ hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
+ } else {
+ cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,
+ hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
+ }
+ res->add_fused_node({kv_rows ? LLM_FUSED_OP_FLASH_ATTN_KV_ROWS : LLM_FUSED_OP_FLASH_ATTN, cur, il});
ggml_flash_attn_ext_add_sinks(cur, sinks);
GGML_ASSERT(n_kv_max >= 0 && n_kv_max <= INT32_MAX);
@@ -3041,20 +3102,20 @@ llm_graph_input_attn_no_cache * llm_graph_context::build_attn_inp_no_cache() con
// note: there is no KV cache, so the KV values are the tokens of the batch, padded like the KV cache (see llama_kv_cache::get_n_kv)
const int64_t n_kv = GGML_PAD(n_tokens, 256);
- inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_kv, n_tokens, 1, 1);
- ggml_set_input(inp->self_kq_mask);
- cb(inp->self_kq_mask, "self_kq_mask", -1);
+ inp->self_attn_mask.raw = ggml_new_tensor_4d(ctx0, type_mask, n_kv, n_tokens, 1, 1);
+ ggml_set_input(inp->self_attn_mask.raw);
+ cb(inp->self_attn_mask.raw, "self_kq_mask", -1);
- inp->self_kq_mask_cnv = inp->self_kq_mask;
+ inp->self_attn_mask.cnv = inp->self_attn_mask.raw;
if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
- inp->self_kq_mask_swa = ggml_new_tensor_4d(ctx0, type_mask, n_kv, n_tokens, 1, 1);
- ggml_set_input(inp->self_kq_mask_swa);
+ inp->self_attn_mask_swa.raw = ggml_new_tensor_4d(ctx0, type_mask, n_kv, n_tokens, 1, 1);
+ ggml_set_input(inp->self_attn_mask_swa.raw);
- inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
+ inp->self_attn_mask_swa.cnv = inp->self_attn_mask_swa.raw;
} else {
- inp->self_kq_mask_swa = nullptr;
- inp->self_kq_mask_swa_cnv = nullptr;
+ inp->self_attn_mask_swa.raw = nullptr;
+ inp->self_attn_mask_swa.cnv = nullptr;
}
return (llm_graph_input_attn_no_cache *) res->add_input(std::move(inp));
@@ -3083,7 +3144,8 @@ ggml_tensor * llm_graph_context::build_attn(
const bool is_swa = hparams.is_swa(il);
- const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
+ const auto & kq_mask = is_swa ? inp->self_attn_mask_swa.cnv : inp->self_attn_mask.cnv;
+ const auto & kv_rows = is_swa ? inp->self_attn_mask_swa.rows : inp->self_attn_mask.rows;
// [TAG_NO_CACHE_PAD]
// TODO: if ubatch.equal_seqs() == true, we can split the three tensors below into ubatch.n_seqs_unq streams
@@ -3104,7 +3166,7 @@ ggml_tensor * llm_graph_context::build_attn(
kq_b = ggml_pad(ctx0, kq_b, n_kv_pad, 0, 0, 0);
}
- ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
@@ -3137,8 +3199,7 @@ static std::unique_ptr<llm_graph_input_attn_kv> build_attn_inp_kv_impl(
inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch);
inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch);
- inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams);
- inp->self_kq_mask_cnv = inp->self_kq_mask;
+ inp->self_attn_mask = build_attn_inp_attn_mask(ctx0, mctx_cur, ubatch, cparams);
}
inp->self_k_rot = mctx_cur->get_k_rot();
@@ -3210,9 +3271,10 @@ ggml_tensor * llm_graph_context::build_attn(
v = mctx_cur->get_v(ctx0, il);
}
- ggml_tensor * kq_mask = inp->get_kq_mask();
+ const auto & kq_mask = inp->self_attn_mask.cnv;
+ const auto & kv_rows = inp->self_attn_mask.rows;
- ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (inp->self_v_rot[il]) {
@@ -3253,8 +3315,8 @@ static std::unique_ptr<llm_graph_input_attn_k> build_attn_inp_k_impl(
inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch);
- inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams);
- inp->self_kq_mask_cnv = inp->self_kq_mask;
+ inp->self_attn_mask.raw = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams);
+ inp->self_attn_mask.cnv = inp->self_attn_mask.raw;
}
return inp;
@@ -3297,13 +3359,14 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
}
- const auto & kq_mask = inp->get_kq_mask();
+ const auto & kq_mask = inp->self_attn_mask.cnv;
+ const auto & kv_rows = inp->self_attn_mask.rows;
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
- ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
@@ -3377,9 +3440,9 @@ ggml_tensor * llm_graph_context::build_attn(
// top_k indexes the cells, and the indexer mask of the same type is built the same way
// let's use it, so when no other layer reads the MLA mask is not allocated
- ggml_tensor * kq_mask = inp->get_kq_mask_mla();
- if (inp->get_kq_mask_lid()->type == kq_mask->type) {
- kq_mask = inp->get_kq_mask_lid();
+ ggml_tensor * kq_mask = inp->self_attn_mask_mla.cnv;
+ if (inp->self_attn_mask_lid.cnv->type == kq_mask->type) {
+ kq_mask = inp->self_attn_mask_lid.cnv;
}
ggml_tensor * kq_mask_top_k = nullptr;
@@ -3420,7 +3483,7 @@ ggml_tensor * llm_graph_context::build_attn(
const int64_t n_embd_v = v_cur ? v_cur->ne[0] : v_mla->ne[0];
ggml_tensor * v = ggml_view_4d(ctx0, k, n_embd_v, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
- ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, top_k->ne[0], kq_scale, il);
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, nullptr, sinks, v_mla, top_k->ne[0], kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
@@ -3501,13 +3564,14 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il));
}
- const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
+ const auto & kq_mask = is_swa ? inp->self_attn_mask_swa.cnv : inp->self_attn_mask.cnv;
+ const auto & kv_rows = is_swa ? inp->self_attn_mask_swa.rows : inp->self_attn_mask.rows;
ggml_tensor * q = q_cur;
ggml_tensor * k = use_kv_cur ? k_cur : mctx_cur->get_k(ctx0, il);
ggml_tensor * v = use_kv_cur ? v_cur : mctx_cur->get_v(ctx0, il);
- ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (v_rot) {
@@ -3571,14 +3635,15 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
}
- const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
+ const auto & kq_mask = is_swa ? inp->self_attn_mask_swa.cnv : inp->self_attn_mask.cnv;
+ const auto & kv_rows = is_swa ? inp->self_attn_mask_swa.rows : inp->self_attn_mask.rows;
// MLA-style attention: the cached K is used as V
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
- ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (k_rot) {
@@ -3604,10 +3669,10 @@ llm_graph_input_attn_cross * llm_graph_context::build_attn_inp_cross() const {
// flash attention requires an f16 mask
const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
- inp->cross_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_enc, n_tokens, 1, 1);
- ggml_set_input(inp->cross_kq_mask);
+ inp->cross_attn_mask.raw = ggml_new_tensor_4d(ctx0, type_mask, n_enc, n_tokens, 1, 1);
+ ggml_set_input(inp->cross_attn_mask.raw);
- inp->cross_kq_mask_cnv = inp->cross_kq_mask;
+ inp->cross_attn_mask.cnv = inp->cross_attn_mask.raw;
return (llm_graph_input_attn_cross *) res->add_input(std::move(inp));
}
@@ -3631,13 +3696,14 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_build_forward_expand(gf, k_cur);
ggml_build_forward_expand(gf, v_cur);
- const auto & kq_mask = inp->get_kq_mask_cross();
+ const auto & kq_mask = inp->cross_attn_mask.cnv;
+ const auto & kv_rows = inp->cross_attn_mask.rows;
ggml_tensor * q = q_cur;
ggml_tensor * k = k_cur;
ggml_tensor * v = v_cur;
- ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
+ ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
cb(cur, "kqv_out", il);
if (wo) {
@@ -3667,8 +3733,8 @@ static std::unique_ptr<llm_graph_input_attn_k_dsa> build_attn_inp_k_dsa_impl(
{
inp->self_k_idxs_mla = mctx_cur->get_mla()->build_input_k_idxs(ctx0, ubatch);
- inp->self_kq_mask_mla = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla(), ubatch, cparams);
- inp->self_kq_mask_mla_cnv = inp->self_kq_mask_mla;
+ inp->self_attn_mask_mla.raw = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla(), ubatch, cparams);
+ inp->self_attn_mask_mla.cnv = inp->self_attn_mask_mla.raw;
}
{
@@ -3678,8 +3744,8 @@ static std::unique_ptr<llm_graph_input_attn_k_dsa> build_attn_inp_k_dsa_impl(
auto cparams_copy = cparams;
cparams_copy.flash_attn = cparams.fused_lid;
- inp->self_kq_mask_lid = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy);
- inp->self_kq_mask_lid_cnv = inp->self_kq_mask_lid;
+ inp->self_attn_mask_lid.raw = build_attn_inp_kq_mask(ctx0, mctx_cur->get_lid(), ubatch, cparams_copy);
+ inp->self_attn_mask_lid.cnv = inp->self_attn_mask_lid.raw;
inp->self_k_rot_lid = mctx_cur->get_lid()->get_k_rot();
}
@@ -3705,8 +3771,8 @@ llm_graph_input_attn_k_dsa_iswa * llm_graph_context::build_attn_inp_k_dsa_iswa()
inp_swa->self_k_idxs = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);
- inp_swa->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
- inp_swa->self_kq_mask_cnv = inp_swa->self_kq_mask;
+ inp_swa->self_attn_mask.raw = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
+ inp_swa->self_attn_mask.cnv = inp_swa->self_attn_mask.raw;
auto inp = std::make_unique<llm_graph_input_attn_k_dsa_iswa>(std::move(inp_dsa), std::move(inp_swa), mctx_cur);
@@ -3727,8 +3793,8 @@ llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_
inp->self_k_idxs = mctx_base->build_input_k_idxs(ctx0, ubatch);
inp->self_v_idxs = mctx_base->build_input_v_idxs(ctx0, ubatch);
- inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_base, ubatch, cparams);
- inp->self_kq_mask_cnv = inp->self_kq_mask;
+ inp->self_attn_mask.raw = build_attn_inp_kq_mask(ctx0, mctx_base, ubatch, cparams);
+ inp->self_attn_mask.cnv = inp->self_attn_mask.raw;
}
inp->self_k_rot = mctx_base->get_k_rot();
@@ -3753,8 +3819,7 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const
inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch);
inp->self_v_idxs = mctx_cur->get_base()->build_input_v_idxs(ctx0, ubatch);
- inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams);
- inp->self_kq_mask_cnv = inp->self_kq_mask;
+ inp->self_attn_mask = build_attn_inp_attn_mask(ctx0, mctx_cur->get_base(), ubatch, cparams);
}
{
@@ -3763,8 +3828,7 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const
inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);
inp->self_v_idxs_swa = mctx_cur->get_swa()->build_input_v_idxs(ctx0, ubatch);
- inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
- inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
+ inp->self_attn_mask_swa = build_attn_inp_attn_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
}
inp->self_k_rot = mctx_cur->get_base()->get_k_rot();
@@ -3784,8 +3848,8 @@ llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const {
{
inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch);
- inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams);
- inp->self_kq_mask_cnv = inp->self_kq_mask;
+ inp->self_attn_mask.raw = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams);
+ inp->self_attn_mask.cnv = inp->self_attn_mask.raw;
}
{
@@ -3793,8 +3857,8 @@ llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const {
inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch);
- inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
- inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa;
+ inp->self_attn_mask_swa.raw = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams);
+ inp->self_attn_mask_swa.cnv = inp->self_attn_mask_swa.raw;
}
inp->self_k_rot = mctx_cur->get_base()->get_k_rot();
@@ -3815,8 +3879,8 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "DSV4 expects SWA raw cache");
inp_raw->self_k_idxs = raw_ctx->build_input_k_idxs(ctx0, ubatch);
- inp_raw->self_kq_mask = dsv4_build_raw_kq_mask(ctx0, raw_ctx, ubatch, cparams, n_stream);
- inp_raw->self_kq_mask_cnv = inp_raw->self_kq_mask;
+ inp_raw->self_attn_mask.raw = dsv4_build_raw_kq_mask(ctx0, raw_ctx, ubatch, cparams, n_stream);
+ inp_raw->self_attn_mask.cnv = inp_raw->self_attn_mask.raw;
inp_raw->self_k_rot = raw_ctx->get_k_rot();
auto inp = std::make_unique<llm_graph_input_dsv4>(cparams, std::move(inp_raw), mctx_cur);
@@ -3999,16 +4063,16 @@ llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa()
inp_attn->self_k_idxs = attn_ctx->get_base()->build_input_k_idxs(ctx0, ubatch);
inp_attn->self_v_idxs = attn_ctx->get_base()->build_input_v_idxs(ctx0, ubatch);
- inp_attn->self_kq_mask = build_attn_inp_kq_mask(ctx0, attn_ctx->get_base(), ubatch, cparams);
- inp_attn->self_kq_mask_cnv = inp_attn->self_kq_mask;
+ inp_attn->self_attn_mask.raw = build_attn_inp_kq_mask(ctx0, attn_ctx->get_base(), ubatch, cparams);
+ inp_attn->self_attn_mask.cnv = inp_attn->self_attn_mask.raw;
}
{
inp_attn->self_k_idxs_swa = attn_ctx->get_swa()->build_input_k_idxs(ctx0, ubatch);
inp_attn->self_v_idxs_swa = attn_ctx->get_swa()->build_input_v_idxs(ctx0, ubatch);
- inp_attn->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, attn_ctx->get_swa(), ubatch, cparams);
- inp_attn->self_kq_mask_swa_cnv = inp_attn->self_kq_mask_swa;
+ inp_attn->self_attn_mask_swa.raw = build_attn_inp_kq_mask(ctx0, attn_ctx->get_swa(), ubatch, cparams);
+ inp_attn->self_attn_mask_swa.cnv = inp_attn->self_attn_mask_swa.raw;
}
inp_attn->self_k_rot = attn_ctx->get_base()->get_k_rot();
diff --git a/src/llama-graph.h b/src/llama-graph.h
index f9497de5a..b528ee11d 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -47,6 +47,7 @@ enum llm_graph_type {
enum llm_fused_op {
LLM_FUSED_OP_FLASH_ATTN,
+ LLM_FUSED_OP_FLASH_ATTN_KV_ROWS,
LLM_FUSED_OP_GDN_AR,
LLM_FUSED_OP_GDN_CH,
LLM_FUSED_OP_LIGHTNING_INDEXER,
@@ -314,6 +315,18 @@ public:
const llama_cross * cross;
};
+// KQ mask and the optional per-sequence cell rows of a unified cache (see ggml_flash_attn_ext_rows)
+// the mask has two shapes depending on whether `rows` is used:
+// rows != nullptr: [n_kv_rows, n_batch/n_slices, 1, n_slices]
+// rows == nullptr: [n_kv, n_batch/n_stream, 1, n_stream]
+// TODO: cnv is redundant (always aliases raw) except dflash overwrites it with a derived block mask [TAG_ATTN_MASK_CNV]
+struct llm_attn_mask {
+ ggml_tensor * raw = nullptr;
+ ggml_tensor * cnv = nullptr; // same shape
+ ggml_tensor * rows = nullptr; // I32 [n_kv_rows, n_slices]
+ uint32_t n_kv = 0;
+};
+
class llm_graph_input_attn_no_cache : public llm_graph_input_i {
public:
llm_graph_input_attn_no_cache(const llama_hparams & hparams, const llama_cparams & cparams) :
@@ -324,14 +337,9 @@ public:
void set_input(const llama_ubatch * ubatch) override;
- ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
- ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
-
// n_tokens == n_batch, n_kv == n_tokens padded to 256
- ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ llm_attn_mask self_attn_mask;
+ llm_attn_mask self_attn_mask_swa;
const llama_hparams hparams;
const llama_cparams cparams;
@@ -356,13 +364,10 @@ public:
ggml_tensor * get_k_idxs() const { return self_k_idxs; }
ggml_tensor * get_v_idxs() const { return self_v_idxs; }
- ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
-
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]
- ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ llm_attn_mask self_attn_mask;
// note: assumes v_rot^2 == I
std::vector<ggml_tensor *> self_k_rot;
@@ -400,12 +405,9 @@ public:
ggml_tensor * get_k_idxs() const { return self_k_idxs; }
- ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
-
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
- ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ llm_attn_mask self_attn_mask;
const llama_hparams hparams;
const llama_cparams cparams;
@@ -435,16 +437,11 @@ public:
ggml_tensor * get_k_idxs_mla() const { return self_k_idxs_mla; }
ggml_tensor * get_k_idxs_lid() const { return self_k_idxs_lid; }
- ggml_tensor * get_kq_mask_mla() const { return self_kq_mask_mla_cnv; }
- ggml_tensor * get_kq_mask_lid() const { return self_kq_mask_lid; }
-
ggml_tensor * self_k_idxs_mla = nullptr; // I64 [n_batch]
ggml_tensor * self_k_idxs_lid = nullptr; // I64 [n_batch]
- ggml_tensor * self_kq_mask_mla = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_mla_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_lid = nullptr; // F32 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_lid_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ llm_attn_mask self_attn_mask_mla;
+ llm_attn_mask self_attn_mask_lid;
std::vector<ggml_tensor *> self_k_rot_lid;
@@ -521,18 +518,13 @@ public:
ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }
ggml_tensor * get_v_idxs_swa() const { return self_v_idxs_swa; }
- ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
- ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
-
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]
ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]
ggml_tensor * self_v_idxs_swa = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa]
- ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ llm_attn_mask self_attn_mask;
+ llm_attn_mask self_attn_mask_swa;
std::vector<ggml_tensor *> self_k_rot;
std::vector<ggml_tensor *> self_v_rot;
@@ -565,16 +557,11 @@ public:
ggml_tensor * get_k_idxs() const { return self_k_idxs; }
ggml_tensor * get_k_idxs_swa() const { return self_k_idxs_swa; }
- ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
- ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
-
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
ggml_tensor * self_k_idxs_swa = nullptr; // I64 [n_batch]
- ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_swa_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ llm_attn_mask self_attn_mask;
+ llm_attn_mask self_attn_mask_swa;
std::vector<ggml_tensor *> self_k_rot;
std::vector<ggml_tensor *> self_k_rot_swa;
@@ -599,12 +586,10 @@ public:
void set_input(const llama_ubatch * ubatch);
ggml_tensor * get_k_idxs() const { return self_k_idxs; }
- ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch]
- ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]
- ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream]
+ llm_attn_mask self_attn_mask;
std::vector<ggml_tensor *> self_k_rot;
@@ -669,10 +654,8 @@ public:
void set_input(const llama_ubatch * ubatch) override;
- ggml_tensor * get_kq_mask_cross() const { return cross_kq_mask_cnv; }
- ggml_tensor * cross_kq_mask = nullptr; // F32/F16 [n_outputs_enc, n_batch, 1, 1]
- ggml_tensor * cross_kq_mask_cnv = nullptr; // F32/F16 [n_outputs_enc, n_batch, 1, 1]
+ llm_attn_mask cross_attn_mask;
const llama_cross * cross = nullptr;
};
@@ -1261,6 +1244,7 @@ struct llm_graph_context {
ggml_tensor * v, // [n_embd_head_v, n_head_v, n_tokens] (v_trans = false)
ggml_tensor * kq_b,
ggml_tensor * kq_mask,
+ ggml_tensor * kv_rows, // [n_kv, n_stream], k and v are shared by the streams
ggml_tensor * sinks, // [n_head_q]
ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v]
int64_t n_kv_max,
diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp
index c02e15bdd..10da95ae8 100644
--- a/src/llama-kv-cache-dsv4.cpp
+++ b/src/llama-kv-cache-dsv4.cpp
@@ -2082,8 +2082,8 @@ void llama_kv_cache_dsv4_raw_context::set_input_k_idxs(ggml_tensor * dst) const
kv_swa->set_input_k_idxs(dst, &ubatches_write[i_next], sinfos_write[i_next]);
}
-void llama_kv_cache_dsv4_raw_context::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const {
- kv_swa->set_input_kq_mask(dst, ubatch, causal_attn);
+void llama_kv_cache_dsv4_raw_context::set_input_attn_mask(const llm_attn_mask & attn_mask, const llama_ubatch * ubatch, bool causal_attn) const {
+ kv_swa->set_input_attn_mask(attn_mask, ubatch, causal_attn);
}
void llama_kv_cache_dsv4_raw_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const {
diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h
index 5adc9fb43..1c8d23ed2 100644
--- a/src/llama-kv-cache-dsv4.h
+++ b/src/llama-kv-cache-dsv4.h
@@ -226,7 +226,7 @@ public:
std::vector<ggml_tensor *> get_k_rot() const;
void set_input_k_idxs(ggml_tensor * dst) const;
- void set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
+ void set_input_attn_mask(const llm_attn_mask & attn_mask, const llama_ubatch * ubatch, bool causal_attn) const;
// token history of the raw (window) cells, for host-side n-gram hashing (deepseek41 Engram)
void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector<llama_token> & res) const;
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 6d5a57633..84f547ae1 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -1324,6 +1324,98 @@ uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
return result;
}
+uint32_t llama_kv_cache::get_n_kv_slices(const llama_ubatch & ubatch) {
+ if (ubatch.n_seqs_unq == 1) {
+ return 1;
+ }
+
+ // each sequence set of an equal split is a single sequence
+ if (ubatch.equal_seqs() && ubatch.n_seqs == ubatch.n_seqs_unq) {
+ return ubatch.n_seqs;
+ }
+
+ // runs of equal length with a single sequence each, f.ex. one token per sequence or the drafts of several slots
+ // there are n_seqs_unq runs, so each run has a different sequence
+ if (ubatch.n_tokens % ubatch.n_seqs_unq == 0) {
+ const uint32_t n_tps = ubatch.n_tokens/ubatch.n_seqs_unq;
+
+ bool is_runs = true;
+ for (uint32_t i = 0; i < ubatch.n_tokens && is_runs; ++i) {
+ is_runs = ubatch.n_seq_id[i] == 1 && ubatch.seq_id[i][0] == ubatch.seq_id[i - i%n_tps][0];
+ }
+
+ if (is_runs) {
+ return ubatch.n_seqs_unq;
+ }
+ }
+
+ // cannot split into per-sequence slices - a single slice over the cells of all the sequences of the ubatch
+ return 1;
+}
+
+// per slice: the sequences whose cells form the list of the slice and the min position of the cells that are not masked by the SWA
+static std::vector<llama_kv_rows_slice> get_kv_rows_slices(const llama_ubatch & ubatch, uint32_t n_slices, uint32_t n_swa, llama_swa_type swa_type) {
+ std::vector<llama_kv_rows_slice> res(n_slices);
+
+ if (n_slices == 1) {
+ res[0].seq_ids.assign(ubatch.seq_id_unq, ubatch.seq_id_unq + ubatch.n_seqs_unq);
+ } else {
+ const uint32_t n_tps = ubatch.n_tokens/n_slices;
+
+ for (uint32_t s = 0; s < n_slices; ++s) {
+ res[s].seq_ids = { ubatch.seq_id[s*n_tps][0] };
+ }
+ }
+
+ // the cells below these positions are masked by the SWA for all tokens of the slice
+ // only for SWA types that mask more cells as the token position grows
+ if (n_swa == 0 || (swa_type != LLAMA_SWA_TYPE_STANDARD && swa_type != LLAMA_SWA_TYPE_CHUNKED)) {
+ return res;
+ }
+
+ const uint32_t n_tps = ubatch.n_tokens/n_slices;
+
+ for (uint32_t s = 0; s < n_slices; ++s) {
+ llama_pos p = ubatch.pos[s*n_tps];
+ for (uint32_t i = 1; i < n_tps; ++i) {
+ p = std::min(p, ubatch.pos[s*n_tps + i]);
+ }
+
+ res[s].p_min = swa_type == LLAMA_SWA_TYPE_STANDARD ? p - (llama_pos) n_swa + 1 : (p/(llama_pos) n_swa)*(llama_pos) n_swa;
+ }
+
+ return res;
+}
+
+uint32_t llama_kv_cache::get_n_kv_rows(const llama_ubatch & ubatch) const {
+ if (n_stream != 1) {
+ return 0;
+ }
+
+ const uint32_t n_slices = get_n_kv_slices(ubatch);
+
+ const auto & cells = v_cells[0];
+
+ const uint32_t n_pad_cur = 256;
+
+ uint32_t n_max = 0;
+
+ const auto slices = get_kv_rows_slices(ubatch, n_slices, n_swa, swa_type);
+
+ for (uint32_t s = 0; s < n_slices; ++s) {
+ // a cell shared by several sequences of the slice is counted once per sequence, an upper bound is enough
+ uint32_t n = 0;
+ for (const llama_seq_id seq_id : slices[s].seq_ids) {
+ n += slices[s].p_min < 0 ? cells.seq_n_cells(seq_id) : cells.seq_n_cells(seq_id, slices[s].p_min);
+ }
+
+ n_max = std::max(n_max, std::min(n, cells.used_max_p1()));
+ }
+
+ // pad like n_kv so that the graph can be reused
+ return std::min(GGML_PAD(n_max, n_pad_cur), get_size());
+}
+
ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const {
const int32_t ikv = map_layer_ids.at(il);
@@ -1574,7 +1666,7 @@ void llama_kv_cache::set_input_k_shift(ggml_tensor * dst) const {
}
}
-struct args_set_input_kq_mask {
+struct args_set_input_attn_mask {
const llama_hparams & hparams;
const llama_ubatch * ubatch;
@@ -1587,10 +1679,13 @@ struct args_set_input_kq_mask {
int64_t n_kv;
int64_t n_stream;
int64_t n_tps;
+
+ // cell of each mask column per stream, nullptr if the columns are the cells
+ const int32_t * rows;
};
template<typename T, bool causal, bool swa, bool is_2d, bool alibi>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
+static void set_input_attn_mask_impl(const args_set_input_attn_mask & args, T * data) {
//const auto & hparams = args.hparams;
const auto & ubatch = args.ubatch;
@@ -1604,6 +1699,8 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data
const int64_t n_stream = args.n_stream;
const int64_t n_tps = args.n_tps;
+ const int32_t * rows_all = args.rows;
+
const T mask_keep = llama_cast<T>(0.0f);
const T mask_drop = llama_cast<T>(-INFINITY);
@@ -1622,6 +1719,8 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data
std::unordered_map<llama_seq_id, uint32_t> seq_srct;
std::unordered_map<llama_seq_id, std::vector<uint32_t>> seq_idxs;
+ const int32_t * rows = rows_all ? rows_all + s*n_kv : nullptr;
+
for (uint32_t ii = 0; ii < n_tps; ++ii) {
const uint32_t i = s*n_tps + ii;
@@ -1666,7 +1765,8 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data
}
for (uint32_t jj = 0; jj < n_kv; ++jj) {
- uint32_t j = jj;
+ // column of the mask
+ uint32_t c = jj;
// we have an exiting mask for this sequence -> update just seq_idxs
if (!alibi) {
@@ -1675,11 +1775,14 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data
break;
}
- j = idxs[jj];
+ c = idxs[jj];
}
}
- if (cells.is_empty(j)) {
+ // cell of the column, negative for padding
+ const int32_t j = rows ? rows[c] : (int32_t) c;
+
+ if (j < 0 || cells.is_empty(j)) {
goto skip;
}
@@ -1694,7 +1797,7 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data
if (!prev) {
// record all cells for which: p0 >= seq_pos_min[seq_id] - n_swa - 32
if (p0 + (int32_t) (n_swa + 32) >= seq_pos_min[seq_id]) {
- idxs.push_back(j);
+ idxs.push_back(c);
}
}
}
@@ -1727,103 +1830,128 @@ static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data
}
if (alibi) {
- data[idst + j] = llama_cast<T>(static_cast<float>(-std::abs(p0 - p1)));
+ data[idst + c] = llama_cast<T>(static_cast<float>(-std::abs(p0 - p1)));
} else {
- data[idst + j] = mask_keep;
+ data[idst + c] = mask_keep;
}
continue;
skip:
- data[idst + j] = mask_drop;
+ data[idst + c] = mask_drop;
}
}
}
}
template<typename T, bool causal, bool swa, bool is_2d>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
+static void set_input_attn_mask_impl(const args_set_input_attn_mask & args, T * data) {
const bool alibi = args.hparams.use_alibi;
if (alibi) {
- set_input_kq_mask_impl<T, causal, swa, is_2d, true> (args, data);
+ set_input_attn_mask_impl<T, causal, swa, is_2d, true> (args, data);
} else {
- set_input_kq_mask_impl<T, causal, swa, is_2d, false>(args, data);
+ set_input_attn_mask_impl<T, causal, swa, is_2d, false>(args, data);
}
}
template<typename T, bool causal, bool swa>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
+static void set_input_attn_mask_impl(const args_set_input_attn_mask & args, T * data) {
const bool is_2d = args.ubatch->is_pos_2d();
if (is_2d) {
- set_input_kq_mask_impl<T, causal, swa, true> (args, data);
+ set_input_attn_mask_impl<T, causal, swa, true> (args, data);
} else {
- set_input_kq_mask_impl<T, causal, swa, false>(args, data);
+ set_input_attn_mask_impl<T, causal, swa, false>(args, data);
}
}
template<typename T, bool causal>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data) {
+static void set_input_attn_mask_impl(const args_set_input_attn_mask & args, T * data) {
const bool swa = args.swa_type != LLAMA_SWA_TYPE_NONE;
if (swa) {
- set_input_kq_mask_impl<T, causal, true> (args, data);
+ set_input_attn_mask_impl<T, causal, true> (args, data);
} else {
- set_input_kq_mask_impl<T, causal, false>(args, data);
+ set_input_attn_mask_impl<T, causal, false>(args, data);
}
}
template<typename T>
-static void set_input_kq_mask_impl(const args_set_input_kq_mask & args, T * data, bool causal_attn) {
+static void set_input_attn_mask_impl(const args_set_input_attn_mask & args, T * data, bool causal_attn) {
if (causal_attn) {
- set_input_kq_mask_impl<T, true> (args, data);
+ set_input_attn_mask_impl<T, true> (args, data);
} else {
- set_input_kq_mask_impl<T, false>(args, data);
+ set_input_attn_mask_impl<T, false>(args, data);
}
}
-void llama_kv_cache::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const {
- const uint32_t n_tokens = ubatch->n_tokens;
+void llama_kv_cache::set_input_attn_mask(const llm_attn_mask & attn_mask, const llama_ubatch * ubatch, bool causal_attn) const {
+ // fill the cell rows of each slice
+ if (attn_mask.rows && attn_mask.rows->buffer) {
+ GGML_ASSERT(ggml_backend_buffer_is_host(attn_mask.rows->buffer));
+ GGML_ASSERT(n_stream == 1);
- GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
+ const int64_t n_rows = attn_mask.rows->ne[0];
+ const int64_t n_slices = attn_mask.rows->ne[1];
- const int64_t n_kv = dst->ne[0];
- const int64_t n_stream = dst->ne[3]; // num streams in the current ubatch
+ const auto & cells = v_cells[0];
+ int32_t * data = (int32_t *) attn_mask.rows->data;
- GGML_ASSERT(n_tokens%n_stream == 0);
+ cells.get_seq_cells(get_kv_rows_slices(*ubatch, n_slices, n_swa, swa_type), data, n_rows);
+ }
- // n_tps == n_tokens_per_stream
- const int64_t n_tps = n_tokens/n_stream;
+ // the mask is left unallocated when the graph only stores K/V without attending
+ if (attn_mask.raw && attn_mask.raw->buffer) {
+ ggml_tensor * dst = attn_mask.raw;
+ const int32_t * rows = attn_mask.rows && attn_mask.rows->buffer ? (const int32_t *) attn_mask.rows->data : nullptr;
- // see llama_non_causal_type
- // only the SWA cache (or the SWA layers of a single cache) become non-causal
- if (!causal_attn && hparams.non_causal_type == LLAMA_NON_CAUSAL_TYPE_SWA_ONLY) {
- causal_attn = swa_type == LLAMA_SWA_TYPE_NONE;
- }
- if (!causal_attn && hparams.non_causal_type == LLAMA_NON_CAUSAL_TYPE_DENSE_ONLY) {
- causal_attn = swa_type != LLAMA_SWA_TYPE_NONE;
- }
+ const uint32_t n_tokens = ubatch->n_tokens;
- //const int64_t t_start = ggml_time_us();
+ const int64_t n_kv = dst->ne[0];
+ const int64_t n_stream = dst->ne[3]; // num streams in the current ubatch
- const args_set_input_kq_mask args = {
- /*.hparams =*/ hparams,
- /*.ubatch =*/ ubatch,
- /*.v_cells =*/ v_cells,
- /*.seq_to_stream =*/ seq_to_stream,
- /*.n_swa =*/ n_swa,
- /*.swa_type =*/ swa_type,
- /*.n_kv =*/ n_kv,
- /*.n_stream =*/ n_stream,
- /*.n_tps =*/ n_tps,
- };
+ GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
- if (dst->type == GGML_TYPE_F16) {
- set_input_kq_mask_impl<ggml_fp16_t>(args, (ggml_fp16_t *) dst->data, causal_attn);
- } else {
- set_input_kq_mask_impl<float>(args, (float *) dst->data, causal_attn);
- }
+ if (rows) {
+ GGML_ASSERT(attn_mask.rows->ne[0] == n_kv && attn_mask.rows->ne[1] == n_stream);
+ }
+
+ GGML_ASSERT(n_tokens%n_stream == 0);
+
+ // n_tps == n_tokens_per_stream
+ const int64_t n_tps = n_tokens/n_stream;
+
+ // see llama_non_causal_type
+ // only the SWA cache (or the SWA layers of a single cache) become non-causal
+ if (!causal_attn && hparams.non_causal_type == LLAMA_NON_CAUSAL_TYPE_SWA_ONLY) {
+ causal_attn = swa_type == LLAMA_SWA_TYPE_NONE;
+ }
+ if (!causal_attn && hparams.non_causal_type == LLAMA_NON_CAUSAL_TYPE_DENSE_ONLY) {
+ causal_attn = swa_type != LLAMA_SWA_TYPE_NONE;
+ }
+
+ //const int64_t t_start = ggml_time_us();
- //const int64_t t_end = ggml_time_us();
+ const args_set_input_attn_mask args = {
+ /*.hparams =*/ hparams,
+ /*.ubatch =*/ ubatch,
+ /*.v_cells =*/ v_cells,
+ /*.seq_to_stream =*/ seq_to_stream,
+ /*.n_swa =*/ n_swa,
+ /*.swa_type =*/ swa_type,
+ /*.n_kv =*/ n_kv,
+ /*.n_stream =*/ n_stream,
+ /*.n_tps =*/ n_tps,
+ /*.rows =*/ rows,
+ };
- //LLAMA_LOG_ERROR("%s: kq mask time: %0.3f ms\n", __func__, (t_end - t_start)/1000.0);
+ if (dst->type == GGML_TYPE_F16) {
+ set_input_attn_mask_impl<ggml_fp16_t>(args, (ggml_fp16_t *) dst->data, causal_attn);
+ } else {
+ set_input_attn_mask_impl<float>(args, (float *) dst->data, causal_attn);
+ }
+
+ //const int64_t t_end = ggml_time_us();
+
+ //LLAMA_LOG_ERROR("%s: kq mask time: %0.3f ms\n", __func__, (t_end - t_start)/1000.0);
+ }
}
void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const {
@@ -2815,6 +2943,9 @@ llama_kv_cache_context::llama_kv_cache_context(
llama_kv_cache * kv) : status(LLAMA_MEMORY_STATUS_SUCCESS), kv(kv) {
n_kv = kv->get_size();
+ // worst case for graph reservation, only a unified cache splits into per-sequence lists
+ n_kv_rows = kv->get_n_stream() == 1 ? n_kv : 0;
+
const uint32_t n_stream = kv->get_n_stream();
// create a dummy slot info - the actual data is irrelevant. we just need to build the graph
@@ -2867,7 +2998,8 @@ bool llama_kv_cache_context::apply() {
}
kv->apply_ubatch(sinfos[i_cur], ubatches[i_cur]);
- n_kv = kv->get_n_kv(sinfos[i_cur]);
+ n_kv = kv->get_n_kv(sinfos[i_cur]);
+ n_kv_rows = kv->get_n_kv_rows(ubatches[i_cur]);
return true;
}
@@ -2886,6 +3018,10 @@ uint32_t llama_kv_cache_context::get_n_kv() const {
return n_kv;
}
+uint32_t llama_kv_cache_context::get_n_kv_rows() const {
+ return n_kv_rows;
+}
+
ggml_type llama_kv_cache_context::type_k() const {
return kv->type_k();
}
@@ -2938,8 +3074,8 @@ void llama_kv_cache_context::set_input_v_idxs(ggml_tensor * dst, const llama_uba
kv->set_input_v_idxs(dst, ubatch, sinfos[i_cur]);
}
-void llama_kv_cache_context::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const {
- kv->set_input_kq_mask(dst, ubatch, causal_attn);
+void llama_kv_cache_context::set_input_attn_mask(const llm_attn_mask & attn_mask, const llama_ubatch * ubatch, bool causal_attn) const {
+ kv->set_input_attn_mask(attn_mask, ubatch, causal_attn);
}
void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const {
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index 6dfd0c524..9ed36c562 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -180,8 +180,8 @@ public:
llama_io_read_i & io,
llama_seq_id seq_id,
llama_state_seq_flags flags,
- slot_info_vec_t * sinfos_out,
- const slot_info_vec_t * sinfos_in);
+ slot_info_vec_t * sinfos_out,
+ const slot_info_vec_t * sinfos_in);
// undo a state_read() of seq_id (-1 for the whole cache) that another memory module failed to complete
void state_clear(llama_seq_id seq_id);
@@ -192,6 +192,15 @@ public:
uint32_t get_n_kv(const slot_info & sinfo) const;
+ // attention over the cells of each sequence of the ubatch instead of the whole cache (see ggml_flash_attn_ext_rows)
+ // the ubatch splits into n_slices groups of tokens with one sequence each
+ // if it cannot (f.ex. interleaved or coupled sequences), it is a single slice over the cells of all its sequences
+ // TODO: allow slices with coupled sequences
+ static uint32_t get_n_kv_slices(const llama_ubatch & ubatch);
+
+ // padded length of the longest cell list of the slices; 0 if kv_rows are not used (multiple streams)
+ uint32_t get_n_kv_rows(const llama_ubatch & ubatch) const;
+
// get views of the current state of the cache
ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
@@ -233,7 +242,7 @@ public:
void set_input_k_shift(ggml_tensor * dst) const;
- void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
+ void set_input_attn_mask (const llm_attn_mask & attn_mask, const llama_ubatch * ubatch, bool causal_attn) const;
void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const;
// true if llama_kv_cell_ext holds information that has to survive a state save/restore
@@ -401,6 +410,7 @@ public:
//
uint32_t get_n_kv() const;
+ uint32_t get_n_kv_rows() const;
ggml_type type_k() const;
ggml_type type_v() const;
@@ -431,7 +441,7 @@ public:
void set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ubatch) const;
void set_input_k_shift (ggml_tensor * dst) const;
- void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
+ void set_input_attn_mask (const llm_attn_mask & attn_mask, const llama_ubatch * ubatch, bool causal_attn) const;
void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const;
// see llama_kv_cache::get_prev_tokens()
@@ -469,4 +479,7 @@ private:
// a heuristic, to avoid attending the full cache if it is not yet utilized
// as the cache gets filled, the benefit from this heuristic disappears
int32_t n_kv;
+
+ // see llama_kv_cache::get_n_kv_rows
+ uint32_t n_kv_rows = 0;
};
diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h
index 9c0d51e4e..85a8e0f75 100644
--- a/src/llama-kv-cells.h
+++ b/src/llama-kv-cells.h
@@ -3,13 +3,37 @@
#include "llama.h"
#include "llama-cparams.h"
+#include <algorithm>
#include <bitset>
#include <cassert>
+#include <cstdint>
#include <cstring>
+#include <iterator>
#include <limits>
#include <set>
#include <vector>
+#if defined(__GNUC__)
+# define LLAMA_ATTRIBUTE_NOINLINE __attribute__((noinline))
+#else
+# define LLAMA_ATTRIBUTE_NOINLINE
+#endif
+
+#if defined(_MSC_VER)
+# include <intrin.h>
+#endif
+
+// index of the lowest set bit, x != 0
+static inline int ctz64(uint64_t x) {
+#if defined(_MSC_VER)
+ unsigned long i;
+ _BitScanForward64(&i, x);
+ return (int) i;
+#else
+ return __builtin_ctzll(x);
+#endif
+}
+
struct llama_kv_cell_ext {
// 2D spatial positions, typically used for M-RoPE
llama_pos x = 0;
@@ -31,6 +55,14 @@ struct llama_kv_cell_ext {
}
};
+// the data of one slice of a ubatch, used to build the cell rows of the kv_rows attention (see ggml_flash_attn_ext_rows)
+struct llama_kv_rows_slice {
+ std::vector<llama_seq_id> seq_ids; // the sequences whose cells form the list of the slice
+
+ // min position of the cells that are not masked by the SWA (-1 if all cells are kept)
+ llama_pos p_min = -1;
+};
+
// meta information about KV cells that can be part of multiple sequences at the same time
// TODO: add unit tests
class llama_kv_cells {
@@ -51,6 +83,7 @@ public:
for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) {
seq_pos[s].clear();
+ std::fill(seq_bits[s].begin(), seq_bits[s].end(), 0);
}
}
@@ -72,6 +105,10 @@ public:
shift.resize(n);
seq.resize(n);
+ for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) {
+ seq_bits[s].clear();
+ }
+
reset();
}
@@ -368,6 +405,63 @@ public:
return -1;
}
+ // number of cells that contain seq_id
+ uint32_t seq_n_cells(llama_seq_id seq_id) const {
+ assert(seq_id >= 0);
+ assert(seq_id < LLAMA_MAX_SEQ);
+
+ return seq_pos[seq_id].size();
+ }
+
+ // number of cells that contain seq_id at a position >= p_min
+ uint32_t seq_n_cells(llama_seq_id seq_id, llama_pos p_min) const {
+ assert(seq_id >= 0);
+ assert(seq_id < LLAMA_MAX_SEQ);
+
+ const auto & sp = seq_pos[seq_id];
+
+ return std::distance(sp.lower_bound({ p_min, 0 }), sp.end());
+ }
+
+ // get the lists of the cells that contain any of the sequences of each slice, in cell order
+ // one list of max `n_per_seq` for each slice, padded with -1
+ // the cells of slice s at a position below slices[s].p_min are skipped
+ LLAMA_ATTRIBUTE_NOINLINE
+ void get_seq_cells(const std::vector<llama_kv_rows_slice> & slices, int32_t * data, int64_t n_per_seq) const {
+ assert(slices.size() >= 1);
+
+ const int64_t n_slices = (int64_t) slices.size();
+
+ const uint32_t n_words = (used_max_p1() + 63)/64;
+
+ for (int64_t s = 0; s < n_slices; ++s) {
+ const llama_pos pm = slices[s].p_min;
+
+ int32_t * dst = data + s*n_per_seq;
+ int32_t * end = dst + n_per_seq;
+
+ for (uint32_t w = 0; w < n_words; ++w) {
+ uint64_t word = 0;
+ for (const llama_seq_id seq_id : slices[s].seq_ids) {
+ const auto & bits = seq_bits[seq_id];
+ if (w < bits.size()) {
+ word |= bits[w];
+ }
+ }
+
+ for (uint64_t b = word; b != 0; b &= b - 1) {
+ const int32_t i = 64*w + ctz64(b);
+ if (pos[i] >= pm) {
+ GGML_ASSERT(dst < end);
+ *dst++ = i;
+ }
+ }
+ }
+
+ std::fill(dst, end, -1);
+ }
+ }
+
// the minimum position of sequence seq_id currently present in any of the cells
// return -1 if the sequence is not present
llama_pos seq_pos_min(llama_seq_id seq_id) const {
@@ -531,16 +625,28 @@ private:
//
std::set<std::pair<llama_pos, uint32_t>> seq_pos[LLAMA_MAX_SEQ];
- // helper functions for updating `seq_pos`, once cell at a time:
+ // bit i of seq_bits[s] is set if cell i carries sequence s (the transpose of `seq`), used to build the cell lists
+ // allocated on the first cell of the sequence, so the copies made by cp() never allocate it
+ std::vector<uint64_t> seq_bits[LLAMA_MAX_SEQ];
+
+ // helper functions for updating `seq_pos` and `seq_bits`, once cell at a time:
void seq_pos_dec(llama_seq_id s, uint32_t i) {
const auto n = seq_pos[s].erase({ pos[i], i });
assert(n == 1);
GGML_UNUSED(n);
+
+ seq_bits[s][i/64] &= ~(1ull << (i%64));
}
void seq_pos_inc(llama_seq_id s, uint32_t i) {
seq_pos[s].insert({ pos[i], i });
+
+ if (seq_bits[s].empty()) {
+ seq_bits[s].assign((pos.size() + 63)/64, 0);
+ }
+
+ seq_bits[s][i/64] |= 1ull << (i%64);
}
// remove cell i
diff --git a/src/models/clef.cpp b/src/models/clef.cpp
index 6f8720913..e9be0ec14 100644
--- a/src/models/clef.cpp
+++ b/src/models/clef.cpp
@@ -321,7 +321,7 @@ public:
using llm_graph_input_attn_no_cache::llm_graph_input_attn_no_cache;
void set_input(const llama_ubatch * ubatch) override {
- const int64_t n_kv = self_kq_mask->ne[0];
+ const int64_t n_kv = self_attn_mask.raw->ne[0];
const int64_t n_tokens = ubatch->n_tokens;
const auto fill_mask = [&](auto * data, auto zero, auto ninf) {
@@ -333,11 +333,11 @@ public:
}
};
- GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));
- if (self_kq_mask->type == GGML_TYPE_F16) {
- fill_mask((ggml_fp16_t *) self_kq_mask->data, ggml_fp32_to_fp16(0.0f), ggml_fp32_to_fp16(-INFINITY));
+ GGML_ASSERT(ggml_backend_buffer_is_host(self_attn_mask.raw->buffer));
+ if (self_attn_mask.raw->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) self_attn_mask.raw->data, ggml_fp32_to_fp16(0.0f), ggml_fp32_to_fp16(-INFINITY));
} else {
- fill_mask((float *) self_kq_mask->data, 0.0f, -INFINITY);
+ fill_mask((float *) self_attn_mask.raw->data, 0.0f, -INFINITY);
}
}
};
@@ -348,11 +348,11 @@ llm_graph_input_attn_no_cache * llama_model_clef::graph::build_attn_inp_causal()
const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
- inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, GGML_PAD(n_tokens, 256), n_tokens, 1, 1);
- ggml_set_input(inp->self_kq_mask);
- cb(inp->self_kq_mask, "self_kq_mask", -1);
+ inp->self_attn_mask.raw = ggml_new_tensor_4d(ctx0, type_mask, GGML_PAD(n_tokens, 256), n_tokens, 1, 1);
+ ggml_set_input(inp->self_attn_mask.raw);
+ cb(inp->self_attn_mask.raw, "self_kq_mask", -1);
- inp->self_kq_mask_cnv = inp->self_kq_mask;
+ inp->self_attn_mask.cnv = inp->self_attn_mask.raw;
return (llm_graph_input_attn_no_cache *) res->add_input(std::move(inp));
}
diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp
index 6ee0d564d..22f516a48 100644
--- a/src/models/deepseek32.cpp
+++ b/src/models/deepseek32.cpp
@@ -271,7 +271,7 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
ggml_tensor * indexer_score = nullptr;
if (cparams.fused_lid) {
- indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
+ indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->self_attn_mask_lid.cnv);
cb(indexer_score, "indexer_score", il);
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
} else {
@@ -305,7 +305,7 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
cb(indexer_score, "indexer_score", il);
// mask indexer scores
- ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
+ ggml_tensor * indexer_kq_mask = inp_attn_dsa->self_attn_mask_lid.cnv;
indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
cb(indexer_score, "indexer_score", il);
}
diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp
index 4ffef08c9..85fe141fd 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -759,14 +759,14 @@ ggml_tensor * llama_model_deepseek4::graph::build_csa_lid_attention(
ggml_tensor * k_all = ggml_concat(ctx0, raw_k, csa_k, 2);
cb(k_all, "csa_k_all", il);
- ggml_tensor * raw_mask = inp_attn->get_kq_mask();
+ ggml_tensor * raw_mask = inp_attn->self_attn_mask.cnv;
ggml_tensor * csa_mask = build_top_k_mask(inp_csa.kq_mask, top_k, "csa_top_k_mask", il);
ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, csa_mask, 0);
cb(kq_mask, "csa_lid_kq_mask", il);
const int64_t n_kv_max = std::min<int64_t>(raw_mask->ne[0], hparams.n_swa) + top_k->ne[0];
- ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, n_kv_max, kq_scale, il);
+ ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, nullptr, sinks, nullptr, n_kv_max, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
@@ -815,13 +815,13 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_attention(
ggml_tensor * k_all = ggml_concat(ctx0, raw_k, hca_k, 2);
cb(k_all, "hca_k_all", il);
- ggml_tensor * raw_mask = inp_attn->get_kq_mask();
+ ggml_tensor * raw_mask = inp_attn->self_attn_mask.cnv;
ggml_tensor * hca_mask = inp_hca.kq_mask;
ggml_tensor * kq_mask = ggml_concat(ctx0, raw_mask, hca_mask, 0);
cb(kq_mask, "hca_kq_mask", il);
- ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, sinks, nullptr, 0, kq_scale, il);
+ ggml_tensor * out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, nullptr, sinks, nullptr, 0, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
@@ -853,11 +853,11 @@ ggml_tensor * llama_model_deepseek4::graph::build_raw_attention(
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il));
- ggml_tensor * kq_mask = inp_attn->get_kq_mask();
+ ggml_tensor * kq_mask = inp_attn->self_attn_mask.cnv;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
- ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, sinks, nullptr, 0, kq_scale, il);
+ ggml_tensor * out = build_attn_mha(q, k, k, nullptr, kq_mask, nullptr, sinks, nullptr, 0, kq_scale, il);
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
}
@@ -1237,7 +1237,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
ggml_tensor * inp_out_ids = build_inp_out_ids();
llm_graph_input_dsv4 * inp_dsv4 = build_inp_dsv4();
llm_graph_input_dsv4_raw * inp_attn = inp_dsv4->get_raw();
- ggml_build_forward_expand(gf, inp_attn->self_kq_mask);
+ ggml_build_forward_expand(gf, inp_attn->self_attn_mask.raw);
const int64_t hc = hparams.dsv4_hc_mult;
ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens);
diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp
index 95c8a36d5..d82998970 100644
--- a/src/models/deepseek41.cpp
+++ b/src/models/deepseek41.cpp
@@ -642,7 +642,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, kv, inp_attn->get_k_idxs(), il));
ggml_tensor * raw_k = inp_attn->mctx->get_k(ctx0, il);
- ggml_tensor * raw_mask = inp_attn->get_kq_mask();
+ ggml_tensor * raw_mask = inp_attn->self_attn_mask.cnv;
const int32_t topk_src = hparams.dsv41_topk_src[il];
const int64_t k_sel = top_k->ne[0];
@@ -713,7 +713,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
cb(kq_mask, "comp_kq_mask_sel", il);
}
- out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, layer.attn_sinks, nullptr, 0, kq_scale, il);
+ out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, nullptr, layer.attn_sinks, nullptr, 0, kq_scale, il);
} else {
// dense attention over [window | all compressed rows], selection applied as a mask
ggml_tensor * comp_k = ggml_view_4d(ctx0, comp_all, n_embd_head, 1, n_comp, n_stream,
@@ -740,7 +740,7 @@ ggml_tensor * llama_model_deepseek41::graph::build_attention41(
// at most window + top-k visible columns per query: lets CUDA FA gather only those
const int64_t n_kv_max = sparse ? std::min<int64_t>(raw_mask->ne[0], hparams.n_swa) + k_sel : 0;
- out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, layer.attn_sinks, nullptr, n_kv_max, kq_scale, il);
+ out = build_attn_mha(q, k_all, k_all, nullptr, kq_mask, nullptr, layer.attn_sinks, nullptr, n_kv_max, kq_scale, il);
}
if (k_rot) {
out = llama_mul_mat_hadamard(ctx0, out, k_rot);
@@ -781,7 +781,7 @@ llama_model_deepseek41::graph::graph(const llama_model & model_base, const llm_g
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
llm_graph_input_dsv4 * inp_dsv4 = build_inp_dsv4();
- ggml_build_forward_expand(gf, inp_dsv4->get_raw()->self_kq_mask);
+ ggml_build_forward_expand(gf, inp_dsv4->get_raw()->self_attn_mask.raw);
auto inp41 = std::make_unique<llm_graph_input_dsv41>(model, inp_dsv4->mctx);
// embedding input has no token ids: it skips Engram, like the reference's masked (image) positions
diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
index fc1ad69b0..209498421 100644
--- a/src/models/dflash.cpp
+++ b/src/models/dflash.cpp
@@ -1380,10 +1380,11 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
// * embd batch -> writes every stage's kv_norm(wkv(main_proj(target features))) into its window ring
// * token batch -> [anchor, noise...] block through the stages, single-pass mHC, no per-head q norm
// window of a block = window of its first query, see n_swa
+// TODO: move this to a new type of `llm_graph_input` and implement this in `set_input` [TAG_ATTN_MASK_CNV]
static void build_dspark_block_mask(llm_graph_context & g, llm_graph_input_attn_k_iswa * inp_attn) {
ggml_context * ctx0 = g.ctx0;
- ggml_tensor * m = inp_attn->self_kq_mask_swa_cnv;
+ ggml_tensor * m = inp_attn->self_attn_mask_swa.cnv;
const int64_t n_tokens = g.n_tokens;
const int64_t n_blocks = g.ubatch.n_seqs_unq;
const int64_t bd = n_blocks > 0 ? n_tokens / n_blocks : 0;
@@ -1391,8 +1392,8 @@ static void build_dspark_block_mask(llm_graph_context & g, llm_graph_input_attn_
ggml_tensor * m4 = ggml_reshape_4d(ctx0, m, m->ne[0], bd, m->ne[1]/bd, m->ne[3]);
ggml_tensor * r0 = ggml_view_4d(ctx0, m4, m4->ne[0], 1, m4->ne[2], m4->ne[3], m4->nb[1], m4->nb[2], m4->nb[3], 0);
m4 = ggml_repeat(ctx0, r0, m4);
- inp_attn->self_kq_mask_swa_cnv = ggml_reshape_4d(ctx0, m4, m->ne[0], m->ne[1], 1, m->ne[3]);
- g.cb(inp_attn->self_kq_mask_swa_cnv, "dspark_block_mask", -1);
+ inp_attn->self_attn_mask_swa.cnv = ggml_reshape_4d(ctx0, m4, m->ne[0], m->ne[1], 1, m->ne[3]);
+ g.cb(inp_attn->self_attn_mask_swa.cnv, "dspark_block_mask", -1);
}
}
diff --git a/src/models/dots3note.cpp b/src/models/dots3note.cpp
index 5269c8ff8..df6fee7e4 100644
--- a/src/models/dots3note.cpp
+++ b/src/models/dots3note.cpp
@@ -256,7 +256,7 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p
ggml_tensor * indexer_score = nullptr;
if (cparams.fused_lid) {
- indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn->get_dsa()->get_kq_mask_lid());
+ indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn->get_dsa()->self_attn_mask_lid.cnv);
cb(indexer_score, "indexer_score", il);
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
} else {
@@ -286,7 +286,7 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p
indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
cb(indexer_score, "indexer_score", il);
- ggml_tensor * indexer_kq_mask = inp_attn->get_dsa()->get_kq_mask_lid();
+ ggml_tensor * indexer_kq_mask = inp_attn->get_dsa()->self_attn_mask_lid.cnv;
indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
cb(indexer_score, "indexer_score", il);
}
diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp
index 9822ef5ee..0bf7a7c85 100644
--- a/src/models/glm-dsa.cpp
+++ b/src/models/glm-dsa.cpp
@@ -305,7 +305,7 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
ggml_tensor * indexer_score = nullptr;
if (cparams.fused_lid) {
- indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
+ indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->self_attn_mask_lid.cnv);
cb(indexer_score, "indexer_score", il);
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
} else {
@@ -339,7 +339,7 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
cb(indexer_score, "indexer_score", il);
// mask indexer scores
- ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
+ ggml_tensor * indexer_kq_mask = inp_attn_dsa->self_attn_mask_lid.cnv;
indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
cb(indexer_score, "indexer_score", il);
}
diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index d0917ed8a..100af6d23 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -1096,7 +1096,7 @@ llama_model_glm5_next::graph::build_dsa_layer(
cb(sel, "kq_mask_dsa", il);
- ggml_tensor * out = build_attn_mha(q_absorbed, k, v, nullptr, sel, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il);
+ ggml_tensor * out = build_attn_mha(q_absorbed, k, v, nullptr, sel, nullptr, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il);
cb(out, "kqv_out", il);
out = ggml_mul_mat(ctx0, layer.wo, out);
diff --git a/src/models/hy-v4.cpp b/src/models/hy-v4.cpp
index 3d456fe85..06833e20b 100644
--- a/src/models/hy-v4.cpp
+++ b/src/models/hy-v4.cpp
@@ -395,7 +395,7 @@ ggml_tensor * llama_model_hy_v4::graph::build_indexer_top_k(
ggml_tensor * score = nullptr;
if (cparams.fused_lid) {
- score = ggml_lightning_indexer(ctx0, iq, ik, iw, inp_attn_dsa->get_kq_mask_lid());
+ score = ggml_lightning_indexer(ctx0, iq, ik, iw, inp_attn_dsa->self_attn_mask_lid.cnv);
cb(score, "indexer_score", il);
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il});
} else {
@@ -408,7 +408,7 @@ ggml_tensor * llama_model_hy_v4::graph::build_indexer_top_k(
score = ggml_mul(ctx0, score, iw);
score = ggml_sum_rows(ctx0, score);
score = ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3));
- score = ggml_add(ctx0, score, inp_attn_dsa->get_kq_mask_lid());
+ score = ggml_add(ctx0, score, inp_attn_dsa->self_attn_mask_lid.cnv);
cb(score, "indexer_score", il);
}
diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp
index dd61f509a..3488b6c77 100644
--- a/src/models/lfm2.cpp
+++ b/src/models/lfm2.cpp
@@ -281,11 +281,11 @@ public:
}
};
- GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));
- if (self_kq_mask->type == GGML_TYPE_F16) {
- fill_mask((ggml_fp16_t *) self_kq_mask->data, ggml_fp32_to_fp16(0.0f), ggml_fp32_to_fp16(-INFINITY));
+ GGML_ASSERT(ggml_backend_buffer_is_host(self_attn_mask.raw->buffer));
+ if (self_attn_mask.raw->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) self_attn_mask.raw->data, ggml_fp32_to_fp16(0.0f), ggml_fp32_to_fp16(-INFINITY));
} else {
- fill_mask((float *) self_kq_mask->data, 0.0f, -INFINITY);
+ fill_mask((float *) self_attn_mask.raw->data, 0.0f, -INFINITY);
}
}
@@ -330,9 +330,9 @@ llama_model_lfm2::graph_decision::graph_decision(const llama_model & model, cons
llm_graph_input_attn_no_cache * inp_attn[2];
for (bool is_head : {false, true}) {
auto inp = std::make_unique<llm_graph_input_attn_media>(hparams, cparams, is_head);
- inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
- ggml_set_input(inp->self_kq_mask);
- inp->self_kq_mask_cnv = inp->self_kq_mask;
+ inp->self_attn_mask.raw = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
+ ggml_set_input(inp->self_attn_mask.raw);
+ inp->self_attn_mask.cnv = inp->self_attn_mask.raw;
inp_attn[is_head] = (llm_graph_input_attn_no_cache *) res->add_input(std::move(inp));
}
diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp
index 1b112b743..0da9d5a44 100644
--- a/src/models/minimax-m3.cpp
+++ b/src/models/minimax-m3.cpp
@@ -262,7 +262,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
if (msa_enabled) {
const auto * mctx_msa = static_cast<const llama_kv_cache_msa_context *>(mctx);
- msa_kqm = inp_attn->get_kq_mask();
+ msa_kqm = inp_attn->self_attn_mask.cnv;
n_kv = msa_kqm->ne[0];
n_tps = msa_kqm->ne[1]; // tokens per stream
ns = msa_kqm->ne[3]; // streams in this ubatch
diff --git a/src/models/mistral-large4.cpp b/src/models/mistral-large4.cpp
index dc7811d4c..dbd1907a4 100644
--- a/src/models/mistral-large4.cpp
+++ b/src/models/mistral-large4.cpp
@@ -289,7 +289,7 @@ llama_model_mistral_large4::graph::graph(const llama_model & model, const llm_gr
ggml_tensor * indexer_score = nullptr;
if (cparams.fused_lid) {
- indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
+ indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->self_attn_mask_lid.cnv);
cb(indexer_score, "indexer_score", il);
res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
} else {
@@ -323,7 +323,7 @@ llama_model_mistral_large4::graph::graph(const llama_model & model, const llm_gr
cb(indexer_score, "indexer_score", il);
// mask indexer scores
- ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
+ ggml_tensor * indexer_kq_mask = inp_attn_dsa->self_attn_mask_lid.cnv;
indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
cb(indexer_score, "indexer_score", il);
}
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index 5de67aa50..60665381c 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -1005,7 +1005,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
sel = ggml_reshape_4d(ctx0, sel, sel->ne[0], sel->ne[1]/n_stream, 1, n_stream);
cb(sel, "kq_mask_qsa", il);
- ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, sel, nullptr, nullptr, n_sel, kq_scale, il);
+ ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, sel, nullptr, nullptr, nullptr, n_sel, kq_scale, il);
cb(cur, "kqv_out", il);
// the rotation is its own inverse, so undo it on the value side of the output
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index e4ebd7c5f..eec0423cb 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -8781,6 +8781,102 @@ struct test_flash_attn_ext_large_logits : public test_flash_attn_ext {
}
};
+// GGML_OP_FLASH_ATTN_EXT with kv_rows: each q slice attends to its own list of K/V cache rows
+struct test_flash_attn_ext_kv_rows : public test_case {
+ const int64_t hs;
+ const int64_t nh; // K/V heads
+ const int64_t gqa;
+ const int64_t n_cells; // rows of the K/V cache
+ const int64_t kv; // list length per slice
+ const int64_t nb; // queries per slice
+ const int64_t n_slices;
+ const ggml_type type_KV;
+
+ std::vector<int32_t> rows;
+
+ std::string vars() override {
+ return VARS_TO_STR8(hs, nh, gqa, n_cells, kv, nb, n_slices, type_KV);
+ }
+
+ double max_nmse_err() override {
+ return 5e-4;
+ }
+
+ test_flash_attn_ext_kv_rows(int64_t hs = 256, int64_t nh = 2, int64_t gqa = 8, int64_t n_cells = 2048, int64_t kv = 512,
+ int64_t nb = 8, int64_t n_slices = 4, ggml_type type_KV = GGML_TYPE_F16)
+ : hs(hs), nh(nh), gqa(gqa), n_cells(n_cells), kv(kv), nb(nb), n_slices(n_slices), type_KV(type_KV) {}
+
+ ggml_tensor * build_graph(ggml_context * ctx) override {
+ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, hs, nb, nh*gqa, n_slices);
+ ggml_set_name(q, "q");
+
+ // cache layout [hs, nh, n_cells], shared by all slices
+ ggml_tensor * k0 = ggml_new_tensor_3d(ctx, type_KV, hs, nh, n_cells);
+ ggml_tensor * v0 = ggml_new_tensor_3d(ctx, type_KV, hs, nh, n_cells);
+ ggml_set_name(k0, "k0");
+ ggml_set_name(v0, "v0");
+
+ ggml_tensor * k = ggml_permute(ctx, k0, 0, 2, 1, 3);
+ ggml_tensor * v = ggml_permute(ctx, v0, 0, 2, 1, 3);
+
+ ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, n_slices);
+ ggml_set_name(m, "m");
+
+ ggml_tensor * r = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, kv, n_slices);
+ ggml_set_name(r, "r");
+
+ ggml_tensor * out = ggml_flash_attn_ext_rows(ctx, q, k, v, m, r, 1.0f/sqrtf(hs), 0.0f, 0.0f);
+ ggml_prec_set_acc(out, GGML_PREC_F32);
+ ggml_set_name(out, "out");
+
+ return out;
+ }
+
+ void initialize_tensors(ggml_context * ctx) override {
+ std::mt19937 gen(0x6B76);
+
+ // per slice: distinct random cells, the tail of the list is padding
+ rows.assign(kv*n_slices, -1);
+ std::vector<int32_t> order(n_cells);
+ for (int64_t s = 0; s < n_slices; ++s) {
+ for (int64_t i = 0; i < n_cells; ++i) {
+ order[i] = i;
+ }
+ std::shuffle(order.begin(), order.end(), gen);
+ const int64_t n_used = std::min<int64_t>(n_cells, kv - (s*37) % std::max<int64_t>(1, kv/2));
+ for (int64_t i = 0; i < n_used; ++i) {
+ rows[s*kv + i] = order[i];
+ }
+ }
+
+ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+ if (t->view_src) {
+ continue;
+ }
+ if (strcmp(t->name, "r") == 0) {
+ ggml_backend_tensor_set(t, rows.data(), 0, rows.size()*sizeof(int32_t));
+ } else if (strcmp(t->name, "m") == 0) {
+ std::vector<float> f(ggml_nelements(t));
+ std::uniform_real_distribution<float> dis(-1.0f, 1.0f);
+ for (int64_t s = 0; s < n_slices; ++s) {
+ for (int64_t j = 0; j < nb; ++j) {
+ for (int64_t i = 0; i < kv; ++i) {
+ const bool pad = rows[s*kv + i] < 0;
+ const bool drop = (i + 3*j + 5*s) % 11 == 0;
+ f[(s*nb + j)*kv + i] = pad || drop ? -INFINITY : dis(gen);
+ }
+ }
+ }
+ std::vector<ggml_fp16_t> h(f.size());
+ ggml_fp32_to_fp16_row(f.data(), h.data(), f.size());
+ ggml_backend_tensor_set(t, h.data(), 0, h.size()*sizeof(ggml_fp16_t));
+ } else {
+ init_tensor_uniform(t);
+ }
+ }
+ }
+};
+
// GGML_OP_CROSS_ENTROPY_LOSS
struct test_cross_entropy_loss : public test_case {
const ggml_type type;
@@ -12286,6 +12382,18 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_flash_attn_ext(128, 128, 1, { 8, 1}, 4096, 4, true, false, 8.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
+ // flash attention over a list of K/V rows per slice (unified KV cache)
+ for (int64_t nb : {1, 3, 8, 64, 512}) {
+ for (int64_t n_slices : {1, 4}) {
+ test_cases.emplace_back(new test_flash_attn_ext_kv_rows(256, 2, 8, 4096, 1024, nb, n_slices));
+ test_cases.emplace_back(new test_flash_attn_ext_kv_rows(128, 8, 4, 4096, 512, nb, n_slices));
+ test_cases.emplace_back(new test_flash_attn_ext_kv_rows(128, 4, 1, 1000, 256, nb, n_slices));
+ test_cases.emplace_back(new test_flash_attn_ext_kv_rows( 64, 4, 2, 3000, 768, nb, n_slices, GGML_TYPE_Q8_0));
+ test_cases.emplace_back(new test_flash_attn_ext_kv_rows(512, 2, 8, 4096, 1024, nb, n_slices));
+ test_cases.emplace_back(new test_flash_attn_ext_kv_rows(256, 8, 2, 4096, 1280, nb, n_slices));
+ }
+ }
+
// sparse mask with large batch size
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2048));