Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cpu/ops.cpp (C++)

Code before:
                    for (int tk = 0; tk < KV_TILE_SZ; tk++) {
                        const float mv = slope * GGML_CPU_FP16_TO_FP32(mp_row[ic + tk]);
                        KQ[tq * KV_TILE_SZ + tk] += mv;
                    }
                }
            }

            for (int tq = 0; tq < Q_TILE_SZ; tq++) {
                float * kq_row = KQ + tq * KV_TILE_SZ;

                float tile_max;
                ggml_vec_max_f32(KV_TILE_SZ, &tile_max, kq_row);

                const float Mold = M[tq];
                const float Mnew = fmaxf(Mold, tile_max);

                if (Mnew > Mold) {
                    const float ms = expf(Mold - Mnew);
                    ggml_vec_scale_f32(DV, VKQ32 + tq * DV, ms);
                    S[tq] *= ms;
                }
                M[tq] = Mnew;

                S[tq] += ggml_vec_soft_max_f32(KV_TILE_SZ, kq_row, kq_row, Mnew);
            }

            // Convert V tile to F32 first (if F16), then do MAD
            // On x86, ggml_vec_mad_f16 internall converts F16<->F32 on every load/store, so pre-converting is faster.
            // TODO: on ARM, native f16 should be faster
            if (kv_type == GGML_TYPE_F16) {
                for (int tk = 0; tk < KV_TILE_SZ; tk++) {
                    const ggml_fp16_t * v_row = (const ggml_fp16_t *)((const char *) v->data + ((ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3));
                    ggml_fp16_to_fp32_row(v_row, V32 + tk * DV, DV);
                }
                for (int tq = 0; tq < Q_TILE_SZ; tq++) {
                    float * vkq_row = VKQ32 + tq * DV;
                    for (int tk = 0; tk < KV_TILE_SZ; tk++) {
                        const float p = KQ[tq * KV_TILE_SZ + tk];
                        ggml_vec_mad_f32(DV, vkq_row, V32 + tk * DV, p);
                    }
                }
            } else {
                for (int tq = 0; tq < Q_TILE_SZ; tq++) {
                    float * vkq_row = VKQ32 + tq * DV;
                    for (int tk = 0; tk < KV_TILE_SZ; tk++) {
                        const float p = KQ[tq * KV_TILE_SZ + tk];
                        const float * v_row = (const float *)((const char *) v->data + ((ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3));
                        ggml_vec_mad_f32(DV, vkq_row, v_row, p);
                    }
                }
            }

Review comment (performance):
It might be worth to pre-load the mask tile at the start of the loop and skip the iteration if all values are -INF. This helps for unified kv cache with multiple sequences. Probably you should see a benefit in the prompt processing at larger `-np` with this command:

```bash
llama-batched-bench -m [model] -c 33792 -npp 2048 -ntg 32 -npl 1,2,4 -kvu
```

Can be done in follow-up PR - just leave a TODO for now.

Code after:
                }

                if (can_skip) {
                    continue;
                }
            }

            for (int tq = 0; tq < Q_TILE_SZ; tq++) {
                const void * q_row = (const char *)Q_q + tq * DK * kv_type_size;
                for (int tk = 0; tk < KV_TILE_SZ; tk++) {
                    const void * k_row = (const char *) k->data + ((ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3);
                    float s;
                    kv_vec_dot(DK, &s, 0, k_row, 0, q_row, 0, 1);
                    KQ[tq * KV_TILE_SZ + tk] = s * scale;
                }
            }

            if (logit_softcap != 0.0f) {
                ggml_vec_tanh_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, KQ);
                ggml_vec_scale_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, logit_softcap);
            }

            if (mask) {
                ggml_vec_add_f32(tile_rows * KV_TILE_SZ, KQ, KQ, mask32);
            }

            bool skip[Q_TILE_SZ] = {};

            for (int tq = 0; tq < Q_TILE_SZ; tq++) {
                float * kq_row = KQ + tq * KV_TILE_SZ;

                float tile_max;
                ggml_vec_max_f32(KV_TILE_SZ, &tile_max, kq_row);

                if (tile_max == -INFINITY) {
                    skip[tq] = true;
                    continue;
                }

                const float Mold = M[tq];
                const float Mnew = fmaxf(Mold, tile_max);

                if (Mnew > Mold) {
                    const float ms = expf(Mold - Mnew);
                    ggml_vec_scale_f32(DV, VKQ32 + tq * DV, ms);
                    S[tq] *= ms;
                }
                M[tq] = Mnew;


                S[tq] += ggml_vec_soft_max_f32(KV_TILE_SZ, kq_row, kq_row, Mnew);