From 0893cb46834f6abe64c8ef2da628b8694ca27293 Mon Sep 17 00:00:00 2001
From: Aman Gupta <amangupta052@gmail.com>
Date: Sat, 29 Aug 2026 15:57:11 +0800
Subject: [PATCH] fix PDL issues

---
 ggml/src/ggml-cuda/fattn.cu | 14 ++++++++++----
 src/models/qwen4exp.cpp     |  2 +-
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index 2f1929bbe1..d1e6bb4ee6 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -10,6 +10,8 @@ __launch_bounds__(256, 1)
 static __global__ void flash_attn_mask_to_sparse_indices(
         const half * mask_ptr, int32_t * indices_ptr, const int ne30, const int n_kv_max,
         const int64_t s31, const int64_t s33) {
+    ggml_cuda_pdl_sync();
+
     constexpr int values_per_lane = 8;
     const int tid      = threadIdx.x;
     const int warp     = tid / WARP_SIZE;
@@ -56,7 +58,7 @@ static __global__ void flash_attn_mask_to_sparse_indices(
         }
         __syncthreads();
 
-        const uint32_t lane_mask = lane == 0 ? 0 : (uint32_t(1) << lane) - 1;
+        const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1;
         int warp_item_offset = 0;
 #pragma unroll
         for (int item = 0; item < values_per_lane; ++item) {
@@ -75,13 +77,17 @@ static __global__ void flash_attn_mask_to_sparse_indices(
         __syncthreads();
     }
 
+    ggml_cuda_pdl_lc();
+
     const int count = row_count;
     for (int i = count + tid; i < n_kv_max; i += blockDim.x) {
         indices[i] = -1;
     }
-    if (tid == 0 && count > n_kv_max) {
-        printf("flash attention sparse mask row exceeds n_kv_max (%d > %d)\n", count, n_kv_max);
-        __trap();
+    if (count > n_kv_max) {
+        if (tid == 0) {
+            printf("flash attention sparse mask row exceeds n_kv_max (%d > %d)\n", count, n_kv_max);
+            __trap();
+        }
     }
 }
 #endif
diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp
index abf6a0502f..2074b2f14f 100644
--- a/src/models/qwen4exp.cpp
+++ b/src/models/qwen4exp.cpp
@@ -686,7 +686,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa(
     ggml_tensor * k = mctx_cur->get_k(ctx0, il);
     ggml_tensor * v = mctx_cur->get_v(ctx0, il);
 
-    ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il);
+    ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, 0, kq_scale, il);
     cb(cur, "kqv_out", il);
 
     // the rotation is its own inverse, so undo it on the value side of the output