diff --git a/ggml/src/ggml-cuda/mt_pagedattn.cu b/ggml/src/ggml-cuda/mt_pagedattn.cu index 61752e67cf29..9257824a2213 100644 --- a/ggml/src/ggml-cuda/mt_pagedattn.cu +++ b/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); diff --git a/ggml/src/ggml-cuda/mt_pagedattn_aiter.cu b/ggml/src/ggml-cuda/mt_pagedattn_aiter.cu index 70b6e1980fcc..7a90c0b3cbf2 100644 --- a/ggml/src/ggml-cuda/mt_pagedattn_aiter.cu +++ b/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); } diff --git a/ggml/src/ggml-cuda/mt_pagedattn_decode.cu b/ggml/src/ggml-cuda/mt_pagedattn_decode.cu index 1eb7cac8e73e..2b876dc65038 100644 --- a/ggml/src/ggml-cuda/mt_pagedattn_decode.cu +++ b/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 __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 __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; diff --git a/ggml/src/ggml-cuda/mt_pagedattn_tile.cu b/ggml/src/ggml-cuda/mt_pagedattn_tile.cu index 3e6f512910c5..6c41d2744b5a 100644 --- a/ggml/src/ggml-cuda/mt_pagedattn_tile.cu +++ b/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; diff --git a/ggml/src/ggml-cuda/set-rows.cu b/ggml/src/ggml-cuda/set-rows.cu index 6e67a6ce6a08..f3a789db60ee 100644 --- a/ggml/src/ggml-cuda/set-rows.cu +++ b/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; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 159fcc7cd437..c3c4a91920c7 100644 --- a/src/llama-graph.cpp +++ b/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 == + // , 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]; diff --git a/src/llama-kv-cache-paged.cpp b/src/llama-kv-cache-paged.cpp index dcff541fda75..450c3075f098 100644 --- a/src/llama-kv-cache-paged.cpp +++ b/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");