CUDA: fix FA tg at long context for CC >= 8.9 (#13852)

JohannesGaessler · web-flow · commit a68247439bd6 · 2025-05-28T13:33:37.000+02:00
diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh
@@ -623,8 +623,8 @@ static __global__ void flash_attn_combine_results(
     __builtin_assume(tid < D);
 
     extern __shared__ float2 meta[];
-    if (tid < 2*parallel_blocks) {
-        ((float *) meta)[threadIdx.x] = ((const float *)VKQ_meta) [blockIdx.z*(2*parallel_blocks) + tid];
+    for (int i = tid; i < 2*parallel_blocks; i += D) {
+        ((float *) meta)[i] = ((const float *)VKQ_meta) [blockIdx.z*(2*parallel_blocks) + i];
     }
 
     __syncthreads();

Original file line number	Diff line number	Diff line change
`@@ -623,8 +623,8 @@ static __global__ void flash_attn_combine_results(`
`623`	`623`	`__builtin_assume(tid < D);`
`624`	`624`
`625`	`625`	`extern __shared__ float2 meta[];`
`626`		`- if (tid < 2*parallel_blocks) {`
`627`		`- ((float ) meta)[threadIdx.x] = ((const float )VKQ_meta) [blockIdx.z(2parallel_blocks) + tid];`
	`626`	`+ for (int i = tid; i < 2*parallel_blocks; i += D) {`
	`627`	`+ ((float ) meta)[i] = ((const float )VKQ_meta) [blockIdx.z(2parallel_blocks) + i];`
`628`	`628`	`}`
`629`	`629`
`630`	`630`	`__syncthreads();`