Repository: ggerganov/llama.cpp
File: tests/test-backend-ops.cpp (C++)

Code before:
                t = ggml_new_tensor_4d(ctx, type, ne_perm[0], ne_perm[1], ne_perm[2], ne_perm[3]);
            }
            if (permute != std::array<int32_t, 4>{0, 1, 2, 3}) {
                t = ggml_permute(ctx, t, permute[0], permute[1], permute[2], permute[3]);
            }
            return t;
        };

        ggml_tensor * q = create_permuted(GGML_TYPE_F32, hsk_padded, nb, nh*nr23[0], nr23[1], false);
        ggml_set_name(q, "q");

        ggml_tensor * k = create_permuted(type_KV,       hsk_padded, kv, nh,         nr23[1], true); // the K tensor is usually a view of the K cache
        ggml_set_name(k, "k");

        ggml_tensor * v = nullptr;
        if (hsk_padded == 576 && hsv_padded == 512) {
            // TODO: this branch should become a separate test case parameter instead of hardcoding this for these head shapes

            // in this branch, the V cache is sub-view of the K cache. this is used by some MLA-based models
            // for more info:
            //   - https://github.com/ggml-org/llama.cpp/pull/13435
            //   - https://github.com/ggml-org/llama.cpp/pull/18953#issuecomment-3774948392
            v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], 0);

            // note: this is the currently assumed layout by the CUDA FA implementation
            //       however, this layout is problematic, because the offset can become very inconveniet for quantized KV types (i.e. not multiple of 16)
            //       so the models/deepseek2.cpp implementation has been updated to concat the PE data at the end of the row, resulting in the layout above
            //       the CUDA implementation has to be updated to not make an assumption abou the layout, and instead use the V tensor as it is
            //v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], ggml_row_size(type_KV, hsk_padded - hsv_padded)/ggml_blck_size(type_KV));
        } else {
            v = create_permuted(type_KV,       hsv_padded, kv, nh,         nr23[1], true); // the V tensor is usually a view of the V cache
        }
        ggml_set_name(v, "v");

        ggml_tensor * m = nullptr;
        if (mask) {
            m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, nr23[1]);
            ggml_set_name(m, "m");
        }

        ggml_tensor * s = nullptr;
        if (sinks) {
            s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, q->ne[2]);
            ggml_set_name(s, "s");
        }

        ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hsk), max_bias, logit_softcap);
        ggml_flash_attn_ext_add_sinks(out, s);
        ggml_flash_attn_ext_set_prec (out, prec);
        ggml_set_name(out, "out");


Review comment (bug):
I'm not sure what you mean here. Moving the K-only data to the front does not fix any issues with memory alignment. On master the check for memory alignment causes quantized KV cache to fall back to CPU but that is a bug with the check, see https://github.com/ggml-org/llama.cpp/pull/19023 .

Code after:
                t = ggml_new_tensor_4d(ctx, type, ne_perm[0], ne_perm[1], ne_perm[2], ne_perm[3]);
            }
            if (permute != std::array<int32_t, 4>{0, 1, 2, 3}) {
                t = ggml_permute(ctx, t, permute[0], permute[1], permute[2], permute[3]);
            }
            return t;
        };

        ggml_tensor * q = create_permuted(GGML_TYPE_F32, hsk_padded, nb, nh*nr23[0], nr23[1], false);
        ggml_set_name(q, "q");

        ggml_tensor * k = create_permuted(type_KV,       hsk_padded, kv, nh,         nr23[1], true); // the K tensor is usually a view of the K cache
        ggml_set_name(k, "k");

        ggml_tensor * v = nullptr;
        if (hsk_padded == 576 && hsv_padded == 512) {
            // TODO: this branch should become a separate test case parameter instead of hardcoding this for these head shapes

            // in this branch, the V cache is sub-view of the K cache. this is used by some MLA-based models
            // for more info:
            //   - https://github.com/ggml-org/llama.cpp/pull/13435
            //   - https://github.com/ggml-org/llama.cpp/pull/18953#issuecomment-3774948392
            //   - https://github.com/ggml-org/llama.cpp/pull/18986
            v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], 0);
        } else {
            v = create_permuted(type_KV,       hsv_padded, kv, nh,         nr23[1], true); // the V tensor is usually a view of the V cache
        }
        ggml_set_name(v, "v");

        ggml_tensor * m = nullptr;
        if (mask) {
            m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, nr23[1]);
            ggml_set_name(m, "m");
        }

        ggml_tensor * s = nullptr;
        if (sinks) {
            s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, q->ne[2]);
            ggml_set_name(s, "s");
        }

        ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hsk), max_bias, logit_softcap);
        ggml_flash_attn_ext_add_sinks(out, s);
        ggml_flash_attn_ext_set_prec (out, prec);
        ggml_set_name(out, "out");

        return out;
    }

    void initialize_tensors(ggml_context * ctx) override {
        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {