Merge pull request #13 from kmbandy/fix/lfm25-paged-turbo4-gfx803

fix: LFM2.5 paged turbo4 on gfx803 (warp-width, sub-128 head_dim, parallel partials)

Author: kmbandy

ggml/src/ggml-cuda/mt_pagedattn.cu

@@ -421,7 +421,7 @@ __global__ void mt_scatter_kv_turbo4_0_kernel(
     // ---- Step 6: Pack qs (nibble packed, warp-cooperative) ----
     const int      lane            = j % WARP_SIZE;
     const uint8_t  my_nibble       = idx & 0xF;
-    const uint8_t  partner_nibble  = __shfl_sync(0xffffffffu, my_nibble, lane ^ 1);
+    const uint8_t  partner_nibble  = __shfl_sync(0xffffffffu, my_nibble, lane ^ 1, WARP_SIZE);
     if ((j & 1) == 0) {
         blk->qs[j / 2] = my_nibble | (partner_nibble << 4);
     }
@@ -1208,7 +1208,7 @@ void ggml_cuda_op_paged_attn_mt(ggml_backend_cuda_context & ctx, ggml_tensor * d
                         num_seqs, (int) k_cur->ne[2], n_kv_heads, stream);
 
                     const int num_chunks    = paged_attn_decode_num_chunks(max_ctx_len);
-                    const int max_q_len     = avg_q_len;  // see comment above
+                    const int max_q_len     = total_q_tokens;  // partials inner stride: must be >= max per-seq q_len. avg_q_len = total/num_seqs floors to 0 with idle parallel slots (1 active + N idle), collapsing every head/seq/chunk partial offset to 0 -> OOB corruption. total_q_tokens (gate-capped <= 8) is a safe upper bound.
                     const size_t partials_n = (size_t) n_heads * (size_t) num_seqs
                                             * (size_t) num_chunks * (size_t) max_q_len
                                             * (size_t) (HS + 2);

ggml/src/ggml-cuda/mt_pagedattn_aiter.cu

@@ -322,7 +322,7 @@ __global__ void mt_scatter_kv_turbo4_aiter_kernel(
     {
         const int      lane            = j % WARP_SIZE;
         const uint8_t  my_nibble       = idx & 0xF;
-        const uint8_t  partner_nibble  = __shfl_sync(0xffffffffu, my_nibble, lane ^ 1);
+        const uint8_t  partner_nibble  = __shfl_sync(0xffffffffu, my_nibble, lane ^ 1, WARP_SIZE);
         if ((j & 1) == 0) {
             blk->qs[j / 2] = my_nibble | (partner_nibble << 4);
         }

ggml/src/ggml-cuda/mt_pagedattn_decode.cu

@@ -238,7 +238,7 @@ static __device__ __forceinline__ void decode_coop_stage_turbo4(
                 }
             }
         }
-        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0);
+        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0, WARP_SIZE);
 
         uint16_t packed = 0;
         if (blk != nullptr) {
@@ -302,7 +302,7 @@ static __device__ __forceinline__ void decode_coop_stage_turbo3(
                 }
             }
         }
-        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0);
+        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0, WARP_SIZE);
 
         uint8_t qs_byte    = 0;
         uint8_t signs_byte = 0;
@@ -386,7 +386,7 @@ template <typename T>
 __device__ __forceinline__ T decode_warp_reduce_sum(T v) {
     #pragma unroll
     for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
-        v += __shfl_xor_sync(0xFFFFFFFF, v, mask);
+        v += __shfl_xor_sync(0xFFFFFFFF, v, mask, WARP_SIZE);
     }
     return v;
 }
@@ -395,7 +395,7 @@ template <typename T>
 __device__ __forceinline__ T decode_warp_reduce_max(T v) {
     #pragma unroll
     for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
-        v = max(v, __shfl_xor_sync(0xFFFFFFFF, v, mask));
+        v = max(v, __shfl_xor_sync(0xFFFFFFFF, v, mask, WARP_SIZE));
     }
     return v;
 }
@@ -938,7 +938,7 @@ __global__ void mt_paged_attention_decode_kernel_wmma(
             float local_max = -INFINITY;
             #pragma unroll
             for (int l = 0; l < scores.ne; ++l) local_max = max(local_max, scores.x[l]);
-            const float row_max = max(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, 16));
+            const float row_max = max(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, 16, WARP_SIZE));
 
             const float new_max = max(running_max, row_max);
             float rescale = 1.0f;
@@ -959,7 +959,7 @@ __global__ void mt_paged_attention_decode_kernel_wmma(
                 scores.x[l] = e;
                 local_sum  += e;
             }
-            const float row_sum = local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, 16);
+            const float row_sum = local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, 16, WARP_SIZE);
             running_sum += row_sum;
             running_max  = new_max;
 

ggml/src/ggml-cuda/mt_pagedattn_tile.cu

@@ -230,7 +230,7 @@ static __device__ __forceinline__ void coop_stage_turbo4_tile(
         }
 
         // Broadcast norm from lane 0 to all 32 lanes of this warp.
-        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0);
+        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0, WARP_SIZE);
 
         // Each lane reads 2 bytes (4 nibbles = 4 elements). qs is uint8_t[64],
         // 2-byte aligned at qs[2*lane_id] for any lane.
@@ -308,7 +308,7 @@ static __device__ __forceinline__ void coop_stage_turbo3_tile(
         }
 
         // Broadcast norm from lane 0 to all 32 lanes of this warp.
-        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0);
+        norm_f = __shfl_sync(0xFFFFFFFF, norm_f, 0, WARP_SIZE);
 
         // Each lane reads 1 byte of qs (its 4 elements' low2 bits) and 1
         // byte of signs (covers this lane's 4 elements' high1 bits, plus
@@ -476,7 +476,7 @@ __global__ void mt_paged_attention_tile_kernel(
         for (int l = 0; l < scores.ne; ++l) {
             local_max = max(local_max, scores.x[l]);
         }
-        const float row_max = max(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, 16));
+        const float row_max = max(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, 16, WARP_SIZE));
 
         const float new_max = max(running_max, row_max);
 
@@ -502,7 +502,7 @@ __global__ void mt_paged_attention_tile_kernel(
             scores.x[l]   = e;
             local_sum   += e;
         }
-        const float row_sum = local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, 16);
+        const float row_sum = local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, 16, WARP_SIZE);
         running_sum += row_sum;
         running_max  = new_max;
 
@@ -861,7 +861,7 @@ __global__ void mt_paged_attention_tile_mw_kernel(
             for (int l = 0; l < scores.ne; ++l) {
                 local_max = max(local_max, scores.x[l]);
             }
-            const float row_max = max(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, 16));
+            const float row_max = max(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, 16, WARP_SIZE));
 
             const float new_max = max(running_max, row_max);
 
@@ -887,7 +887,7 @@ __global__ void mt_paged_attention_tile_mw_kernel(
                 scores.x[l]   = e;
                 local_sum   += e;
             }
-            const float row_sum = local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, 16);
+            const float row_sum = local_sum + __shfl_xor_sync(0xFFFFFFFF, local_sum, 16, WARP_SIZE);
             running_sum += row_sum;
             running_max  = new_max;
 

ggml/src/ggml-cuda/set-rows.cu

@@ -1068,7 +1068,7 @@ static __global__ void k_set_rows_turbo4(
     const uint8_t my_nibble = idx & 0xF;
     uint8_t qs_byte = 0;
     // Gather nibble from partner thread
-    uint8_t partner_nibble = __shfl_sync(0xffffffff, my_nibble, lane ^ 1);
+    uint8_t partner_nibble = __shfl_sync(0xffffffff, my_nibble, lane ^ 1, WARP_SIZE);
     if (j % 2 == 0) {
         qs_byte = my_nibble | (partner_nibble << 4);
         blk->qs[j / 2] = qs_byte;

src/llama-graph.cpp

@@ -2659,9 +2659,27 @@ ggml_tensor * llm_graph_context::build_attn(
             }
             return x;
         };
-        ggml_tensor * q_cast = to_f16_cont(q_cur);
-        ggml_tensor * k_cast = to_f16_cont(k_cur);
-        ggml_tensor * v_cast = to_f16_cont(v_cur);
+        // Turbo paged kernel quantizes 128-element blocks (QK_TURBO=128); a
+        // head_dim < 128 (e.g. LFM2.5 head_dim 64) is smaller than one block,
+        // which the kernel rejects. Zero-pad each head to 128 so the proven
+        // HS=128 path runs: the turbo WHT identity makes <WHT(Qp),WHT(Kp)> ==
+        // <Q,K>, and the padded V dims contribute zero and are sliced off the
+        // output below. Mirrors the non-paged turbo path. ggml_pad needs F32,
+        // so pad before the F16 cast.
+        const bool paged_turbo_pad =
+            (layer.k->type == GGML_TYPE_TURBO2_0 ||
+             layer.k->type == GGML_TYPE_TURBO3_0 ||
+             layer.k->type == GGML_TYPE_TURBO4_0);
+        const int64_t paged_orig_head = q_cur->ne[0];
+        const bool    paged_pad_head  = paged_turbo_pad && (paged_orig_head % 128 != 0);
+        auto pad_head_to_128 = [&](ggml_tensor * t) -> ggml_tensor * {
+            const int64_t pad = ((paged_orig_head + 127) / 128) * 128 - paged_orig_head;
+            ggml_tensor * x = (t->type == GGML_TYPE_F32) ? t : ggml_cast(ctx0, t, GGML_TYPE_F32);
+            return ggml_pad(ctx0, x, pad, 0, 0, 0);
+        };
+        ggml_tensor * q_cast = to_f16_cont(paged_pad_head ? pad_head_to_128(q_cur) : q_cur);
+        ggml_tensor * k_cast = to_f16_cont(paged_pad_head ? pad_head_to_128(k_cur) : k_cur);
+        ggml_tensor * v_cast = to_f16_cont(paged_pad_head ? pad_head_to_128(v_cur) : v_cur);
 
         // 2) Forward-expand the source nodes so they're fully computed
         //    before paged_attn reads them.
@@ -2692,6 +2710,14 @@ ggml_tensor * llm_graph_context::build_attn(
         //    typically wired only for F32 activations.
         cur = ggml_cast(ctx0, cur, GGML_TYPE_F32);
 
+        // Slice the turbo head padding back off (see paged_pad_head above):
+        // output mirrors q as [padded_head, n_heads, n_tokens]; keep the first
+        // paged_orig_head dims per head.
+        if (paged_pad_head) {
+            cur = ggml_cont(ctx0, ggml_view_3d(ctx0, cur, paged_orig_head, cur->ne[1], cur->ne[2],
+                                               cur->nb[1], cur->nb[2], 0));
+        }
+
         // 6) Reshape to (head_dim*n_heads, n_tokens) for the wo projection.
         const int64_t n_embd_full = q_cur->ne[0] * q_cur->ne[1];
         const int64_t n_tokens    = q_cur->ne[2];

src/llama-kv-cache-paged.cpp

@@ -120,6 +120,15 @@ llama_kv_cache_paged::llama_kv_cache_paged(
         }
     }
 
+    // Turbo cache quantizes 128-element blocks; pad a sub-128 head_dim up to
+    // 128 so each head is exactly one turbo block (matches the graph-level
+    // padding in llm_graph_context::build_attn for the paged path).
+    const bool paged_cache_is_turbo =
+        (type_k == GGML_TYPE_TURBO2_0 || type_k == GGML_TYPE_TURBO3_0 || type_k == GGML_TYPE_TURBO4_0);
+    if (paged_cache_is_turbo && head_dim % 128 != 0) {
+        head_dim = ((head_dim + 127) / 128) * 128;
+    }
+
     GGML_ASSERT(head_dim > 0 && "head_dim must be > 0");
     GGML_ASSERT(n_kv_heads > 0 && "n_kv_heads must be > 0");