Avoid PDL race conditions by disabling __restrict__ when PDL is used (#24030)

* Removes __restrict__ from PDL kernel headers due to incompatibility with
PDL. Adds preprocessor directives based on arch in kernel body to add
__restrict__ to retain performance on older architectures.

* Simplifies new __restrict__ usage via macro

* Add hopper to PDL __restrict__ fix.

Co-authored-by: Oliver Simons <osimons@nvidia.com>

---------

Co-authored-by: Oliver Simons <osimons@nvidia.com>
This commit is contained in:
Andreas Kieslinger
2026-06-03 13:56:42 +02:00
committed by GitHub
co-authored by Oliver Simons
parent 3571fa5435
commit 9e58d4d692
14 changed files with 145 additions and 62 deletions
+17 -9
View File
@@ -788,14 +788,14 @@ static __device__ __forceinline__ void flash_attn_tile_iter(
template<int DKQ, int DV, int ncols1, int ncols2, bool use_logit_softcap> // D == head size
__launch_bounds__(ggml_cuda_fattn_tile_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_tile_get_occupancy(DKQ, DV, ncols1*ncols2))
static __global__ void flash_attn_tile(
const char * __restrict__ Q,
const char * __restrict__ K,
const char * __restrict__ V,
const char * __restrict__ mask,
const char * __restrict__ sinks,
const int * __restrict__ KV_max,
float * __restrict__ dst,
float2 * __restrict__ dst_meta,
const char * Q_ptr,
const char * K_ptr,
const char * V_ptr,
const char * mask_ptr,
const char * sinks_ptr,
const int * KV_max_ptr,
float * dst_ptr,
float2 * dst_meta_ptr,
const float scale,
const float max_bias,
const float m0,
@@ -810,6 +810,14 @@ static __global__ void flash_attn_tile(
const int32_t ne31, const int32_t ne32, const int32_t ne33,
const int32_t nb31, const int32_t nb32, const int64_t nb33) {
#ifdef FLASH_ATTN_AVAILABLE
const char * GGML_CUDA_RESTRICT Q = Q_ptr;
const char * GGML_CUDA_RESTRICT K = K_ptr;
const char * GGML_CUDA_RESTRICT V = V_ptr;
const char * GGML_CUDA_RESTRICT mask = mask_ptr;
const char * GGML_CUDA_RESTRICT sinks = sinks_ptr;
const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr;
float * GGML_CUDA_RESTRICT dst = dst_ptr;
float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
// Skip unused kernel variants for faster compilation:
@@ -1126,7 +1134,7 @@ static __global__ void flash_attn_tile(
}
}
#else
GGML_UNUSED_VARS(Q, K, V, mask, sinks, KV_max, dst, dst_meta, scale,
GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
max_bias, m0, m1, n_head_log2, logit_softcap,
ne00, ne01, ne02, ne03,
nb01, nb02, nb03,