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
@@ -17,14 +17,22 @@ using namespace cub;
|
||||
#endif // __clang__
|
||||
template <size_t splitD, size_t N, size_t L_template>
|
||||
__global__ void __launch_bounds__(splitD, 1)
|
||||
ssm_scan_f32(const float *__restrict__ src0, const float *__restrict__ src1, const float *__restrict__ src2,
|
||||
const float *__restrict__ src3, const float *__restrict__ src4, const float *__restrict__ src5,
|
||||
const int32_t * __restrict__ src6, float * __restrict__ dst,
|
||||
ssm_scan_f32(const float * src0_ptr, const float * src1_ptr, const float * src2_ptr,
|
||||
const float * src3_ptr, const float * src4_ptr, const float * src5_ptr,
|
||||
const int32_t * src6_ptr, float * dst_ptr,
|
||||
const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3,
|
||||
const int src2_nb1, const int src2_nb2, const int src3_nb1,
|
||||
const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3,
|
||||
const int64_t s_off, const int64_t d_inner, const int64_t L_param)
|
||||
{
|
||||
const float * GGML_CUDA_RESTRICT src0 = src0_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src1 = src1_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src2 = src2_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src3 = src3_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src4 = src4_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src5 = src5_ptr;
|
||||
const int32_t * GGML_CUDA_RESTRICT src6 = src6_ptr;
|
||||
float * GGML_CUDA_RESTRICT dst = dst_ptr;
|
||||
const size_t L = L_template == 0 ? L_param : L_template;
|
||||
ggml_cuda_pdl_sync();
|
||||
const float *s0_block = (const float *)((const char *)src0 + src6[blockIdx.x] * src0_nb3 + blockIdx.y * splitD * src0_nb2);
|
||||
@@ -118,13 +126,21 @@ __global__ void __launch_bounds__(splitD, 1)
|
||||
template <int c_factor, int d_state>
|
||||
__global__ void __launch_bounds__(d_state, 1)
|
||||
ssm_scan_f32_group(
|
||||
const float * __restrict__ src0, const float * __restrict__ src1, const float * __restrict__ src2,
|
||||
const float * __restrict__ src3, const float * __restrict__ src4, const float * __restrict__ src5,
|
||||
const int32_t * __restrict__ src6, float * __restrict__ dst,
|
||||
const float * src0_ptr, const float * src1_ptr, const float * src2_ptr,
|
||||
const float * src3_ptr, const float * src4_ptr, const float * src5_ptr,
|
||||
const int32_t * src6_ptr, float * dst_ptr,
|
||||
const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3,
|
||||
const int src2_nb1, const int src2_nb2, const int src3_nb1,
|
||||
const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3,
|
||||
const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok) {
|
||||
const float * GGML_CUDA_RESTRICT src0 = src0_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src1 = src1_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src2 = src2_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src3 = src3_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src4 = src4_ptr;
|
||||
const float * GGML_CUDA_RESTRICT src5 = src5_ptr;
|
||||
const int32_t * GGML_CUDA_RESTRICT src6 = src6_ptr;
|
||||
float * GGML_CUDA_RESTRICT dst = dst_ptr;
|
||||
|
||||
const int warp = threadIdx.x / WARP_SIZE;
|
||||
const int lane = threadIdx.x % WARP_SIZE;
|
||||
|
||||
Reference in New Issue
Block a user