Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/fattn-common.cuh
+16 -9
View File
@@ -718,8 +718,8 @@ static __global__ void flash_attn_mask_to_KV_max(
template<int D, int ncols1, int ncols2> // D == head size
__launch_bounds__(D, 1)
static __global__ void flash_attn_stream_k_fixup_uniform(
float * __restrict__ dst,
const float2 * __restrict__ dst_fixup,
float * dst_ptr,
const float2 * dst_fixup_ptr,
const int ne01, const int ne02,
const int ne12, const int nblocks_stream_k,
const int gqa_ratio,
@@ -729,6 +729,8 @@ static __global__ void flash_attn_stream_k_fixup_uniform(
const uint3 fd_iter_j) {
constexpr int ncols = ncols1*ncols2;
ggml_cuda_pdl_lc();
float * GGML_CUDA_RESTRICT dst = dst_ptr;
const float2 * GGML_CUDA_RESTRICT dst_fixup = dst_fixup_ptr;
const int tile_idx = blockIdx.x; // One block per output tile.
const int j = blockIdx.y;
@@ -800,8 +802,8 @@ static __global__ void flash_attn_stream_k_fixup_uniform(
template <int D, int ncols1, int ncols2> // D == head size
__launch_bounds__(D, 1)
static __global__ void flash_attn_stream_k_fixup_general(
float * __restrict__ dst,
const float2 * __restrict__ dst_fixup,
float * dst_ptr,
const float2 * dst_fixup_ptr,
const int ne01, const int ne02,
const int gqa_ratio,
const int total_work,
@@ -809,6 +811,8 @@ static __global__ void flash_attn_stream_k_fixup_general(
const uint3 fd_iter_k_j_z,
const uint3 fd_iter_k_j,
const uint3 fd_iter_k) {
float * GGML_CUDA_RESTRICT dst = dst_ptr;
const float2 * GGML_CUDA_RESTRICT dst_fixup = dst_fixup_ptr;
constexpr int ncols = ncols1*ncols2;
const int bidx0 = blockIdx.x;
@@ -907,11 +911,14 @@ static __global__ void flash_attn_stream_k_fixup_general(
template<int D> // D == head size
__launch_bounds__(D, 1)
static __global__ void flash_attn_combine_results(
const float * __restrict__ VKQ_parts,
const float2 * __restrict__ VKQ_meta,
float * __restrict__ dst,
const float * VKQ_parts_ptr,
const float2 * VKQ_meta_ptr,
float * dst_ptr,
const int parallel_blocks) {
ggml_cuda_pdl_lc();
const float * GGML_CUDA_RESTRICT VKQ_parts = VKQ_parts_ptr;
const float2 * GGML_CUDA_RESTRICT VKQ_meta = VKQ_meta_ptr;
float * GGML_CUDA_RESTRICT dst = dst_ptr;
// Dimension 0: threadIdx.x
// Dimension 1: blockIdx.x
// Dimension 2: blockIdx.y
@@ -1196,8 +1203,8 @@ void launch_fattn(
GGML_ASSERT(block_dim.x % warp_size == 0);
// disabled PDL enrollment for now due to a compiler bug.
fattn_kernel<<<blocks_num, block_dim, nbytes_shared, main_stream>>>(
ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num, block_dim, nbytes_shared, main_stream);
ggml_cuda_kernel_launch(fattn_kernel, launch_params,
(const char *) Q->data,
K_data,
V_data,