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

Code before:

            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);

                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;

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

                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++) {
                    if (skip[tq]) continue;
                    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++) {
                    if (skip[tq]) continue;
                    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 (suggestion):
I think this is slightly better:

```suggestion
                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);
```

Code after:

            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);
            }

            // 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++) {
                    if (skip[tq]) continue;
                    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++) {
                    if (skip[tq]) continue;
                    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));