ikawrakow/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-02-23 22:54:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fused soft cap and SIMD-ified GeLU (#9)

Browse Source 
* Softcap: WIP

Fuses scale + tanh + scale as used for softcaping in some
models.

Just CPU for now. ~1.4% for PP-512 on Gemma2-9b, no effect on TG.

Somewhat surprisingly the improvement does not increase as I
go to longer contexts. Gemma2 does softcap on K*Q, which grows
quadratically with context length, so I would have thought
the benefit from fusing scale, tanh, scale would increase.
But no, no luck.

* softcap: CUDA

* softcap: CUDA

~1% speedup for Gemma2-9b

* softcap: Metal and NEON

About 1% speedup.

* Simdified gelu

Gives ~1% speedup for Gemma2-9b prompt processing on AVX512/AVX2.
It looks like the gelu operation is memory bound on my CPU's
after SIMD-ifying it. By not using the 128 kb gelu lookup table
we gain a small advantage.
On the M2-Max the lookup table is slightly faster than the SIMD
version, so left the lookup table for ARM_NEON.

* softcap, tanh: avoid NaNs for large arguments (AVX2, AVX512)

Not that I have encountered this in practice, but just to be sure.
This does it for AVX512 and AVX2, still need a guard for ARM_NEON.

* llama-bench: add ability to turn off warmup runs

So we don't need to wait forever on, e.g., benchmarks involving
long contexts.

* softcap, tanh: avoid NaNs for large arguments (NEON)

---------

Co-authored-by: Iwan Kawrakow <iwan.kawrakow@gmail.com>
This commit is contained in:
Kawrakow
2024-08-20 17:15:47 +03:00
committed by GitHub
parent 38dcba95fe
commit 8a10467990
9 changed files with 437 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
src/llama.cpp
View File

@@ -8317,14 +8317,17 @@ static struct ggml_tensor * llm_build_kqv(
//try from phi2
//ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
kq = ggml_tanh(ctx, ggml_scale(ctx, kq, 0.08838834764831845f/30.0f));
kq = ggml_scale(ctx, kq, 30);
//kq = ggml_tanh(ctx, ggml_scale(ctx, kq, 0.08838834764831845f/30.0f));
//kq = ggml_scale(ctx, kq, 30);
kq = ggml_softcap(ctx, kq, 0.08838834764831845f/30.0f, 30.f);
}
if (hparams.attn_soft_cap) {
kq = ggml_scale(ctx, kq, 1.0f / hparams.f_attn_logit_softcapping);
kq = ggml_tanh(ctx, kq);
kq = ggml_scale(ctx, kq, hparams.f_attn_logit_softcapping);
kq = ggml_softcap(ctx, kq, 1.0f / hparams.f_attn_logit_softcapping, hparams.f_attn_logit_softcapping);
//kq = ggml_scale(ctx, kq, 1.0f / hparams.f_attn_logit_softcapping);
//kq = ggml_tanh(ctx, kq);
//kq = ggml_scale(ctx, kq, hparams.f_attn_logit_softcapping);
}
kq = ggml_soft_max_ext(ctx, kq, kq_mask, kq_scale, hparams.f_max_alibi_bias);
@@ -11935,9 +11938,10 @@ struct llm_build_context {
cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
// final logit soft-capping
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
cur = ggml_tanh(ctx0, cur);
cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
cur = ggml_softcap(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping, hparams.f_final_logit_softcapping);
//cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
//cur = ggml_tanh(ctx0, cur);
//cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
cb(cur, "result_output", -1);