nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: FA vec kernels for D=256 quantized KV cache (e.g., Qwen3-Coder-Next) (#1452) 

* 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
This commit is contained in:
Max Homilius 2026-03-20 04:39:39 -04:00 committed by GitHub
parent 0871ab2964
commit 77f8060ca6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
ggml/src/ggml-cuda/fattn-vec-f16.cuh
View File

@ -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<half2>(Q_f + 4*i0, scale, tmp_q_i32, tmp_q_ds);
quantize_q8_1_to_shared<half2>(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1);
}
}
@ -368,8 +368,8 @@ template <int Dk, int Dv, int cols_per_block, ggml_type type_K, ggml_type type_V
void ggml_cuda_flash_attn_ext_vec_f16_case_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
constexpr int nwarps = Dk/WARP_SIZE;
fattn_kernel_t fattn_kernel = flash_attn_vec_ext_f16<Dk, Dv, cols_per_block, type_K, type_V, use_logit_softcap>;
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<Dv, cols_per_block, 1>(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);

11
ggml/src/ggml-cuda/fattn-vec-f32.cuh
View File

@ -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<float2>(Q_f + 4*i0, scale, tmp_q_i32, tmp_q_ds);
quantize_q8_1_to_shared<float2>(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1);
}
}
@ -353,8 +353,8 @@ template <int Dk, int Dv, int cols_per_block, ggml_type type_K, ggml_type type_V
void ggml_cuda_flash_attn_ext_vec_f32_case_impl(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
constexpr int nwarps = Dk/WARP_SIZE;
fattn_kernel_t fattn_kernel = flash_attn_vec_ext_f32<Dk, Dv, cols_per_block, type_K, type_V, use_logit_softcap>;
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<Dv, cols_per_block, 1>(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);