Commit 9d7b2d739 for llama.cpp
commit 9d7b2d7397155e83934b71c01224a69516b19a2c
Author: Sihan Yu <116848194+sihanyu03@users.noreply.github.com>
Date: Wed Nov 11 08:45:01 2026 +0000
cuda : fattn mma head sizes 40 72 (#30285)
* cuda : add MMA FlashAttention support for head sizes 40 and 72
Head sizes 40 and 72 (used e.g. in vision towers of Gemma 3/4 and Qwen3-VL) were excluded from every tensor-core dispatch branch and fell back to the tile kernel.
The fp16 MMA fragment covers 16 halfs in the head direction, so the in-kernel geometry is padded to the next multiple of 16 (40 -> 48, 72 -> 80). Register arrays, shared memory tile strides and MMA loop bounds use the padded size, while global memory loads and stores the real one.
Enabled on NVIDIA Turing and newer while Volta and AMD keep the tile kernel. The RDNA/CDNA config rows and the mma_tile_sizes specialisations are needed because the explicit instantiations are compiled for every target of a translation unit.
* tests : add the head size = 40 encoder perf shape
* cuda : address review comments on the head size 40 and 72 kernels
diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
index 96b596ea6..2ce2fb261 100644
--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh
+++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
@@ -36,11 +36,21 @@ struct fattn_mma_config {
} \
static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config_ampere(const int DKQ, const int DV, const int ncols) {
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 8, 128, 2, 128, 24, 24, 24, 2, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 16, 128, 2, 64, 24, 24, 24, 2, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 32, 128, 2, 64, 24, 24, 24, 2, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 64, 128, 2, 64, 24, 24, 24, 2, true);
+
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 8, 128, 2, 128, 32, 32, 32, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 16, 128, 2, 64, 32, 32, 32, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 32, 128, 2, 64, 32, 32, 32, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 64, 128, 2, 64, 32, 32, 32, 2, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 8, 128, 2, 128, 40, 40, 40, 2, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 16, 128, 2, 64, 40, 40, 40, 2, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 32, 128, 2, 64, 40, 40, 40, 2, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 64, 128, 2, 64, 40, 40, 40, 2, true);
+
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 8, 128, 2, 128, 40, 40, 40, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 16, 128, 2, 64, 40, 40, 40, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 32, 128, 2, 64, 40, 40, 40, 2, true);
@@ -128,11 +138,21 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
}
static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config_rdna(const int DKQ, const int DV, const int ncols) {
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 8, 128, 2, 64, 24, 24, 24, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 16, 128, 2, 64, 24, 24, 24, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 32, 128, 2, 64, 24, 24, 24, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 64, 128, 2, 64, 24, 24, 24, 1, true);
+
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 8, 128, 2, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 16, 128, 2, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 32, 128, 2, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 64, 128, 2, 64, 32, 32, 32, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 8, 64, 2, 32, 40, 40, 40, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 16, 64, 2, 32, 40, 40, 40, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 32, 128, 2, 64, 40, 40, 40, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 64, 128, 2, 64, 40, 40, 40, 1, true);
+
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 8, 64, 2, 32, 40, 40, 40, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 16, 64, 2, 32, 40, 40, 40, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 32, 128, 2, 64, 40, 40, 40, 1, true);
@@ -180,11 +200,21 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
}
static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_config_cdna(const int DKQ, const int DV, const int ncols) {
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 8, 256, 2, 64, 24, 24, 24, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 16, 256, 2, 64, 24, 24, 24, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 32, 256, 2, 64, 24, 24, 24, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 40, 40, 64, 256, 2, 64, 24, 24, 24, 1, true);
+
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 8, 128, 1, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 16, 256, 2, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 32, 256, 2, 64, 32, 32, 32, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 64, 256, 3, 64, 32, 32, 32, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 8, 256, 2, 64, 40, 40, 40, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 16, 256, 2, 64, 40, 40, 40, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 32, 256, 2, 64, 40, 40, 40, 1, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 72, 72, 64, 256, 2, 64, 40, 40, 40, 1, true);
+
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 8, 256, 2, 64, 40, 40, 40, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 16, 256, 2, 64, 40, 40, 40, 1, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE( 80, 80, 32, 256, 2, 64, 40, 40, 40, 1, true);
@@ -619,6 +649,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols);
constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse);
+ constexpr int DKQp = GGML_PAD(DKQ, 16);
+ constexpr int DVp = GGML_PAD(DV, 16);
+
constexpr bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols);
constexpr int stride_tile_K = swz ? nbatch_K2 : nbatch_K2 + 4;
constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : (swz ? nbatch_V2 : nbatch_V2 + 4);
@@ -635,12 +668,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
if constexpr (nstages > 1) {
static_assert(!oob_check, "OOB check incompatible with multi-stage pipeline");
static_assert(!V_is_K_view, "K data reuse not implemented multi-stage loading");
- static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi stage loading");
+ static_assert(nbatch_K2 == DKQp/2, "batching not implemented for multi stage loading");
constexpr bool use_cp_async = true;
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, indices);
+ (V_h2, tile_V, min(nbatch_V2, DV/2), 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;
@@ -653,11 +686,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
// For MLA K and V have the same data.
// Therefore, iterate over K in reverse and later re-use the data if possible.
#pragma unroll
- for (int k0_start = (DKQ/2-1) - (DKQ/2-1) % nbatch_K2; k0_start >= 0; k0_start -= nbatch_K2) {
- const int k0_stop = k0_start + nbatch_K2 < DKQ/2 ? k0_start + nbatch_K2 : DKQ/2;
+ for (int k0_start = (DKQp/2-1) - (DKQp/2-1) % nbatch_K2; k0_start >= 0; k0_start -= nbatch_K2) {
+ const int k0_stop = k0_start + nbatch_K2 < DKQp/2 ? k0_start + nbatch_K2 : DKQp/2;
if constexpr (nstages <= 1) {
- const int k0_diff = k0_stop - k0_start;
+ const int k0_diff = min(k0_stop, DKQ/2) - k0_start; // Only load the real head size from global memory.
constexpr bool use_cp_async = nstages == 1;
flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
(K_h2 + k0_start, tile_K, k0_diff, stride_K, k_VKQ_0, k_VKQ_sup, indices);
@@ -691,7 +724,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
}
}
} else {
- constexpr int stride_tile_Q = DKQ/2 + 4;
+ constexpr int stride_tile_Q = DKQp/2 + 4;
#pragma unroll
for (int k_KQ_0 = k0_start; k_KQ_0 < k0_stop; k_KQ_0 += T_A_KQ::J) {
load_ldmatrix(Q_B[0], tile_Q + (threadIdx.y / np)*(T_B_KQ::I*stride_tile_Q) + k_KQ_0, stride_tile_Q);
@@ -918,7 +951,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
if constexpr (cols_per_warp == 8) {
const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[cols_per_thread - 1]);
#pragma unroll
- for (int i = 0; i < DV/T_C_VKQ::I; ++i) {
+ for (int i = 0; i < DVp/T_C_VKQ::I; ++i) {
#pragma unroll
for (int l = 0; l < T_C_VKQ::ne; ++l) {
VKQ_C[i].x[l] *= KQ_max_scale_h2;
@@ -929,7 +962,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
for (int col = 0; col < cols_per_thread; ++col) {
const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[col], KQ_max_scale[col]);
#pragma unroll
- for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
+ for (int i = 0; i < (DVp/2)/T_C_VKQ::J; ++i) {
#pragma unroll
for (int l0 = 0; l0 < T_C_VKQ::ne; l0 += 2) {
VKQ_C[i].x[l0 + col] *= KQ_max_scale_h2;
@@ -941,7 +974,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[0]);
#pragma unroll
- for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
+ for (int i = 0; i < (DVp/2)/T_C_VKQ::J; ++i) {
#pragma unroll
for (int l = 0; l < T_C_VKQ::ne; ++l) {
VKQ_C[i].x[l] *= KQ_max_scale_h2;
@@ -950,7 +983,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
} else {
static_assert(std::is_same_v<decltype(T_C_VKQ::x), float[T_C_VKQ::ne]>, "bad VKQ type");
#pragma unroll
- for (int i = 0; i < DV/T_C_VKQ::J; ++i) {
+ for (int i = 0; i < DVp/T_C_VKQ::J; ++i) {
#pragma unroll
for (int l = 0; l < T_C_VKQ::ne; ++l) {
VKQ_C[i].x[l] *= KQ_max_scale[0];
@@ -961,7 +994,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
const half2 KQ_max_scale_h2 = make_half2(
KQ_max_scale[(threadIdx.x / 2) % 2], KQ_max_scale[(threadIdx.x / 2) % 2]);
#pragma unroll
- for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
+ for (int i = 0; i < (DVp/2)/T_C_VKQ::J; ++i) {
#pragma unroll
for (int l = 0; l < T_C_VKQ::ne; ++l) {
VKQ_C[i].x[l] *= KQ_max_scale_h2;
@@ -997,19 +1030,19 @@ 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, indices);
+ (K_h2, tile_K, min(nbatch_K2, DKQ/2), stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, indices);
}
}
// Calculate VKQ tile, need to use logical rather than physical elements for i0 due to transposition of V:
#pragma unroll
- for (int i0_start = 0; i0_start < DV; i0_start += 2*nbatch_V2) {
- static_assert(DV % (2*nbatch_V2) == 0, "bad loop size");
+ for (int i0_start = 0; i0_start < DVp; i0_start += 2*nbatch_V2) {
+ static_assert(DVp % (2*nbatch_V2) == 0, "bad loop size");
const int i0_stop = i0_start + 2*nbatch_V2;
if constexpr (nstages <= 1) {
- const int i0_diff = i0_stop - i0_start;
+ const int i0_diff = min(i0_stop, DV) - i0_start; // Only load the real head size from global memory.
if (!V_is_K_view || i0_stop > 2*nbatch_K2) {
constexpr bool use_cp_async = nstages == 1;
flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
@@ -1146,6 +1179,22 @@ template<int ncols> struct mma_tile_sizes<112, ncols> {
using T_B_VKQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major
};
+template<int ncols> struct mma_tile_sizes<40, ncols> {
+ using T_A_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+ using T_B_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+ using T_C_KQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major
+ using T_A_VKQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+ using T_B_VKQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+ using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major
+};
+template<int ncols> struct mma_tile_sizes<72, ncols> {
+ using T_A_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+ using T_B_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+ using T_C_KQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major
+ using T_A_VKQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+ using T_B_VKQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+ using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major
+};
#else
template<int DKQ, int ncols> struct mma_tile_sizes {
using T_A_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR>; // row-major
@@ -1171,6 +1220,22 @@ template<int ncols> struct mma_tile_sizes<112, ncols> {
using T_B_VKQ = tile<16, 8, half2>; // column-major
using T_C_VKQ = tile<16, 8, half2>; // column-major
};
+template<int ncols> struct mma_tile_sizes<40, ncols> {
+ using T_A_KQ = tile<16, 8, half2>; // row-major
+ using T_B_KQ = tile<16, 8, half2>; // column-major
+ using T_C_KQ = tile<16, 16, float>; // column-major
+ using T_A_VKQ = tile<16, 8, half2>; // row-major
+ using T_B_VKQ = tile<16, 8, half2>; // column-major
+ using T_C_VKQ = tile<16, 8, half2>; // column-major
+};
+template<int ncols> struct mma_tile_sizes<72, ncols> {
+ using T_A_KQ = tile<16, 8, half2>; // row-major
+ using T_B_KQ = tile<16, 8, half2>; // column-major
+ using T_C_KQ = tile<16, 16, float>; // column-major
+ using T_A_VKQ = tile<16, 8, half2>; // row-major
+ using T_B_VKQ = tile<16, 8, half2>; // column-major
+ using T_C_VKQ = tile<16, 8, half2>; // column-major
+};
#endif // RDNA3
#elif defined(AMD_MFMA_AVAILABLE)
template<int DKQ, int ncols> struct mma_tile_sizes {
@@ -1240,6 +1305,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols);
constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, use_sparse);
+ constexpr int DKQp = GGML_PAD(DKQ, 16);
+ constexpr int DVp = GGML_PAD(DV, 16);
+
if (cols_per_warp > ncols) {
NO_DEVICE_CODE;
return;
@@ -1248,7 +1316,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
static_assert(nwarps * (cols_per_warp/ncols2) % ncols1 == 0, "bad nwarps");
constexpr bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols);
- constexpr int stride_tile_Q = DKQ/2 + 4;
+ constexpr int stride_tile_Q = DKQp/2 + 4;
constexpr int stride_tile_K = swz ? nbatch_K2 : nbatch_K2 + 4;
constexpr int stride_tile_V = V_is_K_view ? stride_tile_K : (swz ? nbatch_V2 : nbatch_V2 + 4);
constexpr int stride_tile_KV_max = stride_tile_K > stride_tile_V ? stride_tile_K : stride_tile_V;
@@ -1258,17 +1326,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
half2 * tile_V = nstages > 1 ? tile_K + nbatch_fa * stride_tile_K : tile_K;
half * tile_mask = (half *) (nstages > 1 ? tile_V + nbatch_fa * stride_tile_V : tile_V + nbatch_fa * stride_tile_KV_max);
- T_B_KQ Q_B[(Q_in_reg ? DKQ/(2*T_B_KQ::J) : 1)];
+ T_B_KQ Q_B[(Q_in_reg ? DKQp/(2*T_B_KQ::J) : 1)];
#if defined(TURING_MMA_AVAILABLE)
- T_C_VKQ VKQ_C[cols_per_warp == 8 ? DV/T_C_VKQ::I : DV/(2*T_C_VKQ::J)];
+ T_C_VKQ VKQ_C[cols_per_warp == 8 ? DVp/T_C_VKQ::I : DVp/(2*T_C_VKQ::J)];
#elif defined(AMD_WMMA_256B_AVAILABLE)
- T_C_VKQ VKQ_C[DV % 32 != 0 ? DV/T_C_VKQ::J : DV/(2*T_C_VKQ::J)];
+ T_C_VKQ VKQ_C[DVp % 32 != 0 ? DVp/T_C_VKQ::J : DVp/(2*T_C_VKQ::J)];
#elif defined(AMD_MFMA_AVAILABLE)
- T_C_VKQ VKQ_C[ DV/T_C_VKQ::J];
+ T_C_VKQ VKQ_C[ DVp/T_C_VKQ::J];
#elif defined(AMD_WMMA_AVAILABLE)
- T_C_VKQ VKQ_C[ DV/(2*T_C_VKQ::J)];
+ T_C_VKQ VKQ_C[ DVp/(2*T_C_VKQ::J)];
#else // Volta
- T_C_VKQ VKQ_C[ DV/(2*T_C_VKQ::J)];
+ T_C_VKQ VKQ_C[ DVp/(2*T_C_VKQ::J)];
#endif // defined(TURING_MMA_AVAILABLE)
float KQ_rowsum[cols_per_thread] = {0.0f};
@@ -1284,8 +1352,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
const half2 scale_h2 = make_half2(scale, scale);
#pragma unroll
for (int stride_k : {warp_size, warp_size/2, warp_size/4, warp_size/8}) {
- const int k0_start = stride_k == warp_size ? 0 : DKQ/2 - (DKQ/2) % (2*stride_k);
- const int k0_stop = DKQ/2 - (DKQ/2) % (1*stride_k);
+ const int k0_start = stride_k == warp_size ? 0 : DKQp/2 - (DKQp/2) % (2*stride_k);
+ const int k0_stop = DKQp/2 - (DKQp/2) % (1*stride_k);
const int stride_jc = warp_size / stride_k;
if (k0_start == k0_stop) {
@@ -1308,7 +1376,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
- const float2 tmp = Q_f2[(jt*ncols1 + j)*stride_Q1 + c*stride_Q2 + k];
+ const float2 tmp = DKQp == DKQ || k < DKQ/2 ?
+ Q_f2[(jt*ncols1 + j)*stride_Q1 + c*stride_Q2 + k] : make_float2(0.0f, 0.0f);
tile_Q[jc*stride_tile_Q + k] = scale_h2 * make_half2(tmp.x, tmp.y);
}
} else {
@@ -1328,19 +1397,50 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
const int j0 = (threadIdx.y / np) * cols_per_warp;
#pragma unroll
- for (int k0 = 0; k0 < DKQ/2; k0 += T_B_KQ::J) {
+ for (int k0 = 0; k0 < DKQp/2; k0 += T_B_KQ::J) {
load_ldmatrix(Q_B[k0/T_B_KQ::J], tile_Q + j0*stride_tile_Q + k0, stride_tile_Q);
}
}
__syncthreads();
+ // Zero the pad columns of the K/V tiles once, the tile loads only write the first DKQ/2 (DV/2) half2 of a row.
+ // Done after Q_B is loaded because tile_K may alias tile_Q.
+ if constexpr (DKQp != DKQ || DVp != DV) {
+ static_assert(!swz, "pad zero-fill assumes unswizzled tiles");
+ static_assert(nbatch_K2 == DKQp/2 && nbatch_V2 == DVp/2, "K/V tile batching not implemented for padded head sizes");
+ static_assert(DKQ == DV, "pad zero-fill for DKQ != DV requires separate handling of the aliased V tile");
+
+ const half2 zero[4] = {{0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}, {0.0f, 0.0f}};
+ const int tid = threadIdx.y*warp_size + threadIdx.x;
+#pragma unroll
+ for (int i0 = 0; i0 < nbatch_fa; i0 += nwarps*warp_size) {
+ const int i = i0 + tid;
+
+ if (i0 + nwarps*warp_size > nbatch_fa && i >= nbatch_fa) {
+ break;
+ }
+
+#pragma unroll
+ for (int k = DKQ/2; k < DKQp/2; k += 4) {
+ ggml_cuda_memcpy_1<16>(tile_K + i*stride_tile_K + k, zero);
+ }
+ if constexpr (nstages > 1) { // tile_V is only a separate buffer with multi-stage loading.
+#pragma unroll
+ for (int k = DV/2; k < DVp/2; k += 4) {
+ ggml_cuda_memcpy_1<16>(tile_V + i*stride_tile_V + k, zero);
+ }
+ }
+ }
+ __syncthreads();
+ }
+
int kb0 = kb0_start;
// Preload mask and K data for first iteration when using cp_async with multiple stages:
if constexpr (nstages > 1) {
static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading");
- static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline");
+ static_assert(nbatch_K2 == DKQp/2, "batching not implemented for multi-stage pipeline");
constexpr bool use_cp_async = true;
constexpr bool oob_check = false;
constexpr int k_VKQ_sup = nbatch_fa;
@@ -1349,7 +1449,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, indices);
+ (K_h2, tile_K, min(nbatch_K2, DKQ/2), stride_K, kb0*nbatch_fa, k_VKQ_sup, indices);
}
// kb0_start is always < kb0_stop so the last iter can be executed unconditionally.
@@ -1454,7 +1554,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
if constexpr (cols_per_warp == 8) {
const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[cols_per_thread - 1]);
#pragma unroll
- for (int i = 0; i < DV/T_C_VKQ::I; ++i) {
+ for (int i = 0; i < DVp/T_C_VKQ::I; ++i) {
#pragma unroll
for (int l = 0; l < T_C_VKQ::ne; ++l) {
VKQ_C[i].x[l] *= KQ_max_scale_h2;
@@ -1465,7 +1565,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
for (int col = 0; col < cols_per_thread; ++col) {
const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[col], KQ_max_scale[col]);
#pragma unroll
- for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
+ for (int i = 0; i < (DVp/2)/T_C_VKQ::J; ++i) {
#pragma unroll
for (int l0 = 0; l0 < T_C_VKQ::ne; l0 += 2) {
VKQ_C[i].x[l0 + col] *= KQ_max_scale_h2;
@@ -1477,7 +1577,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[0]);
#pragma unroll
- for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
+ for (int i = 0; i < (DVp/2)/T_C_VKQ::J; ++i) {
#pragma unroll
for (int l = 0; l < T_C_VKQ::ne; ++l) {
VKQ_C[i].x[l] *= KQ_max_scale_h2;
@@ -1486,7 +1586,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
} else {
static_assert(std::is_same_v<decltype(T_C_VKQ::x), float[T_C_VKQ::ne]>, "bad VKQ type");
#pragma unroll
- for (int i = 0; i < DV/T_C_VKQ::J; ++i) {
+ for (int i = 0; i < DVp/T_C_VKQ::J; ++i) {
#pragma unroll
for (int l = 0; l < T_C_VKQ::ne; ++l) {
VKQ_C[i].x[l] *= KQ_max_scale[0];
@@ -1497,7 +1597,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
const int col = (threadIdx.x / 2) % 2;
const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[col], KQ_max_scale[col]);
#pragma unroll
- for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
+ for (int i = 0; i < (DVp/2)/T_C_VKQ::J; ++i) {
#pragma unroll
for (int l = 0; l < T_C_VKQ::ne; ++l) {
VKQ_C[i].x[l] *= KQ_max_scale_h2;
@@ -1511,7 +1611,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
// So also write VKQ accumulators to shared memory in column-major format if np == 1.
constexpr int tile_stride = nbatch_combine + 4;
- static_assert((DV/2) % nbatch_combine == 0, "bad nbatch_combine");
+ static_assert((DVp/2) % nbatch_combine == 0, "bad nbatch_combine");
if constexpr (cols_per_warp == 8) {
const int jc_cwmo = (threadIdx.x % (2*T_C_VKQ::J)) / T_C_VKQ::J; // jc combine write meta offset
@@ -1655,7 +1755,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
}
#pragma unroll
- for (int k00 = 0; k00 < DV/2; k00 += nbatch_combine) {
+ for (int k00 = 0; k00 < DVp/2; k00 += nbatch_combine) {
if constexpr (cols_per_warp == 8) {
static_assert(std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>, "bad VKQ type");
const int jc_cwd = threadIdx.y*T_B_KQ::I + T_B_KQ::get_i(-1); // jc combine write data
@@ -1754,6 +1854,10 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
for (int k0 = k0_start; k0 < k0_stop; k0 += stride_k) {
const int k = k0 + (stride_k == warp_size ? threadIdx.x : threadIdx.x % stride_k);
+ if (DVp != DV && k00 + k >= DV/2) {
+ continue;
+ }
+
float2 dstk_val = make_float2(0.0f, 0.0f);
#pragma unroll
for (int ip = 0; ip < np; ++ip) {
@@ -2034,6 +2138,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
const int nwarps = nthreads / warp_size_host;
constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu
+ constexpr int DKQp = GGML_PAD(DKQ, 16);
// KV tile strides must match flash_attn_ext_f16_iter / _process_tile.
const bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols, cc);
@@ -2041,7 +2146,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
const int stride_tile_V = V_is_K_view ? stride_tile_K : (swz ? nbatch_V2 : nbatch_V2 + 4);
const size_t nbytes_shared_KV_1stage = nbatch_fa * std::max(stride_tile_K, stride_tile_V) * sizeof(half2);
const size_t nbytes_shared_KV_2stage = nbatch_fa * (stride_tile_K + stride_tile_V) * sizeof(half2);
- const size_t nbytes_shared_Q = ncols * (DKQ/2 + 4) * sizeof(half2);
+ const size_t nbytes_shared_Q = ncols * (DKQp/2 + 4) * sizeof(half2);
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);
@@ -2126,28 +2231,36 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/ 8, 8); \
extern DECL_FATTN_MMA_F16_CASE(DKQ, DV, (ncols)/16, 16); \
+DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 40, 40, 8)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 64, 64, 8)
+DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 72, 72, 8)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 80, 80, 8)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 96, 96, 8)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112, 8)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128, 8)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256, 8)
+DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 40, 40, 16)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 64, 64, 16)
+DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 72, 72, 16)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 80, 80, 16)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 96, 96, 16)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112, 16)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128, 16)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256, 16)
+DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 40, 40, 32)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 64, 64, 32)
+DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 72, 72, 32)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 80, 80, 32)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 96, 96, 32)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112, 32)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128, 32)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256, 32)
+DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 40, 40, 64)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 64, 64, 64)
+DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 72, 72, 64)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 80, 80, 64)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2( 96, 96, 64)
DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112, 64)
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index 49dfef14c..cf44ab4f5 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -314,10 +314,18 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg
const ggml_tensor * mask = dst->src[3];
switch (Q->ne[0]) {
+ case 40:
+ GGML_ASSERT(V->ne[0] == 40);
+ ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2< 40, 40>(ctx, dst);
+ break;
case 64:
GGML_ASSERT(V->ne[0] == 64);
ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2< 64, 64>(ctx, dst);
break;
+ case 72:
+ GGML_ASSERT(V->ne[0] == 72);
+ ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2< 72, 72>(ctx, dst);
+ break;
case 80:
GGML_ASSERT(V->ne[0] == 80);
ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2< 80, 80>(ctx, dst);
@@ -683,7 +691,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
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;
// If Turing tensor cores are available, use them:
- if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) {
+ if (turing_mma_available(cc)) {
if (can_use_vector_kernel) {
if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) {
// the sparse gather exists only in the MMA kernel: (DKQ, DV, 1, 8) with GQA > 4
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu
index 6ae77bec8..3c861c2c0 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_1-ncols2_8.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 1, 8);
DECL_FATTN_MMA_F16_CASE(64, 64, 1, 8);
+DECL_FATTN_MMA_F16_CASE(72, 72, 1, 8);
DECL_FATTN_MMA_F16_CASE(80, 80, 1, 8);
DECL_FATTN_MMA_F16_CASE(96, 96, 1, 8);
DECL_FATTN_MMA_F16_CASE(112, 112, 1, 8);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu
index 9d3cfd8ed..865015364 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_1.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 16, 1);
DECL_FATTN_MMA_F16_CASE(64, 64, 16, 1);
+DECL_FATTN_MMA_F16_CASE(72, 72, 16, 1);
DECL_FATTN_MMA_F16_CASE(80, 80, 16, 1);
DECL_FATTN_MMA_F16_CASE(96, 96, 16, 1);
DECL_FATTN_MMA_F16_CASE(112, 112, 16, 1);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu
index 0c1a37fd0..8cf304825 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_2.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 16, 2);
DECL_FATTN_MMA_F16_CASE(64, 64, 16, 2);
+DECL_FATTN_MMA_F16_CASE(72, 72, 16, 2);
DECL_FATTN_MMA_F16_CASE(80, 80, 16, 2);
DECL_FATTN_MMA_F16_CASE(96, 96, 16, 2);
DECL_FATTN_MMA_F16_CASE(112, 112, 16, 2);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu
index d2415bfa9..51c91c204 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_16-ncols2_4.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 16, 4);
DECL_FATTN_MMA_F16_CASE(64, 64, 16, 4);
+DECL_FATTN_MMA_F16_CASE(72, 72, 16, 4);
DECL_FATTN_MMA_F16_CASE(80, 80, 16, 4);
DECL_FATTN_MMA_F16_CASE(96, 96, 16, 4);
DECL_FATTN_MMA_F16_CASE(112, 112, 16, 4);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu
index 8eec1d74e..b6f99eb72 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_4.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 2, 4);
DECL_FATTN_MMA_F16_CASE(64, 64, 2, 4);
+DECL_FATTN_MMA_F16_CASE(72, 72, 2, 4);
DECL_FATTN_MMA_F16_CASE(80, 80, 2, 4);
DECL_FATTN_MMA_F16_CASE(96, 96, 2, 4);
DECL_FATTN_MMA_F16_CASE(112, 112, 2, 4);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu
index 9f4bef11a..14c120243 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_2-ncols2_8.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 2, 8);
DECL_FATTN_MMA_F16_CASE(64, 64, 2, 8);
+DECL_FATTN_MMA_F16_CASE(72, 72, 2, 8);
DECL_FATTN_MMA_F16_CASE(80, 80, 2, 8);
DECL_FATTN_MMA_F16_CASE(96, 96, 2, 8);
DECL_FATTN_MMA_F16_CASE(112, 112, 2, 8);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu
index 0543532ea..dbc58be55 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_1.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 32, 1);
DECL_FATTN_MMA_F16_CASE(64, 64, 32, 1);
+DECL_FATTN_MMA_F16_CASE(72, 72, 32, 1);
DECL_FATTN_MMA_F16_CASE(80, 80, 32, 1);
DECL_FATTN_MMA_F16_CASE(96, 96, 32, 1);
DECL_FATTN_MMA_F16_CASE(112, 112, 32, 1);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu
index d5a4d160a..784a1f919 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_32-ncols2_2.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 32, 2);
DECL_FATTN_MMA_F16_CASE(64, 64, 32, 2);
+DECL_FATTN_MMA_F16_CASE(72, 72, 32, 2);
DECL_FATTN_MMA_F16_CASE(80, 80, 32, 2);
DECL_FATTN_MMA_F16_CASE(96, 96, 32, 2);
DECL_FATTN_MMA_F16_CASE(112, 112, 32, 2);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu
index d8e7588cf..7ace51905 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_2.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 4, 2);
DECL_FATTN_MMA_F16_CASE(64, 64, 4, 2);
+DECL_FATTN_MMA_F16_CASE(72, 72, 4, 2);
DECL_FATTN_MMA_F16_CASE(80, 80, 4, 2);
DECL_FATTN_MMA_F16_CASE(96, 96, 4, 2);
DECL_FATTN_MMA_F16_CASE(112, 112, 4, 2);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu
index 3475dfea0..f27640e84 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_4.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 4, 4);
DECL_FATTN_MMA_F16_CASE(64, 64, 4, 4);
+DECL_FATTN_MMA_F16_CASE(72, 72, 4, 4);
DECL_FATTN_MMA_F16_CASE(80, 80, 4, 4);
DECL_FATTN_MMA_F16_CASE(96, 96, 4, 4);
DECL_FATTN_MMA_F16_CASE(112, 112, 4, 4);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu
index 859bea5c5..28a961b6b 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_4-ncols2_8.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 4, 8);
DECL_FATTN_MMA_F16_CASE(64, 64, 4, 8);
+DECL_FATTN_MMA_F16_CASE(72, 72, 4, 8);
DECL_FATTN_MMA_F16_CASE(80, 80, 4, 8);
DECL_FATTN_MMA_F16_CASE(96, 96, 4, 8);
DECL_FATTN_MMA_F16_CASE(112, 112, 4, 8);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu
index 0d7a9c728..1ad67c0d8 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_64-ncols2_1.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 64, 1);
DECL_FATTN_MMA_F16_CASE(64, 64, 64, 1);
+DECL_FATTN_MMA_F16_CASE(72, 72, 64, 1);
DECL_FATTN_MMA_F16_CASE(80, 80, 64, 1);
DECL_FATTN_MMA_F16_CASE(96, 96, 64, 1);
DECL_FATTN_MMA_F16_CASE(112, 112, 64, 1);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu
index 9d5a9976f..695782ca4 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_1.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 8, 1);
DECL_FATTN_MMA_F16_CASE(64, 64, 8, 1);
+DECL_FATTN_MMA_F16_CASE(72, 72, 8, 1);
DECL_FATTN_MMA_F16_CASE(80, 80, 8, 1);
DECL_FATTN_MMA_F16_CASE(96, 96, 8, 1);
DECL_FATTN_MMA_F16_CASE(112, 112, 8, 1);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu
index e48e28c4e..de2659815 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_2.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 8, 2);
DECL_FATTN_MMA_F16_CASE(64, 64, 8, 2);
+DECL_FATTN_MMA_F16_CASE(72, 72, 8, 2);
DECL_FATTN_MMA_F16_CASE(80, 80, 8, 2);
DECL_FATTN_MMA_F16_CASE(96, 96, 8, 2);
DECL_FATTN_MMA_F16_CASE(112, 112, 8, 2);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu
index 684cd25ce..9d9c748ae 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_4.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 8, 4);
DECL_FATTN_MMA_F16_CASE(64, 64, 8, 4);
+DECL_FATTN_MMA_F16_CASE(72, 72, 8, 4);
DECL_FATTN_MMA_F16_CASE(80, 80, 8, 4);
DECL_FATTN_MMA_F16_CASE(96, 96, 8, 4);
DECL_FATTN_MMA_F16_CASE(112, 112, 8, 4);
diff --git a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu
index c975ce6b9..50f17f2f8 100644
--- a/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu
+++ b/ggml/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_8-ncols2_8.cu
@@ -2,7 +2,9 @@
#include "../fattn-mma-f16.cuh"
+DECL_FATTN_MMA_F16_CASE(40, 40, 8, 8);
DECL_FATTN_MMA_F16_CASE(64, 64, 8, 8);
+DECL_FATTN_MMA_F16_CASE(72, 72, 8, 8);
DECL_FATTN_MMA_F16_CASE(80, 80, 8, 8);
DECL_FATTN_MMA_F16_CASE(96, 96, 8, 8);
DECL_FATTN_MMA_F16_CASE(112, 112, 8, 8);
diff --git a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
index 427db00e0..f49a015ae 100755
--- a/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
+++ b/ggml/src/ggml-cuda/template-instances/generate_cu_files.py
@@ -90,10 +90,6 @@ for ncols in [8, 16, 32, 64]:
f.write(SOURCE_FATTN_MMA_START)
for head_size_kq in HEAD_SIZES_KQ:
- if head_size_kq == 40:
- continue
- if head_size_kq == 72:
- continue
# Skip compilation of unused ncols2 values for niche head sizes:
if head_size_kq == 192 and ncols2 not in (8, 16): # MiMo-V2.5
continue
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 35b0e4eaa..7be4977d7 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -13676,6 +13676,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
// Qwen3-VL-8B https://github.com/ggml-org/llama.cpp/issues/17012
test_cases.emplace_back(new test_flash_attn_ext(72, 72, 16, {1, 1}, 5776, 5776, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+ // Same non-causal vision-encoder shape for the other head size padded in-kernel by the CUDA MMA kernels
+ test_cases.emplace_back(new test_flash_attn_ext(40, 40, 16, {1, 1}, 5776, 5776, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
// Sparse flash attention (n_kv_max hint) decode across KV depths.
// Shapes: 576/512 DeepSeek MLA, 512/512 DeepSeek-V4/GLM-5.2, 256/256 gqa12 Qwen QSA.