Repository: ggerganov/llama.cpp
File: src/llama-kv-cache.cpp (C++)

Code before:
    //      xxxx------
    //      xxxxx-----
    //   Non-causal mask:
    //      xxxxx-----
    //      xxxxx-----
    //      xxxxx-----
    // To visualize the mask, see https://github.com/ggml-org/llama.cpp/pull/12615
    // TODO: optimize this section
    for (uint32_t h = 0; h < 1; ++h) {
        for (uint32_t s = 0; s < n_stream; ++s) {
            for (uint32_t ii = 0; ii < n_tps; ++ii) {
                const uint32_t i = s*n_tps + ii;

                const llama_seq_id seq_id = ubatch->seq_id[i][0];

                const auto & cells = v_cells[seq_to_stream[seq_id]];

                const llama_pos p1 = ubatch->pos[i];

                // for M-RoPE
                llama_kv_pos_mrope p1_mrope;
                if (ubatch->has_mrope()) {
                    p1_mrope.y = ubatch->pos[i + ubatch->n_tokens];
                    p1_mrope.x = ubatch->pos[i + ubatch->n_tokens*2];
                }

                const uint64_t idst = n_kv*(h*n_stream*n_tps_pad + s*n_tps_pad + ii);

                for (uint32_t j = 0; j < n_kv; ++j) {
                    if (cells.is_empty(j)) {
                        continue;
                    }

                    // mask the token if not the same sequence
                    if (!cells.seq_has(j, seq_id)) {
                        continue;
                    }

                    const llama_pos p0 = cells.pos_get(j);

                    // mask future tokens
                    if (causal_attn && p0 > p1) {
                        continue;
                    }

                    // M-RoPE causal mask
                    if (causal_attn && ubatch->has_mrope() && p0 == p1) {
                        const auto & p0_mrope = cells.pos_mrope_get(j);
                        if (p0_mrope.is_gt(p1_mrope)) {
                            continue;
                        }

Review comment (suggestion):
I find it confusing to have the order of the positions as `y, x`. It's more canonical to have the dimensions ordered by increasing significance - `x, y, z, ...`. This is also inline with the `ggml` convention for indexing.

I now notice that even the implementation of `ggml_rope_multi` uses this order. I would recommend to update this across the codebase for consistency. Even though it's a breaking change, it's better to do it now, before the mtmd stuff gets more adopted.

Code after:
    // n_tps == n_tokens_per_stream
    const int64_t n_tps     = n_tokens/n_stream;
    const int64_t n_tps_pad = GGML_PAD(n_tps, GGML_KQ_MASK_PAD);

    std::fill(data, data + ggml_nelements(dst), -INFINITY);

    // Use only the previous KV cells of the correct sequence for each token of the ubatch.
    // It's assumed that if a token in the batch has multiple sequences, they are equivalent.
    // Example with a cache of 10 tokens, 2 tokens populated in cache and 3 tokens in batch:
    //   Causal mask:
    //      xxx-------
    //      xxxx------
    //      xxxxx-----
    //   Non-causal mask:
    //      xxxxx-----
    //      xxxxx-----
    //      xxxxx-----
    // To visualize the mask, see https://github.com/ggml-org/llama.cpp/pull/12615
    // TODO: optimize this section
    for (uint32_t h = 0; h < 1; ++h) {
        for (uint32_t s = 0; s < n_stream; ++s) {
            for (uint32_t ii = 0; ii < n_tps; ++ii) {
                const uint32_t i = s*n_tps + ii;

                const llama_seq_id seq_id = ubatch->seq_id[i][0];

                const auto & cells = v_cells[seq_to_stream[seq_id]];

                const llama_pos p1 = ubatch->pos[i];

                // for M-RoPE
                const bool is_2d = ubatch->is_pos_2d();
                const llama_pos p1_x = is_2d ? ubatch->pos[i + ubatch->n_tokens*2] : 0;
                const llama_pos p1_y = is_2d ? ubatch->pos[i + ubatch->n_tokens]   : 0;

                const uint64_t idst = n_kv*(h*n_stream*n_tps_pad + s*n_tps_pad + ii);

                for (uint32_t j = 0; j < n_kv; ++j) {
                    if (cells.is_empty(j)) {
                        continue;
                    }

                    // mask the token if not the same sequence
                    if (!cells.seq_has(j, seq_id)) {
                        continue;
                    }

                    const llama_pos p0 = cells.pos_get(j);

                    // mask future tokens
                    if (causal_attn && p0 > p1) {