furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 02:05:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

CUDA: MLA flash-attention decode on Pascal (vec_f32 K=576/V=512), incl. Q8_0 KV (#2079)

Browse Source 
On GPUs without FP16 tensor cores (Pascal / sm_60, e.g. Tesla P100) MLA
flash-attention decode falls back to the CPU. The !fp16_mma_available path
routes decode to the f16 vector kernel, whose is_supported check requires
K == V head sizes; MLA's absorbed head sizes are 576/512 (asymmetric), so it is
rejected and attention runs on the CPU. With --cpu-moe that recomputes the full
MLA attention on the CPU every decoded token, which dominates decode at long
context.

Route Pascal MLA decode (Q->ne[1] <= 8 && K == 576 && V == 512) to the f32
vector kernel and enable that kernel for the 576/512 case, including Q8_0 KV.

Scope: decode only (batch <= 8). Prefill (batch > 8) and -fa 0 are untouched;
tensor-core GPUs never reach this branch. Aligned head sizes are byte-identical
(the asymmetric/Q8_0 work folds to a no-op at compile time), so no other model
or configuration is affected.

- fattn.cu: route 576/512 decode to vec_f32 in the !fp16_mma dispatch and its
  is_supported mirror.
- fattn-vec-f32.cu/.cuh: accept + instantiate 576/512 (F16 and Q8_0); fix latent
  issues exposed by the first asymmetric/large-head use (KQ-row granularity uses
  FATTN_KQ_STRIDE not Dv; guard the dst store to tid < Dv; guard the softmax exp
  on the KV tail; size Q_i32 by ceil; only convert K/V to F16 when the type is
  F16). All are no-ops for the previously-exercised symmetric cases.
- fattn-vec-f32.cuh / fattn-vec-common.cuh: guard the Q8_0 ragged tail (Dk=576 is
  144 int32 lanes = 4.5 warps) with the ragged-dim idiom; compile-time-constant
  for aligned head dims, so it folds away.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
mb8565
2026-07-06 17:58:50 +02:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 2992537ca3
commit 72201359dd
6 changed files with 112 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/fattn-vec-common.cuh
+32 -16
View File
@@ -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<T, half>::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<T, half>::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<T, 1>(&v, &Q_q8[k_KQ_0/WARP_SIZE], K_q8_0[ib].d, Q_d);
}
sum += vec_dot_q8_0_q8_1_impl<T, 1>(&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 <typename Tds>
// 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<Tds, ni> (lane guards :338 / :366).
template <typename Tds, int ni = WARP_SIZE>
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<Tds, half2>::value) {
((half2 *) yds)[threadIdx.x/QI8_1] = make_half2(d, sum);
} else {
ggml/src/ggml-cuda/fattn-vec-f32.cu
+16 -1
View File
@@ -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
ggml/src/ggml-cuda/fattn-vec-f32.cuh
+39 -10
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/(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<float2>(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<float2, WARP_SIZE>(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<float2, ni>(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 <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 && 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<Dv, cols_per_block, 1>(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<Dv, cols_per_block, 1>(ctx, dst, fattn_kernel, nwarps, nbytes_shared, kq_row_granularity, need_f16_K, need_f16_V);
}
template <int Dk, int Dv, ggml_type type_K, ggml_type type_V>
@@ -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);
ggml/src/ggml-cuda/fattn.cu
+15
View File
@@ -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);
ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-f16-f16.cu
+5
View File
@@ -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);
ggml/src/ggml-cuda/template-instances/fattn-vec-f32-instance-hs576-q8_0-q8_0.cu
+5
View File
@@ -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);