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

Code before:
            const float * pq = (const float *) ((char *) q->data + ((iq1 + tq)*nbq1 + iq2*nbq2 + iq3*nbq3));
            kv_from_float(pq, (char *)Q_q + tq * DK * kv_type_size, DK);
        }
        // Zero-pad remaining rows
        for (int tq = tile_rows; tq < Q_TILE_SZ; tq++) {
            memset((char *)Q_q + tq * DK * kv_type_size, 0, DK * kv_type_size);
        }

        for (int64_t ic = 0; ic < nek1; ic += KV_TILE_SZ) {
            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 (mp) {
                for (int tq = 0; tq < Q_TILE_SZ; tq++) {
                    const ggml_fp16_t * mp_row = (const ggml_fp16_t *)((const char *) mask->data + (iq1 + tq)*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]);
                    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);

Review comment (suggestion):
Fixes out of bounds access:

```suggestion
                for (int tq = 0; tq < tile_rows; tq++) {
```

Code after:

        // v indices
        const int iv3 = iq3 / rv3;
        const int iv2 = iq2 / rv2;

        for (int tq = 0; tq < tile_rows; tq++) {
            const float * pq = (const float *) ((char *) q->data + ((iq1 + tq)*nbq1 + iq2*nbq2 + iq3*nbq3));
            kv_from_float(pq, (char *)Q_q + tq * DK * kv_type_size, DK);
        }
        // Zero-pad remaining rows
        for (int tq = tile_rows; tq < Q_TILE_SZ; tq++) {
            memset((char *)Q_q + tq * DK * kv_type_size, 0, DK * kv_type_size);
        }

        for (int64_t ic = 0; ic < nek1; ic += KV_TILE_SZ) {

            // skip the tile entirely if all the masks are -inf
            if (mask) {
                bool can_skip = true;
                for (int tq = 0; tq < tile_rows; tq++) {
                    const ggml_fp16_t * mp_row = (const ggml_fp16_t *)((const char *) mask->data + (iq1 + tq)*mask->nb[1] + (iq2%mask->ne[2])*mask->nb[2] + (iq3%mask->ne[3])*mask->nb[3]);
                    for (int tk = 0; tk < KV_TILE_SZ; tk++) {
                        mask32[tq * KV_TILE_SZ + tk] = slope * GGML_CPU_FP16_TO_FP32(mp_row[ic + tk]);
                        if (mask32[tq * KV_TILE_SZ + tk] != -INFINITY) {
                            can_skip = false;
                        }
                    }
                }

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