From b820cc8e6f94843d32f92c8ebd7db837dae2bd8b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20G=C3=A4=C3=9Fler?= Date: Wed, 1 Jul 2026 10:55:14 +0200 Subject: [PATCH] CUDA: consistent use of __restrict__ + PDL for FA (#25185) --- ggml/src/ggml-cuda/fattn-common.cuh | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index b76121fef3..e67cc7fdf7 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -664,7 +664,10 @@ constexpr __device__ dequantize_V_t get_dequantize_V() { template __launch_bounds__(FATTN_KQ_STRIDE/2, 1) static __global__ void flash_attn_mask_to_KV_max( - const half2 * __restrict__ mask, int * __restrict__ KV_max, const int ne30, const int64_t s31, const int64_t s33) { + const half2 * mask_ptr, int * KV_max_ptr, const int ne30, const int64_t s31, const int64_t s33) { + const half2 * GGML_CUDA_RESTRICT mask = mask_ptr; + int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; + const int ne31 = gridDim.x; const int tid = threadIdx.x; const int sequence = blockIdx.y; @@ -1099,8 +1102,9 @@ void launch_fattn( const int iter_k = K->ne[1] / FATTN_KQ_STRIDE; KV_max.alloc(ne_KV_max); - flash_attn_mask_to_KV_max<<>> - ((const half2 *) mask->data, KV_max.ptr, iter_k, s31, s33); + ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num_KV_max, block_dim_KV_max, 0, main_stream); + ggml_cuda_kernel_launch(flash_attn_mask_to_KV_max, launch_params, + (const half2 *) mask->data, KV_max.ptr, iter_k, s31, s33); CUDA_CHECK(cudaGetLastError()); }