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:
co-authored by
Oliver Simons
parent
3571fa5435
commit
9e58d4d692
@@ -1703,14 +1703,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
|
||||
template<int DKQ, int DV, int ncols1, int ncols2, bool use_logit_softcap, bool V_is_K_view>
|
||||
__launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2))
|
||||
static __global__ void flash_attn_ext_f16(
|
||||
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,
|
||||
@@ -1726,6 +1726,14 @@ static __global__ void flash_attn_ext_f16(
|
||||
const int32_t nb31, const int32_t nb32, const int64_t nb33) {
|
||||
ggml_cuda_pdl_sync(); // TODO optimize placement
|
||||
#if defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_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:
|
||||
if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) {
|
||||
@@ -1871,7 +1879,7 @@ static __global__ void flash_attn_ext_f16(
|
||||
(Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
|
||||
ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
|
||||
#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,
|
||||
|
||||
Reference in New Issue
Block a user