Commit 6e66cb5d4 for llama.cpp
commit 6e66cb5d4e9a2b4f02820a9f3053422b92ec4d51
Author: Ethan Guo <ethanguo.dev@gmail.com>
Date: Thu Oct 15 14:39:47 2026 +0800
metal : add BF16 variant of the tensor API flash attention kernel (#30491)
* metal : add BF16 variant of the tensor API flash attention kernel
* cont : move tensor FA kernels to kernels/fa_tensor.metal
diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt
index 05cdea156..91691b223 100644
--- a/ggml/src/ggml-metal/CMakeLists.txt
+++ b/ggml/src/ggml-metal/CMakeLists.txt
@@ -45,6 +45,7 @@ set(METALLIB_KERNEL_SOURCES
kernels/fa_q5_0.metal
kernels/fa_q5_1.metal
kernels/fa_q8_0.metal
+ kernels/fa_tensor.metal
kernels/fa_vec_f16.metal
kernels/fa_vec_f32.metal
kernels/fa_vec_q4_0.metal
@@ -233,12 +234,12 @@ else()
VERBATIM
)
- set(AIR_FA_TENSOR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/fa_f16_tensor.air")
+ set(AIR_FA_TENSOR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/fa_tensor_tensor.air")
add_custom_command(
OUTPUT ${AIR_FA_TENSOR}
- COMMAND xcrun -sdk ${METAL_SDK} metal ${XC_FLAGS_TENSOR} -DGGML_METAL_HAS_TENSOR -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/fa_f16.metal -o ${AIR_FA_TENSOR}
- DEPENDS kernels/fa_f16.metal ${METALLIB_KERNELS_FA_SHARED} kernels/common.h kernels/dequantize.h ${METALLIB_COMMON} ggml-metal-impl.h
- COMMENT "Compiling kernels/fa_f16.metal (tensor API)"
+ COMMAND xcrun -sdk ${METAL_SDK} metal ${XC_FLAGS_TENSOR} -DGGML_METAL_HAS_TENSOR -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/fa_tensor.metal -o ${AIR_FA_TENSOR}
+ DEPENDS kernels/fa_tensor.metal kernels/common.h ${METALLIB_COMMON} ggml-metal-impl.h
+ COMMENT "Compiling kernels/fa_tensor.metal (tensor API)"
VERBATIM
)
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index b70129de8..2c7972b0a 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -1866,7 +1866,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_t
const int32_t dk = (int32_t) op->src[1]->ne[0];
const int32_t dv = (int32_t) op->src[2]->ne[0];
- snprintf(base, 256, "kernel_flash_attn_ext_tensor_f16_dk%d_dv%d", dk, dv);
+ const char * type = op->src[1]->type == GGML_TYPE_BF16 ? "bf16" : "f16";
+
+ snprintf(base, 256, "kernel_flash_attn_ext_tensor_%s_dk%d_dv%d", type, dk, dv);
snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d", base, has_mask, has_sinks, has_bias, has_scap);
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index 13aecbdd3..d1bce8bad 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -118,6 +118,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi
X(FA_Q5_0, fa_q5_0) \
X(FA_Q5_1, fa_q5_1) \
X(FA_Q8_0, fa_q8_0) \
+ X(FA_TENSOR, fa_tensor) \
X(FA_VEC_F16, fa_vec_f16) \
X(FA_VEC_F32, fa_vec_f32) \
X(FA_VEC_Q4_0, fa_vec_q4_0) \
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 115d73695..53c981618 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -3161,7 +3161,7 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool
return false;
}
- if (op->src[1]->type != GGML_TYPE_F16 && !ggml_metal_op_flash_attn_ext_use_kv_f16(op)) {
+ if (op->src[1]->type != GGML_TYPE_F16 && op->src[1]->type != GGML_TYPE_BF16 && !ggml_metal_op_flash_attn_ext_use_kv_f16(op)) {
return false;
}
@@ -3678,7 +3678,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
/*.logit_softcap =*/ logit_softcap,
};
- // shared memory layout: queries (half), scores (float), probabilities (half), row scale (float), rescale flag (int)
+ // shared memory layout: queries (16-bit), scores (float), probabilities (16-bit), row scale (float), rescale flag (int)
const size_t smem = GGML_PAD(nqptg*ne00*sizeof(ggml_fp16_t) + nqptg*ncpsg*(sizeof(float) + sizeof(ggml_fp16_t)) + nqptg*sizeof(float) + sizeof(int32_t), 16);
auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_tensor(lib, op, has_mask, has_sinks, has_bias, has_scap);
diff --git a/ggml/src/ggml-metal/kernels/fa_f16.metal b/ggml/src/ggml-metal/kernels/fa_f16.metal
index 0f44b419b..6e30b09e3 100644
--- a/ggml/src/ggml-metal/kernels/fa_f16.metal
+++ b/ggml/src/ggml-metal/kernels/fa_f16.metal
@@ -75,291 +75,3 @@ template [[host_name("kernel_flash_attn_ext_bf16_dk576_dv512")]] kernel flash_at
#undef FA_TYPES
#undef FA_TYPES_BF
#undef FA_TYPES_F32
-
-#ifdef GGML_METAL_HAS_TENSOR
-
-constant bool FC_flash_attn_ext_tensor_has_mask [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 0)]];
-constant bool FC_flash_attn_ext_tensor_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 1)]];
-constant bool FC_flash_attn_ext_tensor_has_bias [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 2)]];
-constant bool FC_flash_attn_ext_tensor_has_scap [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 3)]];
-
-// ref: https://arxiv.org/pdf/2307.08691.pdf
-template<
- short DK, // K head size
- short DV, // V head size
- short Q = OP_FLASH_ATTN_EXT_TENSOR_NQPSG, // queries per threadgroup
- short C = OP_FLASH_ATTN_EXT_TENSOR_NCPSG, // cache items per threadgroup
- short NSG = OP_FLASH_ATTN_EXT_TENSOR_NSG> // number of simd groups
-kernel void kernel_flash_attn_ext_tensor(
- constant ggml_metal_kargs_flash_attn_ext & args,
- device const char * q,
- device const char * k,
- device const char * v,
- device const char * mask,
- device const char * sinks,
- device const char * blk,
- device char * dst,
- threadgroup char * shmem [[threadgroup(0)]],
- uint3 tgpig [[threadgroup_position_in_grid]],
- ushort tiisg [[thread_index_in_simdgroup]],
- ushort sgitg [[simdgroup_index_in_threadgroup]]) {
- constexpr short NW = N_SIMDWIDTH;
- constexpr short NT = NW*NSG;
- constexpr short NQ = Q/NSG;
- constexpr short NC = C/NW; // columns per thread
-
- static_assert(DK % 4 == 0, "DK must be divisible by 4");
- static_assert(Q % NSG == 0, "Q must be divisible by NSG");
- static_assert(C % NW == 0, "C must be divisible by NW");
-
- const int iq3 = tgpig[2];
- const int iq2 = tgpig[1];
- const int iq1 = tgpig[0]*Q;
-
- const short tiitg = sgitg*NW + tiisg;
-
- threadgroup half * sq = (threadgroup half *) shmem; // [Q, DK] queries
- threadgroup float * ss = (threadgroup float *) (sq + Q*DK); // [Q, C] scores
- threadgroup half * sp = (threadgroup half *) (ss + Q*C); // [Q, C] probabilities
- threadgroup float * sr = (threadgroup float *) (sp + Q*C); // [Q] per-row scale of O
- threadgroup int * sf = (threadgroup int *) (sr + Q); // [1] last iteration (ic0 + 1) that rescaled O
-
- q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03;
-
- {
- const int ikv2 = iq2/(args.ne02/args.ne_12_2);
- const int ikv3 = iq3/(args.ne03/args.ne_12_3);
-
- k += ikv2*args.nb12 + ikv3*args.nb13;
- v += ikv2*args.nb22 + ikv3*args.nb23;
- }
-
- // with softcap the scale is small (scale/softcap), so it is applied to the scores to keep the precision of Q
- const float qscale = FC_flash_attn_ext_tensor_has_scap ? 1.0f : args.scale;
-
- // load the queries, with the scale folded in
- for (int i = tiitg; i < Q*DK/4; i += NT) {
- const int j = i/(DK/4);
-
- float4 q4 = 0.0f;
- if (iq1 + j < args.ne01) {
- q4 = ((device const float4 *) (q + j*args.nb01))[i%(DK/4)];
- }
-
- ((threadgroup half4 *) sq)[i] = (half4) (q4*qscale);
- }
-
- device const half * pm[NQ];
-
- FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
- const short j = jj*NSG + sgitg;
-
- pm[jj] = (device const half *) (mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
- }
-
- {
- const int nblk1 = (args.ne01 + Q - 1)/Q;
- const int nblk0 = (args.ne11 + C - 1)/C;
-
- blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0;
- }
-
- float M[NQ];
- float S[NQ];
-
- FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
- M[jj] = -FLT_MAX/2;
- S[jj] = 0.0f;
- }
-
- float slope = 1.0f;
-
- // ALiBi
- if (FC_flash_attn_ext_tensor_has_bias) {
- const short h = iq2;
-
- const float base = h < args.n_head_log2 ? args.m0 : args.m1;
- const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
-
- slope = pow(base, exph);
- }
-
- const int sk = args.ns10;
- const int sv = args.ns20;
-
- auto tq = tensor(sq, dextents<int32_t, 2>(DK, Q));
- auto ts = tensor(ss, dextents<int32_t, 2>(C, Q));
- auto tp = tensor(sp, dextents<int32_t, 2>(C, Q));
-
- mpp::tensor_ops::matmul2d<
- mpp::tensor_ops::matmul2d_descriptor(Q, C, DK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply),
- execution_simdgroups<NSG>> mm_qk;
-
- mpp::tensor_ops::matmul2d<
- mpp::tensor_ops::matmul2d_descriptor(Q, DV, C, false, false, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
- execution_simdgroups<NSG>> mm_pv;
-
- auto tv0 = tensor((device half *) v, dextents<int32_t, 2>(DV, C), array<int, 2>({1, sv}));
-
- // the O matrix from the paper
- auto co = mm_pv.template get_destination_cooperative_tensor<decltype(tp), decltype(tv0), float>();
-
- FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
- if (co.is_valid_element(i)) {
- co[i] = 0.0f;
- }
- }
-
- if (tiitg == 0) {
- sf[0] = 0;
- }
-
- threadgroup_barrier(mem_flags::mem_threadgroup);
-
- // the host guarantees ne11 % C == 0
- for (int ic0 = 0, ic = 0; ic < args.ne11; ++ic0, ic += C) {
- char blk_cur = 1;
-
- if (FC_flash_attn_ext_tensor_has_mask) {
- blk_cur = blk[ic0];
-
- if (blk_cur == 0) {
- continue;
- }
- }
-
- // Q*K^T
- {
- auto tk = tensor((device half *) (k + (uint64_t) ic*args.nb11), dextents<int32_t, 2>(DK, C), array<int, 2>({1, sk}));
-
- mm_qk.run(tq, tk, ts);
- }
-
- threadgroup_barrier(mem_flags::mem_threadgroup);
-
- // online softmax
- FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
- const short j = jj*NSG + sgitg;
-
- float s[NC];
-
- FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
- s[ii] = ss[j*C + ii*NW + tiisg];
- }
-
- if (FC_flash_attn_ext_tensor_has_scap) {
- FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
- s[ii] = args.logit_softcap*precise::tanh(s[ii]*args.scale);
- }
- }
-
- if (FC_flash_attn_ext_tensor_has_mask && blk_cur != 2 && iq1 + j < args.ne31) {
- FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
- s[ii] += slope*(float) pm[jj][ic + ii*NW + tiisg];
- }
- }
-
- float m = M[jj];
-
- FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
- m = max(m, s[ii]);
- }
-
- m = simd_max(m);
-
- // lazy rescaling: move the running max only when it grows by more than 8 (e^8 fits in half)
- float ms = 1.0f;
-
- if (m > M[jj] + 8.0f) {
- ms = exp(M[jj] - m);
- M[jj] = m;
-
- if (tiisg == 0) {
- sf[0] = ic0 + 1;
- }
- }
-
- float sum = 0.0f;
-
- FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
- // the sum uses the same rounded values as P*V
- const half p = (half) exp(s[ii] - M[jj]);
-
- sp[j*C + ii*NW + tiisg] = p;
-
- sum += (float) p;
- }
-
- S[jj] = S[jj]*ms + simd_sum(sum);
-
- if (tiisg == 0) {
- sr[j] = ms;
- }
- }
-
- threadgroup_barrier(mem_flags::mem_threadgroup);
-
- // O = diag(ms)*O + P*V
- if (sf[0] == ic0 + 1) {
- FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
- if (co.is_valid_element(i)) {
- co[i] *= sr[co.get_multidimensional_index(i)[1]];
- }
- }
- }
-
- {
- auto tv = tensor((device half *) (v + (uint64_t) ic*args.nb21), dextents<int32_t, 2>(DV, C), array<int, 2>({1, sv}));
-
- mm_pv.run(tp, tv, co);
- }
-
- threadgroup_barrier(mem_flags::mem_threadgroup);
- }
-
- FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
- const short j = jj*NSG + sgitg;
-
- // the sink only adds to the denominator - its rescale of O is folded into the final scale
- float ms = 1.0f;
-
- if (FC_flash_attn_ext_tensor_has_sinks) {
- const float s = ((device const float *) sinks)[iq2];
- const float m = max(M[jj], s);
-
- ms = exp(M[jj] - m);
-
- S[jj] = S[jj]*ms + exp(s - m);
- }
-
- if (tiisg == 0) {
- sr[j] = S[jj] == 0.0f ? 0.0f : ms/S[jj];
- }
- }
-
- threadgroup_barrier(mem_flags::mem_threadgroup);
-
- FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
- if (co.is_valid_element(i)) {
- co[i] *= sr[co.get_multidimensional_index(i)[1]];
- }
- }
-
- // store to global memory - rows past ne01 are clipped by the tensor extents
- device float * pdst = (device float *) dst + ((uint64_t) iq3*args.ne2*args.ne1 + iq2 + (uint64_t) iq1*args.ne1)*DV;
-
- auto td = tensor(pdst, dextents<int32_t, 2>(DV, args.ne01 - iq1), array<int, 2>({1, args.ne1*DV}));
-
- co.store(td);
-}
-
-typedef decltype(kernel_flash_attn_ext_tensor<64, 64>) flash_attn_ext_tensor_t;
-
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk64_dv64" )]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<64, 64>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk128_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<128, 128>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk192_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<192, 128>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk256_dv256")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<256, 256>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk512_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<512, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
-template [[host_name("kernel_flash_attn_ext_tensor_f16_dk576_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<576, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
-
-#endif // GGML_METAL_HAS_TENSOR
diff --git a/ggml/src/ggml-metal/kernels/fa_tensor.metal b/ggml/src/ggml-metal/kernels/fa_tensor.metal
new file mode 100644
index 000000000..4a9ed30e3
--- /dev/null
+++ b/ggml/src/ggml-metal/kernels/fa_tensor.metal
@@ -0,0 +1,298 @@
+#include "common.h"
+
+#ifdef GGML_METAL_HAS_TENSOR
+
+constant bool FC_flash_attn_ext_tensor_has_mask [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 0)]];
+constant bool FC_flash_attn_ext_tensor_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 1)]];
+constant bool FC_flash_attn_ext_tensor_has_bias [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 2)]];
+constant bool FC_flash_attn_ext_tensor_has_scap [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 3)]];
+
+// ref: https://arxiv.org/pdf/2307.08691.pdf
+template<
+ typename T, // Q, K, V and P types
+ typename T4,
+ short DK, // K head size
+ short DV, // V head size
+ short Q = OP_FLASH_ATTN_EXT_TENSOR_NQPSG, // queries per threadgroup
+ short C = OP_FLASH_ATTN_EXT_TENSOR_NCPSG, // cache items per threadgroup
+ short NSG = OP_FLASH_ATTN_EXT_TENSOR_NSG> // number of simd groups
+kernel void kernel_flash_attn_ext_tensor(
+ constant ggml_metal_kargs_flash_attn_ext & args,
+ device const char * q,
+ device const char * k,
+ device const char * v,
+ device const char * mask,
+ device const char * sinks,
+ device const char * blk,
+ device char * dst,
+ threadgroup char * shmem [[threadgroup(0)]],
+ uint3 tgpig [[threadgroup_position_in_grid]],
+ ushort tiisg [[thread_index_in_simdgroup]],
+ ushort sgitg [[simdgroup_index_in_threadgroup]]) {
+ constexpr short NW = N_SIMDWIDTH;
+ constexpr short NT = NW*NSG;
+ constexpr short NQ = Q/NSG;
+ constexpr short NC = C/NW; // columns per thread
+
+ static_assert(DK % 4 == 0, "DK must be divisible by 4");
+ static_assert(Q % NSG == 0, "Q must be divisible by NSG");
+ static_assert(C % NW == 0, "C must be divisible by NW");
+
+ const int iq3 = tgpig[2];
+ const int iq2 = tgpig[1];
+ const int iq1 = tgpig[0]*Q;
+
+ const short tiitg = sgitg*NW + tiisg;
+
+ threadgroup T * sq = (threadgroup T *) shmem; // [Q, DK] queries
+ threadgroup float * ss = (threadgroup float *) (sq + Q*DK); // [Q, C] scores
+ threadgroup T * sp = (threadgroup T *) (ss + Q*C); // [Q, C] probabilities
+ threadgroup float * sr = (threadgroup float *) (sp + Q*C); // [Q] per-row scale of O
+ threadgroup int * sf = (threadgroup int *) (sr + Q); // [1] last iteration (ic0 + 1) that rescaled O
+
+ q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03;
+
+ {
+ const int ikv2 = iq2/(args.ne02/args.ne_12_2);
+ const int ikv3 = iq3/(args.ne03/args.ne_12_3);
+
+ k += ikv2*args.nb12 + ikv3*args.nb13;
+ v += ikv2*args.nb22 + ikv3*args.nb23;
+ }
+
+ // with softcap the scale is small (scale/softcap), so it is applied to the scores to keep the precision of Q
+ const float qscale = FC_flash_attn_ext_tensor_has_scap ? 1.0f : args.scale;
+
+ // load the queries, with the scale folded in
+ for (int i = tiitg; i < Q*DK/4; i += NT) {
+ const int j = i/(DK/4);
+
+ float4 q4 = 0.0f;
+ if (iq1 + j < args.ne01) {
+ q4 = ((device const float4 *) (q + j*args.nb01))[i%(DK/4)];
+ }
+
+ ((threadgroup T4 *) sq)[i] = (T4) (q4*qscale);
+ }
+
+ device const half * pm[NQ];
+
+ FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
+ const short j = jj*NSG + sgitg;
+
+ pm[jj] = (device const half *) (mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
+ }
+
+ {
+ const int nblk1 = (args.ne01 + Q - 1)/Q;
+ const int nblk0 = (args.ne11 + C - 1)/C;
+
+ blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0;
+ }
+
+ float M[NQ];
+ float S[NQ];
+
+ FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
+ M[jj] = -FLT_MAX/2;
+ S[jj] = 0.0f;
+ }
+
+ float slope = 1.0f;
+
+ // ALiBi
+ if (FC_flash_attn_ext_tensor_has_bias) {
+ const short h = iq2;
+
+ const float base = h < args.n_head_log2 ? args.m0 : args.m1;
+ const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
+
+ slope = pow(base, exph);
+ }
+
+ const int sk = args.ns10;
+ const int sv = args.ns20;
+
+ auto tq = tensor(sq, dextents<int32_t, 2>(DK, Q));
+ auto ts = tensor(ss, dextents<int32_t, 2>(C, Q));
+ auto tp = tensor(sp, dextents<int32_t, 2>(C, Q));
+
+ mpp::tensor_ops::matmul2d<
+ mpp::tensor_ops::matmul2d_descriptor(Q, C, DK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply),
+ execution_simdgroups<NSG>> mm_qk;
+
+ mpp::tensor_ops::matmul2d<
+ mpp::tensor_ops::matmul2d_descriptor(Q, DV, C, false, false, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
+ execution_simdgroups<NSG>> mm_pv;
+
+ auto tv0 = tensor((device T *) v, dextents<int32_t, 2>(DV, C), array<int, 2>({1, sv}));
+
+ // the O matrix from the paper
+ auto co = mm_pv.template get_destination_cooperative_tensor<decltype(tp), decltype(tv0), float>();
+
+ FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
+ if (co.is_valid_element(i)) {
+ co[i] = 0.0f;
+ }
+ }
+
+ if (tiitg == 0) {
+ sf[0] = 0;
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ // the host guarantees ne11 % C == 0
+ for (int ic0 = 0, ic = 0; ic < args.ne11; ++ic0, ic += C) {
+ char blk_cur = 1;
+
+ if (FC_flash_attn_ext_tensor_has_mask) {
+ blk_cur = blk[ic0];
+
+ if (blk_cur == 0) {
+ continue;
+ }
+ }
+
+ // Q*K^T
+ {
+ auto tk = tensor((device T *) (k + (uint64_t) ic*args.nb11), dextents<int32_t, 2>(DK, C), array<int, 2>({1, sk}));
+
+ mm_qk.run(tq, tk, ts);
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ // online softmax
+ FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
+ const short j = jj*NSG + sgitg;
+
+ float s[NC];
+
+ FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+ s[ii] = ss[j*C + ii*NW + tiisg];
+ }
+
+ if (FC_flash_attn_ext_tensor_has_scap) {
+ FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+ s[ii] = args.logit_softcap*precise::tanh(s[ii]*args.scale);
+ }
+ }
+
+ if (FC_flash_attn_ext_tensor_has_mask && blk_cur != 2 && iq1 + j < args.ne31) {
+ FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+ s[ii] += slope*(float) pm[jj][ic + ii*NW + tiisg];
+ }
+ }
+
+ float m = M[jj];
+
+ FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+ m = max(m, s[ii]);
+ }
+
+ m = simd_max(m);
+
+ // lazy rescaling: move the running max only when it grows by more than 8 (e^8 fits in half)
+ float ms = 1.0f;
+
+ if (m > M[jj] + 8.0f) {
+ ms = exp(M[jj] - m);
+ M[jj] = m;
+
+ if (tiisg == 0) {
+ sf[0] = ic0 + 1;
+ }
+ }
+
+ float sum = 0.0f;
+
+ FOR_UNROLL (short ii = 0; ii < NC; ++ii) {
+ // the sum uses the same rounded values as P*V
+ const T p = (T) exp(s[ii] - M[jj]);
+
+ sp[j*C + ii*NW + tiisg] = p;
+
+ sum += (float) p;
+ }
+
+ S[jj] = S[jj]*ms + simd_sum(sum);
+
+ if (tiisg == 0) {
+ sr[j] = ms;
+ }
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ // O = diag(ms)*O + P*V
+ if (sf[0] == ic0 + 1) {
+ FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
+ if (co.is_valid_element(i)) {
+ co[i] *= sr[co.get_multidimensional_index(i)[1]];
+ }
+ }
+ }
+
+ {
+ auto tv = tensor((device T *) (v + (uint64_t) ic*args.nb21), dextents<int32_t, 2>(DV, C), array<int, 2>({1, sv}));
+
+ mm_pv.run(tp, tv, co);
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ }
+
+ FOR_UNROLL (short jj = 0; jj < NQ; ++jj) {
+ const short j = jj*NSG + sgitg;
+
+ // the sink only adds to the denominator - its rescale of O is folded into the final scale
+ float ms = 1.0f;
+
+ if (FC_flash_attn_ext_tensor_has_sinks) {
+ const float s = ((device const float *) sinks)[iq2];
+ const float m = max(M[jj], s);
+
+ ms = exp(M[jj] - m);
+
+ S[jj] = S[jj]*ms + exp(s - m);
+ }
+
+ if (tiisg == 0) {
+ sr[j] = S[jj] == 0.0f ? 0.0f : ms/S[jj];
+ }
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) {
+ if (co.is_valid_element(i)) {
+ co[i] *= sr[co.get_multidimensional_index(i)[1]];
+ }
+ }
+
+ // store to global memory - rows past ne01 are clipped by the tensor extents
+ device float * pdst = (device float *) dst + ((uint64_t) iq3*args.ne2*args.ne1 + iq2 + (uint64_t) iq1*args.ne1)*DV;
+
+ auto td = tensor(pdst, dextents<int32_t, 2>(DV, args.ne01 - iq1), array<int, 2>({1, args.ne1*DV}));
+
+ co.store(td);
+}
+
+typedef decltype(kernel_flash_attn_ext_tensor<half, half4, 64, 64>) flash_attn_ext_tensor_t;
+
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk64_dv64" )]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 64, 64>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk128_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 128, 128>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk192_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 192, 128>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk256_dv256")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 256, 256>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk512_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 512, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
+template [[host_name("kernel_flash_attn_ext_tensor_f16_dk576_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<half, half4, 576, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
+
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk64_dv64" )]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 64, 64>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk128_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 128, 128>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk192_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 192, 128>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk256_dv256")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 256, 256>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk512_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 512, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
+template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk576_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<bfloat, bfloat4, 576, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>;
+
+#endif // GGML_METAL_HAS_TENSOR
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 0d7a7ed81..e61823084 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -11813,11 +11813,18 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {2, 1}, 1024, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, {2, 1}, 1024, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+ // BF16 K/V for the other head sizes of the Metal tensor API kernel
+ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {12, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16));
+ test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, { 8, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16));
+ test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, { 4, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16));
+ test_cases.emplace_back(new test_flash_attn_ext(576, 512, 4, { 4, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 1, 2, 3}, true, true));
+
// FLASH_ATTN_EXT: large logits
test_cases.emplace_back(new test_flash_attn_ext_large_logits( 64, 64, 16, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext_large_logits(128, 128, 8, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext_large_logits(256, 256, 4, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext_large_logits(256, 256, 4, {4, 1}, 1024, 75, true, false, 0, 10.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
+ test_cases.emplace_back(new test_flash_attn_ext_large_logits(256, 256, 4, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16));
test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3}));
test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1}));