diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
index a95dc4dc..2a558bce 100644
--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh
+++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
@@ -480,32 +480,83 @@ static __device__ __forceinline__ void flash_attn_ext_f16_store_q4_0_block(
}
// Dequantize one q8_0 block (32 weights) given its scale and the 32 int8 bytes as 8 words, store 16 half2.
+// Each byte pair becomes the half2 1024 + (q + 128) by placing it under the exponent byte 0x64, minus 1152 that is
+// exactly q. d*q is exact in float (11 x 8 significant bits), so the correctly rounded half product equals the
+// rounding of d*float(q) that dequantize_q8_0 -> half does: the tile is bit-identical, at a fraction of the cost.
static __device__ __forceinline__ void flash_attn_ext_f16_store_q8_0_block(
- half2 * const __restrict__ dst, const float d, const uint32_t * const __restrict__ q) {
+ half2 * const __restrict__ dst, const half d, const uint32_t * const __restrict__ q) {
+ const half2 d2 = __half2half2(d);
+ const half2 bias = make_half2(1152.0f, 1152.0f);
#pragma unroll
for (int m = 0; m < 8; ++m) {
- const uint32_t v = q[m];
- dst[2*m ] = __floats2half2_rn(d*float(int(int8_t(v & 0xFF))), d*float(int(int8_t((v >> 8) & 0xFF))));
- dst[2*m + 1] = __floats2half2_rn(d*float(int(int8_t((v >> 16) & 0xFF))), d*float(int(int8_t((v >> 24) & 0xFF))));
+ const uint32_t u = q[m] ^ 0x80808080u; // int8 -> uint8 offset by 128
+ const uint32_t lo = __byte_perm(u, 0x64646464u, 0x4140);
+ const uint32_t hi = __byte_perm(u, 0x64646464u, 0x4342);
+ dst[2*m ] = __hmul2(__hsub2(*((const half2 *) &lo), bias), d2);
+ dst[2*m + 1] = __hmul2(__hsub2(*((const half2 *) &hi), bias), d2);
+ }
+}
+
+// A q4_0 / q8_0 K or V row is read as block pairs (64 weights): a q4_0 pair is 36 bytes and a q8_0 pair 68 bytes,
+// both 4-byte aligned at even block offsets, so the pair is read as aligned 32-bit words and the odd block's bytes
+// are realigned with byte permutes.
+template<ggml_type type_KV>
+struct flash_attn_ext_f16_q_pair {
+ static constexpr int bs = 32; // weights per block, both types
+ static constexpr int ts = type_KV == GGML_TYPE_Q4_0 ? 18 : 34; // bytes per block
+ static constexpr int nwords = 2*ts/4; // 32-bit words per block pair: 9 or 17
+ static_assert((2*ts) % 4 == 0, "block pair is not word aligned");
+};
+
+template<ggml_type type_KV>
+static __device__ __forceinline__ void flash_attn_ext_f16_load_q_pair(
+ const char * const __restrict__ src, uint32_t * const __restrict__ w) {
+#pragma unroll
+ for (int k = 0; k < flash_attn_ext_f16_q_pair<type_KV>::nwords; ++k) {
+ w[k] = __ldg((const uint32_t *) src + k);
+ }
+}
+
+// Dequantize a block pair from its words into 32 half2 (64 weights).
+template<ggml_type type_KV>
+static __device__ __forceinline__ void flash_attn_ext_f16_dequantize_q_pair(
+ const uint32_t * const __restrict__ w, half2 * const __restrict__ dst) {
+ static_assert(type_KV == GGML_TYPE_Q4_0 || type_KV == GGML_TYPE_Q8_0, "unsupported quantized K/V type");
+ constexpr int bs = flash_attn_ext_f16_q_pair<type_KV>::bs;
+ if constexpr (type_KV == GGML_TYPE_Q4_0) {
+ // block 0: d = bytes 0-1, qs = bytes 2-17. block 1: d = bytes 18-19, qs = bytes 20-35 (= w5..w8).
+ const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF)));
+ const float d1 = __half2float(__ushort_as_half((unsigned short) (w[4] >> 16)));
+ flash_attn_ext_f16_store_q4_0_block(dst, d0,
+ __byte_perm(w[0], w[1], 0x5432), __byte_perm(w[1], w[2], 0x5432),
+ __byte_perm(w[2], w[3], 0x5432), __byte_perm(w[3], w[4], 0x5432));
+ flash_attn_ext_f16_store_q4_0_block(dst + bs/2, d1, w[5], w[6], w[7], w[8]);
+ } else {
+ // block 0: d = bytes 0-1, qs = bytes 2-33. block 1: d = bytes 34-35, qs = bytes 36-67 (= w9..w16).
+ const half d0 = __ushort_as_half((unsigned short) (w[0] & 0xFFFF));
+ const half d1 = __ushort_as_half((unsigned short) (w[8] >> 16));
+ uint32_t q0[8];
+#pragma unroll
+ for (int k = 0; k < 8; ++k) {
+ q0[k] = __byte_perm(w[k], w[k + 1], 0x5432);
+ }
+ flash_attn_ext_f16_store_q8_0_block(dst, d0, q0);
+ flash_attn_ext_f16_store_q8_0_block(dst + bs/2, d1, w + 9);
}
}
// Loads a [nbatch_fa x ne] tile of a q4_0 / q8_0 K or V cache and dequantizes it into the F16 shared tile, in the
// same layout the F16 loader produces. KV points at byte 0 of tile row 0, e0 / ne are the element offset and count
-// within a row (multiples of 64), stride_KV is the row stride in bytes. Each thread handles two adjacent blocks
-// (64 weights): a q4_0 pair is 36 bytes and a q8_0 pair 68 bytes, both 4-byte aligned at even block offsets, so the
-// pair is read as aligned 32-bit words and the odd block's bytes are realigned with byte permutes.
+// within a row (multiples of 64), stride_KV is the row stride in bytes. Each thread handles one block pair at a time.
template<ggml_type type_KV, int stride_tile, int nwarps, int nbatch_fa, bool oob_check>
static __device__ __forceinline__ void flash_attn_ext_f16_load_tile_q(
const char * const __restrict__ KV, half2 * const __restrict__ tile_KV,
const int e0, const int ne, const int stride_KV, const int i_sup) {
- static_assert(type_KV == GGML_TYPE_Q4_0 || type_KV == GGML_TYPE_Q8_0, "unsupported quantized K/V type");
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int nthreads = nwarps*warp_size;
- constexpr int bs = 32; // weights per block, both types
- constexpr int ts = type_KV == GGML_TYPE_Q4_0 ? 18 : 34; // bytes per block
- constexpr int nwords = 2*ts/4; // 32-bit words per block pair: 9 or 17
- static_assert((2*ts) % 4 == 0, "block pair is not word aligned");
+ constexpr int bs = flash_attn_ext_f16_q_pair<type_KV>::bs;
+ constexpr int ts = flash_attn_ext_f16_q_pair<type_KV>::ts;
+ constexpr int nwords = flash_attn_ext_f16_q_pair<type_KV>::nwords;
const int pairs_per_row = ne / (2*bs);
const int ntotal = nbatch_fa * pairs_per_row;
@@ -523,32 +574,89 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile_q(
continue;
}
- const uint32_t * src = (const uint32_t *) (KV + int64_t(i)*stride_KV + ((e0 + p*2*bs)/bs)*ts);
uint32_t w[nwords];
-#pragma unroll
- for (int k = 0; k < nwords; ++k) {
- w[k] = __ldg(src + k);
+ flash_attn_ext_f16_load_q_pair<type_KV>(KV + int64_t(i)*stride_KV + ((e0 + p*2*bs)/bs)*ts, w);
+ flash_attn_ext_f16_dequantize_q_pair<type_KV>(w, dst);
+ }
+}
+
+// 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.
+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 &&
+ ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ) % 16 == 0 && ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV) % 16 == 0;
+}
+
+static __host__ bool ggml_cuda_fattn_mma_kv_prefetch(
+ const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) {
+ return cp_async_available(cc) && ggml_cuda_fattn_mma_kv_prefetch_shape(type_KV, DKQ, DV, ncols1, ncols2);
+}
+
+static constexpr __device__ bool ggml_cuda_fattn_mma_kv_prefetch(
+ const ggml_type type_KV, const int DKQ, const int DV, const int ncols1, const int ncols2) {
+#ifdef CP_ASYNC_AVAILABLE
+ return ggml_cuda_fattn_mma_kv_prefetch_shape(type_KV, DKQ, DV, ncols1, ncols2);
+#else
+ GGML_UNUSED_VARS(type_KV, DKQ, DV, ncols1, ncols2);
+ return false;
+#endif // CP_ASYNC_AVAILABLE
+}
+
+// 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(
+ const char * const __restrict__ KV, char * const __restrict__ raw, const int stride_KV) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nthreads = nwarps*warp_size;
+ constexpr int cpr = row_bytes/16; // 16-byte chunks per row
+ constexpr int nchunks = nbatch_fa*cpr;
+ static_assert(row_bytes % 16 == 0, "rows must be whole 16-byte chunks");
+
+ const unsigned int raw_32 = ggml_cuda_cvta_generic_to_shared(raw);
+ // Not unrolled: the 64-bit addresses of all chunks would spill the Q and VKQ registers at head size 256.
+#pragma unroll 1
+ for (int c0 = 0; c0 < nchunks; c0 += nthreads) {
+ const int c = c0 + threadIdx.y*warp_size + threadIdx.x;
+ if (nchunks % nthreads == 0 || c < nchunks) {
+ const int i = c / cpr;
+ const int k = c % cpr;
+ cp_async_cg_16<0>(raw_32 + i*row_bytes + 16*k, KV + int64_t(i)*stride_KV + 16*k);
}
+ }
+}
- if constexpr (type_KV == GGML_TYPE_Q4_0) {
- // block 0: d = bytes 0-1, qs = bytes 2-17. block 1: d = bytes 18-19, qs = bytes 20-35 (= w5..w8).
- const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF)));
- const float d1 = __half2float(__ushort_as_half((unsigned short) (w[4] >> 16)));
- flash_attn_ext_f16_store_q4_0_block(dst, d0,
- __byte_perm(w[0], w[1], 0x5432), __byte_perm(w[1], w[2], 0x5432),
- __byte_perm(w[2], w[3], 0x5432), __byte_perm(w[3], w[4], 0x5432));
- flash_attn_ext_f16_store_q4_0_block(dst + bs/2, d1, w[5], w[6], w[7], w[8]);
- } else {
- // block 0: d = bytes 0-1, qs = bytes 2-33. block 1: d = bytes 34-35, qs = bytes 36-67 (= w9..w16).
- const float d0 = __half2float(__ushort_as_half((unsigned short) (w[0] & 0xFFFF)));
- const float d1 = __half2float(__ushort_as_half((unsigned short) (w[8] >> 16)));
- uint32_t q0[8];
-#pragma unroll
- for (int k = 0; k < 8; ++k) {
- q0[k] = __byte_perm(w[k], w[k + 1], 0x5432);
+// Dequantizes a raw [nbatch_fa x D] tile from shared memory into the F16 tile, block pair by block pair.
+template<ggml_type type_KV, int D, int stride_tile, int nwarps, int nbatch_fa>
+static __device__ __forceinline__ void flash_attn_ext_f16_dequantize_raw_q(
+ const char * const __restrict__ raw, half2 * const __restrict__ tile_KV) {
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nthreads = nwarps*warp_size;
+ constexpr int row_bytes = ggml_cuda_fattn_mma_q_row_bytes(type_KV, D);
+ constexpr int ts = flash_attn_ext_f16_q_pair<type_KV>::ts;
+ constexpr int nwords = flash_attn_ext_f16_q_pair<type_KV>::nwords;
+ constexpr int ppr = D/64; // block pairs per row
+ constexpr int ntotal = nbatch_fa*ppr;
+
+#pragma unroll
+ for (int t0 = 0; t0 < ntotal; t0 += nthreads) {
+ const int t = t0 + threadIdx.y*warp_size + threadIdx.x;
+ if (ntotal % nthreads == 0 || t < ntotal) {
+ const int i = t / ppr;
+ const int p = t % ppr;
+ const uint32_t * src = (const uint32_t *) (raw + i*row_bytes + p*(2*ts));
+ uint32_t w[nwords];
+#pragma unroll
+ for (int k = 0; k < nwords; ++k) {
+ w[k] = src[k];
}
- flash_attn_ext_f16_store_q8_0_block(dst, d0, q0);
- flash_attn_ext_f16_store_q8_0_block(dst + bs/2, d1, w + 9);
+ flash_attn_ext_f16_dequantize_q_pair<type_KV>(w, tile_KV + i*stride_tile + p*32);
}
}
}
@@ -683,6 +791,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
static_assert(type_KV == GGML_TYPE_F16 || ((2*nbatch_K2) % 64 == 0 && (2*nbatch_V2) % 64 == 0 && !V_is_K_view),
"quantized K/V loading needs whole block pairs per pass");
+ 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 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.
+ char * raw_K = (char *) (tile_mask + ncols1*(nbatch_fa + 8));
+ char * raw_V = raw_K + nbatch_fa*row_bytes_K;
+
const int k_VKQ_0 = kb0 * nbatch_fa;
#if defined(TURING_MMA_AVAILABLE)
T_C_KQ KQ_C[nbatch_fa/(np*(cols_per_warp == 8 ? T_C_KQ::I : T_C_KQ::J))];
@@ -701,6 +818,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
__syncthreads();
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);
} else {
constexpr bool use_cp_async = nstages == 1;
if (ncols2 > 1 || mask_h) {
@@ -724,6 +848,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
if (use_cp_async) {
cp_async_wait_all();
}
+ } else if constexpr (kv_prefetch) {
+ flash_attn_ext_f16_dequantize_raw_q<type_KV, DKQ, stride_tile_K, nwarps, nbatch_fa>(raw_K, tile_K);
} else {
flash_attn_ext_f16_load_tile_q<type_KV, stride_tile_K, nwarps, nbatch_fa, oob_check>
((const char *) K_h2 + int64_t(k_VKQ_0)*stride_K, tile_K, 2*k0_start, 2*k0_diff, stride_K, k_VKQ_sup);
@@ -783,7 +909,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
}
}
- if constexpr (nstages <= 1) {
+ if constexpr (nstages <= 1 && !kv_prefetch) {
__syncthreads(); // Only needed if tile_K == tile_V.
}
}
@@ -1081,6 +1207,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
if (use_cp_async) {
cp_async_wait_all();
}
+ } 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();
+ __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);
+ 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);
+ }
+ flash_attn_ext_f16_dequantize_raw_q<type_KV, DV, stride_tile_V, nwarps, nbatch_fa>(raw_V, tile_V);
} else {
flash_attn_ext_f16_load_tile_q<type_KV, stride_tile_V, nwarps, nbatch_fa, oob_check>
((const char *) V_h2 + int64_t(k_VKQ_0)*stride_V, tile_V, i0_start, i0_diff, stride_V, k_VKQ_sup);
@@ -1131,7 +1268,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
}
#endif // defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
- if constexpr (nstages <= 1) {
+ // With prefetch the barrier at the start of the next iteration (or after the loop) protects tile_V.
+ if constexpr (nstages <= 1 && !(kv_prefetch && !last_iter)) {
__syncthreads(); // Only needed if tile_K == tile_V.
}
}
@@ -1376,6 +1514,15 @@ 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.
+ 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);
+ }
+
// Preload mask and K data for first iteration when using cp_async with multiple stages:
if constexpr (nstages > 1) {
static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline");
@@ -2051,7 +2198,11 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
const size_t nbytes_shared_mask = ncols1 * (nbatch_fa/2 + 4) * sizeof(half2);
const size_t nbytes_shared_combine = nwarps*cols_per_warp * (nbatch_combine + 4) * sizeof(half2);
- const size_t nbytes_shared_KV = nstages <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage;
+ // Raw staging tiles of the quantized prompt-tile pipeline, behind the mask tile.
+ const size_t nbytes_shared_raw = ggml_cuda_fattn_mma_kv_prefetch(type_KV, DKQ, DV, ncols1, ncols2, cc) ?
+ nbatch_fa * (ggml_cuda_fattn_mma_q_row_bytes(type_KV, DKQ) + ggml_cuda_fattn_mma_q_row_bytes(type_KV, DV)) : 0;
+
+ const size_t nbytes_shared_KV = (nstages <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage) + nbytes_shared_raw;
const size_t nbytes_shared_total = std::max(nbytes_shared_combine, Q_in_reg ?
std::max(nbytes_shared_Q, nbytes_shared_KV + nbytes_shared_mask) :
@@ -2140,9 +2291,11 @@ static inline bool ggml_cuda_fattn_mma_kv_native_supported(const ggml_tensor * d
if (!(Q->ne[0] == 128 || Q->ne[0] == 256) || V->ne[0] != Q->ne[0]) {
return false;
}
- // Rows must be 32-bit addressable for the block-pair loads.
+ // Rows must be 32-bit addressable for the block-pair loads, and whole 16-byte chunks for the cp.async copies of
+ // the prompt tiles at head size 256 (see ggml_cuda_fattn_mma_kv_prefetch).
+ const size_t align = Q->ne[0] == 256 ? 16 : 4;
for (const ggml_tensor * t : {K, V}) {
- if (t->nb[1] % 4 != 0 || t->nb[2] % 4 != 0 || t->nb[3] % 4 != 0) {
+ if (t->nb[1] % align != 0 || t->nb[2] % align != 0 || t->nb[3] % align != 0 || (uintptr_t) t->data % align != 0) {
return false;
}
}
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index 62e8e9b8..39655868 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -42,6 +42,21 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 64/ncols2, ncols2, type_KV>(ctx, dst);
}
+// Whether launch_fattn gives every output tile of ncols1 x ncols2 columns to one block, i.e. does not use stream-k
+// (mirrors its rule, with the occupancy the MMA config targets).
+static bool ggml_cuda_fattn_mma_whole_tiles(const ggml_tensor * dst, const int DKQ, const int DV, const int ncols1, const int ncols2, const int cc) {
+ if (cc >= GGML_CUDA_CC_ADA_LOVELACE || amd_wmma_available(cc)) {
+ return false;
+ }
+ const ggml_tensor * Q = dst->src[0];
+ const ggml_tensor * K = dst->src[1];
+ const int64_t gqa_ratio = Q->ne[2] / K->ne[2];
+ const int64_t ntiles = (Q->ne[1] + ncols1 - 1)/ncols1 * ((gqa_ratio + ncols2 - 1)/ncols2) * K->ne[2] * Q->ne[3];
+ const int64_t max_blocks = int64_t(ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2, cc)) * ggml_cuda_info().devices[ggml_cuda_get_device()].nsm;
+ const int64_t nwaves = (ntiles + max_blocks - 1) / max_blocks;
+ return 100 * ntiles / (max_blocks*nwaves) >= 75;
+}
+
template <int DKQ, int DV, ggml_type type_KV = GGML_TYPE_F16>
static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
@@ -97,6 +112,24 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con
}
}
+ // Prompt batches over a quantized cache: tile the Q heads of a KV head without padding columns (24 over 4 KV heads
+ // as 3 pairs, not 6 of 8), so no tensor-core work is wasted and every K/V tile is streamed fewer times. Each
+ // column's arithmetic does not depend on ncols2, so the result is the same as long as neither tiling splits a
+ // column's KV range across blocks (stream-k): small batches and decode/verify keep the 8-wide tile.
+ if constexpr (ggml_cuda_fattn_mma_kv_native(type_KV)) {
+ if (use_gqa_opt && Q->ne[1] > 16 && gqa_ratio > 4 && gqa_ratio % 8 != 0 && gqa_ratio % 2 == 0) {
+ if (gqa_ratio % 4 == 0) {
+ if (ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 16, 4, cc) && ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 8, 8, cc)) {
+ ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 4, type_KV>(ctx, dst);
+ return;
+ }
+ } else if (ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 32, 2, cc) && ggml_cuda_fattn_mma_whole_tiles(dst, DKQ, DV, 8, 8, cc)) {
+ ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 2, type_KV>(ctx, dst);
+ return;
+ }
+ }
+ }
+
if (use_gqa_opt && gqa_ratio > 4) {
ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<DKQ, DV, 8, type_KV>(ctx, dst);
return;
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index c93cacf5..57ead864 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -10171,6 +10171,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
+ // prompt-shaped cases with the same shape: the quantized cache is tiled two query heads per column group
+ for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, }) {
+ for (int nb : { 9, 33, 64, 300, 1024, }) {
+ for (int kv : { 512, 4096, }) {
+ test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV));
+ }
+ }
+ }
+
// prefill-shaped cases with long KV (nb >= 32, kv >= 1024): covers the
// XMX/GEMM-accelerated SYCL FA path which only activates for these shapes.
for (int kv : { 1024, 2048, }) {
@@ -10746,6 +10755,24 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 10000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 20000, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+ // prefill at depth with the Qwen 3.x full-attention shape (24 query heads over 4 KV heads)
+ for (int kv : { 17408, 33792, 66560, }) {
+ for (int nb : { 512, 1024, }) {
+ for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, }) {
+ test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV));
+ }
+ }
+ }
+
+ // prompt-shaped cases with the same shape: the quantized cache is tiled two query heads per column group
+ for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0, }) {
+ for (int nb : { 9, 33, 64, 300, 1024, }) {
+ for (int kv : { 512, 4096, }) {
+ test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, nb, true, false, 0, 0, GGML_PREC_F32, type_KV, type_KV));
+ }
+ }
+ }
+
for (int kv : { 4096, 8192, 16384, }) {
for (int hs : { 64, 128, }) {
for (int nr : { 1, 4, }) {