From 77f8060ca6c83e4cca7e1c23932f3353c93ecd68 Mon Sep 17 00:00:00 2001 From: Max Homilius Date: Fri, 20 Mar 2026 04:39:39 -0400 Subject: [PATCH] fix: FA vec kernels for D=256 quantized KV cache (e.g., Qwen3-Coder-Next) (#1452) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: FA vec kernels for D=256 quantized KV cache Three bugs prevented Flash Attention from working with quantized (q8_0) KV cache on models with head dimension 256, such as Qwen3-Coder-Next: 1. need_f16_K/V used dimension-based logic (Dk != 128) that forced q8_0→f16 conversion at D=256. The kernel then read f16 data as q8_0 blocks. Changed to type-based logic (type_K == GGML_TYPE_F16) matching mainline. 2. quantize_q8_1_to_shared output pointers were not advanced between loop iterations, so the second iteration overwrote the first half of Q data in shared memory. Added i0 offset to output pointers. 3. Q_i32 register array in vec_f32 kernel was sized to 1 instead of 2 for D=256 due to an errant comparison (Dk >= expr evaluates to boolean 1, not expr). Removed the comparison. Also adds extern template declarations for (256, Q8_0, Q8_0) so the kernel is compiled and linked. Tested on Pascal (GTX 1080, CC 6.1) with Qwen3-Coder-Next (D=256, n_head_kv=2) using --cache-type-k q8_0 --cache-type-v q8_0. * fix: use dimension-based need_f16 for D=256 vec kernels --- ggml/src/ggml-cuda/fattn-vec-f16.cuh | 9 +++++---- ggml/src/ggml-cuda/fattn-vec-f32.cuh | 11 ++++++----- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-vec-f16.cuh b/ggml/src/ggml-cuda/fattn-vec-f16.cuh index 2eb92b19..27e2ecb6 100644 --- a/ggml/src/ggml-cuda/fattn-vec-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-vec-f16.cuh @@ -146,7 +146,7 @@ static __global__ void flash_attn_vec_ext_f16( const float * Q_f = (const float *) (Q + j*nb01); #pragma unroll for (int i0 = 0; i0 < Dk/sizeof(int); i0 += WARP_SIZE) { - quantize_q8_1_to_shared(Q_f + 4*i0, scale, tmp_q_i32, tmp_q_ds); + quantize_q8_1_to_shared(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1); } } @@ -368,8 +368,8 @@ template ; - constexpr bool need_f16_K = Dk != 128 && Dk != 192; - constexpr bool need_f16_V = Dv != 128 && Dv != 64; + constexpr bool need_f16_K = Dk != 128 && Dk != 192 && Dk != 256; + constexpr bool need_f16_V = Dv != 64 && Dv != 128 && Dv != 256; constexpr size_t nbytes_shared = 0; launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, Dv, need_f16_K, need_f16_V); } @@ -502,7 +502,8 @@ extern DECL_FATTN_VEC_F16_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_F16); extern DECL_FATTN_VEC_F16_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_F16); extern DECL_FATTN_VEC_F16_CASE(128, GGML_TYPE_F16, GGML_TYPE_F16); -extern DECL_FATTN_VEC_F16_CASE(256, GGML_TYPE_F16, GGML_TYPE_F16); +extern DECL_FATTN_VEC_F16_CASE(256, GGML_TYPE_F16, GGML_TYPE_F16); +extern DECL_FATTN_VEC_F16_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); extern DECL_FATTN_VEC_F16_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16); extern DECL_FATTN_VEC_F16_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/fattn-vec-f32.cuh b/ggml/src/ggml-cuda/fattn-vec-f32.cuh index db1dc132..af8b4698 100644 --- a/ggml/src/ggml-cuda/fattn-vec-f32.cuh +++ b/ggml/src/ggml-cuda/fattn-vec-f32.cuh @@ -115,7 +115,7 @@ static __global__ void flash_attn_vec_ext_f32( // Convert Q to float2 (f16 K) or q8_1 (quantized K) and store in registers: float2 Q_f2[ncols][Dk/(2*WARP_SIZE)]; - int Q_i32[ncols][Dk/(sizeof(int)*QK8_1) == 0 ? 1 : Dk >= Dk/(sizeof(int)*QK8_1)]; + int Q_i32[ncols][Dk/(sizeof(int)*QK8_1) == 0 ? 1 : Dk/(sizeof(int)*QK8_1)]; float2 Q_ds[ncols][Dk/QK8_1 == 0 ? 1 : Dk/QK8_1]; if (Q_q8_1) { #pragma unroll @@ -147,7 +147,7 @@ static __global__ void flash_attn_vec_ext_f32( const float * Q_f = (const float *) (Q + j*nb01); #pragma unroll for (int i0 = 0; i0 < Dk/sizeof(int); i0 += WARP_SIZE) { - quantize_q8_1_to_shared(Q_f + 4*i0, scale, tmp_q_i32, tmp_q_ds); + quantize_q8_1_to_shared(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1); } } @@ -353,8 +353,8 @@ template ; - constexpr bool need_f16_K = Dk != 128; - constexpr bool need_f16_V = Dv != 128 && Dv != 64; + constexpr bool need_f16_K = Dk != 128 && Dk != 256; + constexpr bool need_f16_V = Dv != 64 && Dv != 128 && Dv != 256; constexpr size_t nbytes_shared = 0; launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, Dv, need_f16_K, need_f16_V); } @@ -477,7 +477,8 @@ extern DECL_FATTN_VEC_F32_CASE(128, GGML_TYPE_Q5_1, GGML_TYPE_F16); extern DECL_FATTN_VEC_F32_CASE(128, GGML_TYPE_Q8_0, GGML_TYPE_F16); extern DECL_FATTN_VEC_F32_CASE(128, GGML_TYPE_F16, GGML_TYPE_F16); -extern DECL_FATTN_VEC_F32_CASE(256, GGML_TYPE_F16, GGML_TYPE_F16); +extern DECL_FATTN_VEC_F32_CASE(256, GGML_TYPE_F16, GGML_TYPE_F16); +extern DECL_FATTN_VEC_F32_CASE(256, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); extern DECL_FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16); extern DECL_FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0);