Commit 7f24aacde for llama.cpp
commit 7f24aacde32af15559747050b7073e7dc34fe02e
Author: Pascal <admin@serveurperso.com>
Date: Wed Nov 11 11:55:19 2026 +0100
cpu: accumulate flash attention V in F32 (#29810)
* cpu: accumulate flash attention V in F32 in the one-chunk kernel
The one-chunk kernel accumulated an F16 V cache in an F16 running sum
that is only normalized at the end, which overflows past 65504 on long
contexts: inf, then NaN on every following token.
Every V type now accumulates in F32, like the tiled kernel. F16 rows go
through ggml_vec_mad_f16_f32, which converts and multiplies in the same
register (F16C/AVX512, NEON, RVV) with no round trip through a temporary
F32 row, so decode keeps master speed on x86 and loses about 4% at 4k
context on ARM, where master used native F16 FMA.
* cpu: widen the F16 V rows like ggml_vec_dot_f16_f32
ggml_vec_mad_f16_f32 required F16C and FMA, so AVX builds without FMA
fell back to a scalar loop and ran 3x slower than master. It now
follows the guards of ggml_vec_dot_f16_f32: AVX512, then AVX with F16C
through GGML_F32x8_FMA, then NEON and RVV, then the scalar loop.
diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp
index 00a814051..c987a087c 100644
--- a/ggml/src/ggml-cpu/ops.cpp
+++ b/ggml/src/ggml-cpu/ops.cpp
@@ -8911,14 +8911,9 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
float * VKQ32 = (float *) params->wdata + ith*(1*DK + 2*DV + CACHE_LINE_SIZE_F32); // FP32 VKQ accumulator
float * V32 = (VKQ32 + 1*DV); // (temporary) FP32 V buffer
- ggml_fp16_t * VKQ16 = (ggml_fp16_t *) (VKQ32 + 1*DV); // (temporary) FP16 VKQ accumulator
ggml_fp16_t * Q_q = (ggml_fp16_t *) (VKQ32 + 2*DV); // (temporary) buffer for Q converted to quantized/FP16
- if (v->type == GGML_TYPE_F16) {
- memset(VKQ16, 0, DV*sizeof(ggml_fp16_t));
- } else {
- memset(VKQ32, 0, DV*sizeof(float));
- }
+ memset(VKQ32, 0, DV*sizeof(float));
const ggml_fp16_t * mp = mask ? (ggml_fp16_t *)((char *) mask->data + iq1*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]) : NULL;
@@ -8970,53 +8965,32 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
const char * v_data = ((const char *) v->data + (ik1*nbv1 + iv2*nbv2 + iv3*nbv3));
- if (v->type == GGML_TYPE_F16) {
- if (s > M) {
- // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f
- M = s;
- ms = expf(Mold - M);
-
- // V = V*expf(Mold - M)
- ggml_vec_scale_f16(DV, VKQ16, ms);
- } else {
- // no new maximum, ms == 1.0f, vs != 1.0f
- vs = expf(s - M);
- }
+ if (s > M) {
+ // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f
+ M = s;
+ ms = expf(Mold - M);
- // V += v*expf(s - M)
- ggml_vec_mad_f16(DV, VKQ16, (const ggml_fp16_t *) v_data, vs);
+ // V = V*expf(Mold - M)
+ ggml_vec_scale_f32(DV, VKQ32, ms);
} else {
- if (s > M) {
- // s is new maximum, ms < 1.0f, vs == expf(s - s) == 1.0f
- M = s;
- ms = expf(Mold - M);
-
- // V = V*expf(Mold - M)
- ggml_vec_scale_f32(DV, VKQ32, ms);
- } else {
- // no new maximum, ms == 1.0f, vs != 1.0f
- vs = expf(s - M);
- }
+ // no new maximum, ms == 1.0f, vs != 1.0f
+ vs = expf(s - M);
+ }
- // V += v*expf(s - M)
- if (v_to_float) {
- v_to_float(v_data, V32, DV);
- ggml_vec_mad_f32(DV, VKQ32, V32, vs);
- } else {
- // V is F32
- ggml_vec_mad_f32(DV, VKQ32, (const float *) v_data, vs);
- }
+ // V += v*expf(s - M)
+ if (v->type == GGML_TYPE_F16) {
+ ggml_vec_mad_f16_f32(DV, VKQ32, (const ggml_fp16_t *) v_data, vs);
+ } else if (v_to_float) {
+ v_to_float(v_data, V32, DV);
+ ggml_vec_mad_f32(DV, VKQ32, V32, vs);
+ } else {
+ // V is F32
+ ggml_vec_mad_f32(DV, VKQ32, (const float *) v_data, vs);
}
S = S*ms + vs; // scale and increment sum with partial sum
}
- if (v->type == GGML_TYPE_F16) {
- for (int64_t d = 0; d < DV; ++d) {
- VKQ32[d] = GGML_CPU_FP16_TO_FP32(VKQ16[d]);
- }
- }
-
// sinks - apply only on the first kv-chunk
if (sinks && ic_start == 0) {
const float s = ((float *)((char *) sinks->data))[h];
diff --git a/ggml/src/ggml-cpu/vec.h b/ggml/src/ggml-cpu/vec.h
index 6598016d3..c54ed3f07 100644
--- a/ggml/src/ggml-cpu/vec.h
+++ b/ggml/src/ggml-cpu/vec.h
@@ -436,6 +436,41 @@ inline static void ggml_vec_mad_f32(const int n, float * GGML_RESTRICT y, const
#endif
}
+// y (F32) += x (F16)*v
+// on AVX512, AVX with F16C, NEON and RVV, x is widened to F32 in-register and fused into the multiply-add
+inline static void ggml_vec_mad_f16_f32(const int n, float * GGML_RESTRICT y, const ggml_fp16_t * GGML_RESTRICT x, const float v) {
+ int i = 0;
+#if defined(__AVX512F__)
+ const __m512 vx = _mm512_set1_ps(v);
+ for (; i + 16 <= n; i += 16) {
+ const __m512 ax = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x + i)));
+ _mm512_storeu_ps(y + i, _mm512_fmadd_ps(ax, vx, _mm512_loadu_ps(y + i)));
+ }
+#elif defined(__AVX__) && defined(__F16C__)
+ const __m256 vx = _mm256_set1_ps(v);
+ for (; i + 8 <= n; i += 8) {
+ const __m256 ax = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x + i)));
+ _mm256_storeu_ps(y + i, GGML_F32x8_FMA(_mm256_loadu_ps(y + i), ax, vx));
+ }
+#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_FMA) && defined(__ARM_FP16_FORMAT_IEEE)
+ const float32x4_t vx = vdupq_n_f32(v);
+ for (; i + 8 <= n; i += 8) {
+ const float16x8_t ax = vld1q_f16((const __fp16 *)(x + i));
+ vst1q_f32(y + i + 0, vfmaq_f32(vld1q_f32(y + i + 0), vcvt_f32_f16(vget_low_f16 (ax)), vx));
+ vst1q_f32(y + i + 4, vfmaq_f32(vld1q_f32(y + i + 4), vcvt_f32_f16(vget_high_f16(ax)), vx));
+ }
+#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfhmin)
+ for (int vl; i < n; i += vl) {
+ vl = __riscv_vsetvl_e16m2(n - i);
+ const vfloat32m4_t ax = __riscv_vfwcvt_f_f_v_f32m4(__riscv_vle16_v_f16m2((const _Float16 *)x + i, vl), vl);
+ __riscv_vse32_v_f32m4(y + i, __riscv_vfmacc_vf_f32m4(__riscv_vle32_v_f32m4(y + i, vl), v, ax, vl), vl);
+ }
+#endif
+ for (; i < n; ++i) {
+ y[i] += GGML_CPU_FP16_TO_FP32(x[i])*v;
+ }
+}
+
inline static void ggml_vec_mad_f16(const int n, ggml_fp16_t * GGML_RESTRICT y, const ggml_fp16_t * GGML_RESTRICT x, const float v) {
#if defined(GGML_SIMD) && defined(__ARM_FEATURE_SVE)
const int sve_register_length = svcntb() * 8;