diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index b5042840..58620e74 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -10,6 +10,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
const ggml_tensor * Q = dst->src[0];
+ // batch-invariant mode: one tile width for 1 to 4 queries, so a token verified in a speculative batch
+ // attends with the same arithmetic as a token decoded alone
+ if constexpr (ncols2 <= 8) {
+ if (ggml_cuda_batch_invariant() && turing_mma_available(cc) && Q->ne[1] <= 4 && Q->ne[1]*ncols2 <= 32) {
+ ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 32/ncols2, ncols2>(ctx, dst);
+ return;
+ }
+ }
+
if constexpr (ncols2 <= 8) {
if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) {
ggml_cuda_flash_attn_ext_mma_f16_case<DKQ, DV, 8/ncols2, ncols2>(ctx, dst);
@@ -464,6 +473,11 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
// speculative batch attends with the same arithmetic as a token decoded alone (the whole-model
// guarantee is 1 to 4 queries, see ggml_cuda_batch_invariant)
if (ggml_cuda_batch_invariant() && Q->ne[1] <= 8 && Q->ne[3] == 1) {
+ // with an F16 cache the tensor-core kernel is much faster at depth (39 vs 54 ms per token at 64K on
+ // an RTX 3060); ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1 keeps it batch-invariant up to 4 queries
+ if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && gqa_opt_applies && Q->ne[1] <= 4) {
+ return BEST_FATTN_KERNEL_MMA_F16;
+ }
return BEST_FATTN_KERNEL_VEC;
}
if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) {