diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
index 2a558bce..c82647ee 100644
--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh
+++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
@@ -580,17 +580,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile_q(
}
}
-// Prompt-sized tiles over a quantized cache are software pipelined: the raw rows of the next K or V tile and the next
-// mask tile are copied into shared memory with cp.async while the tensor cores work on the current tile, then
-// dequantized shared -> shared. Needs whole 16-byte aligned rows (head size 256: 272 bytes of q8_0, 144 of q4_0)
-// and no OOB rows (ncols2 > 1 implies a padded cache); the staging buffers fit in the Q tile's shared memory.
+// Tiles of 32 or more columns over a quantized cache (prompt tiles and the 4x8 decode/verify tile) are software
+// pipelined: the raw rows of the next K or V tile and the next mask tile are copied into shared memory with cp.async
+// while the tensor cores work on the current tile, then dequantized shared -> shared. Needs whole 16-byte aligned
+// rows (head size 256: 272 bytes of q8_0, 144 of q4_0) and no OOB rows (ncols2 > 1 implies a padded cache). The
+// staging buffers keep the 2 blocks per SM of these tiles.
static constexpr __host__ __device__ int ggml_cuda_fattn_mma_q_row_bytes(const ggml_type type_KV, const int D) {
return D/32 * (type_KV == GGML_TYPE_Q4_0 ? 18 : 34);
}
static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_kv_prefetch_shape(
const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) {
- return ggml_cuda_fattn_mma_kv_native(type_KV) && ncols1*ncols2 >= 64 && ncols2 > 1 &&
+ return ggml_cuda_fattn_mma_kv_native(type_KV) && ncols1*ncols2 >= 32 && ncols2 > 1 &&
ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ) % 16 == 0 && ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV) % 16 == 0;
}
@@ -609,6 +610,14 @@ static constexpr __device__ bool ggml_cuda_fattn_mma_kv_prefetch(
#endif // CP_ASYNC_AVAILABLE
}
+// Decode/verify tiles (32 columns) have little math to hide a load behind, so they start the next tile's raw K right
+// after its staging buffer is consumed and raw V one phase early (same buffers, same barriers). Prompt tiles keep the
+// later start: on them the longer-lived addresses cost registers and spills.
+static constexpr __host__ __device__ bool ggml_cuda_fattn_mma_kv_lead(
+ const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) {
+ return ggml_cuda_fattn_mma_kv_prefetch_shape(type_KV, DKQ, DV, ncols1, ncols2) && ncols1*ncols2 <= 32;
+}
+
// Starts copying the raw rows of a [nbatch_fa x row_bytes] quantized K/V tile into shared memory, row after row.
template<int row_bytes, int nwarps, int nbatch_fa>
static __device__ __forceinline__ void flash_attn_ext_f16_load_raw_q(
@@ -794,6 +803,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
constexpr bool kv_prefetch = ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2);
static_assert(!kv_prefetch || (nstages == 0 && nbatch_K2 == DKQ/2 && nbatch_V2 == DV/2 && !oob_check),
"K/V prefetch needs whole rows per pass and no OOB rows");
+ constexpr bool kv_lead = ggml_cuda_fattn_mma_kv_lead(type_KV, DKQ, DV, ncols1, ncols2);
constexpr int row_bytes_K = ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ);
constexpr int row_bytes_V = ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV);
// Raw staging tiles behind the mask tile, see flash_attn_ext_f16_process_tile.
@@ -819,12 +829,18 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check>
(V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, nbatch_V2, stride_V, k_VKQ_sup);
} else if constexpr (kv_prefetch) {
- // Raw K and the mask of this tile were started by the previous iteration (or before the first one);
- // the barrier also retires every warp's VKQ reads of tile_V before tile_K (the same memory) is written.
- cp_async_wait_all();
- __syncthreads();
- flash_attn_ext_f16_load_raw_q<row_bytes_V, nwarps, nbatch_fa>
- ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, raw_V, stride_V);
+ // Raw K and the mask of this tile were started by the previous iteration (or before the first one). With kv_lead
+ // raw V was started after them, so all but the newest cp.async group must be done, otherwise raw V starts here.
+ // The barrier also retires every warp's VKQ reads of tile_V before tile_K (the same memory) is written.
+ if constexpr (kv_lead) {
+ cp_async_wait_group<1>();
+ __syncthreads();
+ } else {
+ cp_async_wait_all();
+ __syncthreads();
+ flash_attn_ext_f16_load_raw_q<row_bytes_V, nwarps, nbatch_fa>
+ ((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, raw_V, stride_V);
+ }
} else {
constexpr bool use_cp_async = nstages == 1;
if (ncols2 > 1 || mask_h) {
@@ -855,6 +871,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
((const char *) K_h2 + int64_t(k_VKQ_0)*stride_K, tile_K, 2*k0_start, 2*k0_diff, stride_K, k_VKQ_sup);
}
__syncthreads();
+ if constexpr (kv_lead && !last_iter) {
+ // raw_K is free again: start the next tile's K now, so that it has the whole iteration to arrive.
+ flash_attn_ext_f16_load_raw_q<row_bytes_K, nwarps, nbatch_fa>
+ ((const char *) K_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_K, raw_K, stride_K);
+ cp_async_commit_group();
+ }
}
// Calculate tile of KQ:
@@ -1207,6 +1229,21 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
if (use_cp_async) {
cp_async_wait_all();
}
+ } else if constexpr (kv_lead) {
+ // Raw V of this tile is the oldest pending group (the next raw K may still be in flight). Every
+ // warp is past KQ and the softmax, so the mask tile can take the next tile.
+ if constexpr (last_iter) {
+ cp_async_wait_group<0>();
+ } else {
+ cp_async_wait_group<1>();
+ }
+ __syncthreads();
+ if constexpr (!last_iter) {
+ flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, true, oob_check>
+ (mask_h + k_VKQ_0 + nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01);
+ cp_async_commit_group();
+ }
+ flash_attn_ext_f16_dequantize_raw_q<type_KV, DV, stride_tile_V, nwarps, nbatch_fa>(raw_V, tile_V);
} else if constexpr (kv_prefetch) {
// Every warp is past KQ and the softmax, so raw K and the mask tile can take the next tile.
cp_async_wait_all();
@@ -1223,6 +1260,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, i0_start, i0_diff, stride_V, k_VKQ_sup);
}
__syncthreads();
+ if constexpr (kv_lead && !last_iter) {
+ // raw_V is free again: start the next tile's V.
+ flash_attn_ext_f16_load_raw_q<row_bytes_V, nwarps, nbatch_fa>
+ ((const char *) V_h2 + int64_t(k_VKQ_0 + nbatch_fa)*stride_V, raw_V, stride_V);
+ cp_async_commit_group();
+ }
}
}
const half2 * tile_V_i = !V_is_K_view || i0_stop > 2*nbatch_K2 ? tile_V : tile_V + i0_start/2;
@@ -1514,13 +1557,20 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
int kb0 = kb0_start;
- // Quantized prompt tiles: start the raw K and mask tiles of the first iteration, see flash_attn_ext_f16_iter.
+ // Quantized K/V tiles: start the raw K and mask tiles of the first iteration (with kv_lead also raw V, as a second
+ // cp.async group), see flash_attn_ext_f16_iter.
if constexpr (ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2)) {
char * raw_K = (char *) (tile_mask + ncols1*(nbatch_fa + 8));
flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, true, false>
(mask_h + kb0*nbatch_fa, tile_mask, stride_mask, nbatch_fa, jt*ncols1, ne01);
flash_attn_ext_f16_load_raw_q<ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ), nwarps, nbatch_fa>
((const char *) K_h2 + int64_t(kb0)*nbatch_fa*stride_K, raw_K, stride_K);
+ if constexpr (ggml_cuda_fattn_mma_kv_lead(type_KV, DKQ, DV, ncols1, ncols2)) {
+ cp_async_commit_group();
+ flash_attn_ext_f16_load_raw_q<ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV), nwarps, nbatch_fa>
+ ((const char *) V_h2 + int64_t(kb0)*nbatch_fa*stride_V, raw_K + nbatch_fa*ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ), stride_V);
+ cp_async_commit_group();
+ }
}
// Preload mask and K data for first iteration when using cp_async with multiple stages: