diff --git a/ggml/src/ggml-cuda/fattn-vec-common.cuh b/ggml/src/ggml-cuda/fattn-vec-common.cuh index c6dda7ea1..3cc52ce28 100644 --- a/ggml/src/ggml-cuda/fattn-vec-common.cuh +++ b/ggml/src/ggml-cuda/fattn-vec-common.cuh @@ -326,21 +326,28 @@ static __device__ __forceinline__ T vec_dot_fattn_vec_KQ_q8_0( for (int k_KQ_0 = 0; k_KQ_0 < Dk/sizeof(int); k_KQ_0 += WARP_SIZE) { const int k_KQ = k_KQ_0 + threadIdx.x; - const int ib = k_KQ / QI8_0; - const int iqs = k_KQ % QI8_0; + // Ragged-tail guard (mirrors mainline llama.cpp fattn-vec.cuh:169). For a chunk that fully + // fits (k_KQ_0 + WARP_SIZE <= Dk/sizeof(int)) the predicate is a compile-time constant true + // and folds away -> byte-identical for aligned Dk (128/256). For the ragged tail of Dk=576 + // (144 int32 = 4.5*WARP_SIZE) the out-of-bounds lanes (k_KQ >= 144) skip the K-block read + // (ib would be 18/19 > last valid block 17 -> OOB -> NaN) and contribute 0 to the warp sum. + if (k_KQ_0 + WARP_SIZE <= Dk/(int)sizeof(int) || k_KQ < Dk/(int)sizeof(int)) { + const int ib = k_KQ / QI8_0; + const int iqs = k_KQ % QI8_0; - const int v = get_int_b2(K_q8_0[ib].qs, iqs); + const int v = get_int_b2(K_q8_0[ib].qs, iqs); - T Q_d; - if (std::is_same::value) { - const half2 * Q_ds = (const half2 *) Q_ds_v; - Q_d = __low2half(Q_ds[k_KQ_0/WARP_SIZE]); - } else { - const float2 * Q_ds = (const float2 *) Q_ds_v; - Q_d = Q_ds[k_KQ_0/WARP_SIZE].x; + T Q_d; + if (std::is_same::value) { + const half2 * Q_ds = (const half2 *) Q_ds_v; + Q_d = __low2half(Q_ds[k_KQ_0/WARP_SIZE]); + } else { + const float2 * Q_ds = (const float2 *) Q_ds_v; + Q_d = Q_ds[k_KQ_0/WARP_SIZE].x; + } + + sum += vec_dot_q8_0_q8_1_impl(&v, &Q_q8[k_KQ_0/WARP_SIZE], K_q8_0[ib].d, Q_d); } - - sum += vec_dot_q8_0_q8_1_impl(&v, &Q_q8[k_KQ_0/WARP_SIZE], K_q8_0[ib].d, Q_d); } return sum; @@ -388,14 +395,21 @@ static __device__ __forceinline__ T vec_dot_fattn_vec_KQ_f16( return sum; } -template +// ni = number of valid lanes for this chunk (compile-time). For a chunk that fully covers +// WARP_SIZE int32 words ni == WARP_SIZE and every guard below is provably always-true, so the +// compiler eliminates it (byte-identical to the unguarded version). For a ragged tail chunk +// (e.g. Dk=576 -> the i0=128 chunk has only 144-128=16 valid int32 words) ni < WARP_SIZE and the +// guards engage: the out-of-bounds lanes read 0 (no OOB Q load), and skip the yq32 / yds writes +// (no clobber of the trailing q8_1 scale region / no OOB scale write). Mirrors mainline +// llama.cpp fattn-common.cuh quantize_q8_1_to_shared (lane guards :338 / :366). +template static __device__ __forceinline__ void quantize_q8_1_to_shared( const float * __restrict__ x, const float scale, int * __restrict__ yq32, void * __restrict__ yds) { float vals[sizeof(int)] = {0.0f}; #pragma unroll for (int l = 0; l < sizeof(int); ++l) { - vals[l] = scale * x[4*threadIdx.x + l]; + vals[l] = (ni == WARP_SIZE || threadIdx.x < ni) ? scale * x[4*threadIdx.x + l] : 0.0f; } float amax = fabsf(vals[0]); @@ -422,8 +436,10 @@ static __device__ __forceinline__ void quantize_q8_1_to_shared( } } - yq32[threadIdx.x] = q32; - if (threadIdx.x % QI8_1 == 0) { + if (ni == WARP_SIZE || threadIdx.x < ni) { + yq32[threadIdx.x] = q32; + } + if (threadIdx.x % QI8_1 == 0 && (ni == WARP_SIZE || threadIdx.x < ni)) { if (std::is_same::value) { ((half2 *) yds)[threadIdx.x/QI8_1] = make_half2(d, sum); } else { diff --git a/ggml/src/ggml-cuda/fattn-vec-f32.cu b/ggml/src/ggml-cuda/fattn-vec-f32.cu index c1a5ad2c9..cd018c2bc 100644 --- a/ggml/src/ggml-cuda/fattn-vec-f32.cu +++ b/ggml/src/ggml-cuda/fattn-vec-f32.cu @@ -79,6 +79,9 @@ void ggml_cuda_flash_attn_ext_vec_f32(ggml_backend_cuda_context & ctx, ggml_tens FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16) FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0) + + FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16) + FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0) #else FATTN_VEC_F32_CASE(128, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0) @@ -97,6 +100,9 @@ void ggml_cuda_flash_attn_ext_vec_f32(ggml_backend_cuda_context & ctx, ggml_tens FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_F16, GGML_TYPE_F16) FATTN_VEC_F32_CASE_DKDV(192, 128, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0) + FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16) + FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0) + #endif // GGML_CUDA_FA_ALL_QUANTS on_no_fattn_vec_case(Q->ne[0], V->ne[0]); @@ -106,8 +112,17 @@ bool ggml_cuda_fattn_vec_f32_is_supported([[maybe_unused]] ggml_backend_cuda_con auto K = dst->src[1]; auto V = dst->src[2]; if (K->ne[0] != V->ne[0]) { - if (K->ne[0] != 192 || V->ne[2] != 128) return false; + const bool is_192_128 = K->ne[0] == 192 && V->ne[0] == 128; + const bool is_576_512 = K->ne[0] == 576 && V->ne[0] == 512; + if (!is_192_128 && !is_576_512) return false; if (K->type != V->type) return false; + // Asymmetric (Dk != Dv) vec_f32 path. F16 KV is always fine (its dot loops stride Dk/2, + // a clean multiple of 32). Q8_0 KV is now supported: the per-warp Q->q8_1 quantization, + // the shared->register load, and the quantized-K dot all carry a compile-time-gated + // per-lane bounds guard (mirroring mainline llama.cpp's fattn-vec ragged-tail idiom) so + // the ragged final warp iteration for Dk=576 (144 int32 = 4.5*WARP_SIZE) and Dk=192 + // (48 int32 = 1.5*WARP_SIZE) no longer reads Q/K out of bounds. The guard folds to a + // no-op for aligned Dk (128/256), keeping those instances byte-identical. return K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0; } #ifdef GGML_CUDA_FA_ALL_QUANTS diff --git a/ggml/src/ggml-cuda/fattn-vec-f32.cuh b/ggml/src/ggml-cuda/fattn-vec-f32.cuh index fd3b6fa45..d07ba359d 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/(sizeof(int)*QK8_1)]; + int Q_i32[ncols][(Dk/(int)sizeof(int) + WARP_SIZE - 1)/WARP_SIZE]; float2 Q_ds[ncols][Dk/QK8_1 == 0 ? 1 : Dk/QK8_1]; if (Q_q8_1) { #pragma unroll @@ -145,9 +145,21 @@ static __global__ void flash_attn_vec_ext_f32( } const float * Q_f = (const float *) (Q + j*nb01); + // Full WARP_SIZE-wide int32 chunks: ni == WARP_SIZE so quantize_q8_1_to_shared's lane + // guards fold away -> byte-identical to the original unguarded loop for aligned Dk. + constexpr int Dk_i32 = Dk/(int)sizeof(int); + constexpr int Dk_full = (Dk_i32/WARP_SIZE)*WARP_SIZE; #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 + i0, tmp_q_ds + i0/QI8_1); + for (int i0 = 0; i0 < Dk_full; i0 += WARP_SIZE) { + quantize_q8_1_to_shared(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1); + } + // Ragged tail: only compiled for Dk whose int32 count is not a multiple of WARP_SIZE + // (Dk=576 -> 16 valid lanes, Dk=192 -> 16 valid lanes). `if constexpr` drops this + // entirely for aligned Dk (128/256), so the aligned path stays byte-identical. + if constexpr (Dk_i32 % WARP_SIZE != 0) { + constexpr int i0 = Dk_full; + constexpr int ni = Dk_i32 - Dk_full; + quantize_q8_1_to_shared(Q_f + 4*i0, scale, tmp_q_i32 + i0, tmp_q_ds + i0/QI8_1); } } @@ -162,8 +174,17 @@ static __global__ void flash_attn_vec_ext_f32( for (int i0 = 0; i0 < Dk/sizeof(int); i0 += WARP_SIZE) { const int i = i0 + threadIdx.x; - Q_i32[j][i0/WARP_SIZE] = tmp_q_i32[i]; - Q_ds[j][i0/WARP_SIZE] = tmp_q_ds[i/QI8_1]; + // Ragged-tail guard (mirrors mainline fattn-vec.cuh:169). Aligned Dk: predicate + // folds to always-true -> byte-identical. Dk=576 ragged tail: OOB lanes (i >= 144) + // load 0 rather than reading the trailing shared region (tmp_q_i32[144..159] is now + // never written by the guarded quantize, so it is stale) and feed 0 to the dot. + if (i0 + WARP_SIZE <= Dk/(int)sizeof(int) || i < Dk/(int)sizeof(int)) { + Q_i32[j][i0/WARP_SIZE] = tmp_q_i32[i]; + Q_ds[j][i0/WARP_SIZE] = tmp_q_ds[i/QI8_1]; + } else { + Q_i32[j][i0/WARP_SIZE] = 0; + Q_ds[j][i0/WARP_SIZE] = make_float2(0.0f, 0.0f); + } } } @@ -257,7 +278,7 @@ static __global__ void flash_attn_vec_ext_f32( const float KQ_max_scale = expf(kqmax[j] - kqmax_new_j); kqmax[j] = kqmax_new_j; - const float val = expf(KQ[j*Dk + tid] - kqmax[j]); + const float val = (FATTN_KQ_STRIDE % Dk != 0 && k_VKQ_0 + tid >= ne11) ? 0.0f : expf(KQ[j*Dk + tid] - kqmax[j]); kqsum[j] = kqsum[j]*KQ_max_scale + val; KQ[j*Dk + tid] = val; @@ -344,7 +365,9 @@ static __global__ void flash_attn_vec_ext_f32( if (gridDim.y == 1) { dst_val /= kqsum[j_VKQ]; } - dst[(((sequence*ne01 + ic0 + j_VKQ)*ne02 + head)*gridDim.y + blockIdx.y)*Dv + tid] = dst_val; + if (tid < Dv) { + dst[(((sequence*ne01 + ic0 + j_VKQ)*ne02 + head)*gridDim.y + blockIdx.y)*Dv + tid] = dst_val; + } } if (gridDim.y != 1 && tid < ncols && (ncols <= 2 || ic0 + tid < ne01)) { @@ -359,10 +382,13 @@ template ; - constexpr bool need_f16_K = Dk != 128 && Dk != 256; - constexpr bool need_f16_V = Dv != 64 && Dv != 128 && Dv != 256; + constexpr bool need_f16_K = (Dk != 128 && Dk != 256) && type_K == GGML_TYPE_F16; + constexpr bool need_f16_V = (Dv != 64 && Dv != 128 && Dv != 256) && type_V == GGML_TYPE_F16; constexpr size_t nbytes_shared = 0; - launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, Dv, need_f16_K, need_f16_V); + // MLA (Dv=512): the KV cache is padded to FATTN_KQ_STRIDE (256), not to Dv, so use 256 as the KQ-row + // granularity. Passing Dv=512 asserts n_kv % 512 == 0, which fails for odd multiples of 256. + constexpr int kq_row_granularity = (Dv == 512) ? FATTN_KQ_STRIDE : Dv; + launch_fattn(ctx, dst, fattn_kernel, nwarps, nbytes_shared, kq_row_granularity, need_f16_K, need_f16_V); } template @@ -488,3 +514,6 @@ 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); + +extern DECL_FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16); +extern DECL_FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0); diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index a9de08224..02b35c4d0 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -23,6 +23,14 @@ static inline bool mma_better_than_turing(const int cc) { return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) > CC_TURING; } +// On GPUs without fp16 MMA (Pascal/sm_60), the MLA absorbed head sizes 576/512 are asymmetric, so the f16 +// vector kernel (which requires K==V) rejects them and decode falls back to the CPU. Route MLA decode +// (batch <= 8) to the f32 vector kernel to keep it on the GPU. Single-sourced so the dispatch and the +// is_supported check cannot drift. +static inline bool is_pascal_mla_absorbed_decode(const ggml_tensor * Q, const ggml_tensor * K, const ggml_tensor * V) { + return Q->ne[1] <= 8 && K->ne[0] == 576 && V->ne[0] == 512; +} + void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; const ggml_tensor * Q = dst->src[0]; @@ -76,6 +84,10 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst } if (!fp16_mma_available(cc)) { + if (is_pascal_mla_absorbed_decode(Q, K, V)) { + ggml_cuda_flash_attn_ext_vec_f32(ctx, dst); + return; + } if (precision == GGML_PREC_DEFAULT) { if (Q->ne[1] <= 8 || Q->ne[0] == 256) { ggml_cuda_flash_attn_ext_vec_f16(ctx, dst); @@ -171,6 +183,9 @@ bool ggml_cuda_fattn_is_supported(ggml_backend_cuda_context & ctx, const ggml_te } if (!fp16_mma_available(cc)) { + if (is_pascal_mla_absorbed_decode(Q, K, V)) { + return ggml_cuda_fattn_vec_f32_is_supported(ctx, dst); + } if (precision == GGML_PREC_DEFAULT) { if (Q->ne[1] <= 8 || Q->ne[0] == 256) { return ggml_cuda_fattn_vec_f16_is_supported(ctx, dst); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-f16-f16.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-f16-f16.cu new file mode 100644 index 000000000..696c34283 --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-f16-f16.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec-f32.cuh" + +DECL_FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_F16, GGML_TYPE_F16); diff --git a/ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-q8_0-q8_0.cu b/ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-q8_0-q8_0.cu new file mode 100644 index 000000000..0145e213b --- /dev/null +++ b/ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-q8_0-q8_0.cu @@ -0,0 +1,5 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#include "../fattn-vec-f32.cuh" + +DECL_FATTN_VEC_F32_CASE_DKDV(576, 512, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0);