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(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(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)) {