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

Code before:

    if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
        GGML_ASSERT(self_kq_mask_swa);
        GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer));

        float * data = (float *) self_kq_mask_swa->data;

        for (int h = 0; h < 1; ++h) {
            for (int i1 = 0; i1 < n_tokens; ++i1) {
                const llama_seq_id s1 = ubatch->seq_id[i1][0];

                for (int i0 = 0; i0 < n_tokens; ++i0) {
                    float f = -INFINITY;

                    for (int s = 0; s < ubatch->n_seq_id[i0]; ++s) {
                        const llama_seq_id s0 = ubatch->seq_id[i0][0];

                        if (s0 != s1) {
                            continue; // skip different sequences
                        }

                        if (cparams.causal_attn && ubatch->pos[i0] > ubatch->pos[i1]) {
                            continue; // skip future tokens for causal attention
                        }

                        if (llama_hparams::is_masked_swa(hparams.n_swa, hparams.swa_type, ubatch->pos[i0], ubatch->pos[i1])) {
                            continue; // skip masked tokens for SWA
                        }

                        // TODO: reimplement this like in llama_kv_cache_unified
                        if (hparams.use_alibi) {
                            f = -std::abs(ubatch->pos[i0] - ubatch->pos[i1]);
                        } else {
                            f = 0.0f;
                        }
                    }
                    data[h*(n_kv*n_tokens) + i1*n_kv + i0] = f;
                }
            }
        }
        if (debug) {
            print_mask(data, n_tokens, n_kv, hparams.n_swa, hparams.swa_type);
        }
    }
}

void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) {
    mctx->set_input_k_idxs(self_k_idxs, ubatch);
    mctx->set_input_v_idxs(self_v_idxs, ubatch);

    mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);

Review comment (performance):
Could this this for loop and the for loop above be extracted into a separate function that both could call? This line is only this part that differs I think. For example:
```console
static void compute_attn_mask(
        float * data,
        const llama_ubatch * ubatch,
        const llama_hparams & hparams,
        const llama_cparams & cparams,
        int64_t n_tokens,
        int64_t n_kv,
        bool apply_swa_mask) {
    for (int h = 0; h < 1; ++h) {
        for (int i1 = 0; i1 < n_tokens; ++i1) {
            const llama_seq_id s1 = ubatch->seq_id[i1][0];

            for (int i0 = 0; i0 < n_tokens; ++i0) {
                float f = -INFINITY;

                for (int s = 0; s < ubatch->n_seq_id[i0]; ++s) {
                    const llama_seq_id s0 = ubatch->seq_id[i0][0];

                    if (s0 != s1) {
                        continue; // skip different sequences
                    }

                    if (cparams.causal_attn && ubatch->pos[i0] > ubatch->pos[i1]) {
                        continue; // skip future tokens for causal attention
                    }

                    if (apply_swa_mask && llama_hparams::is_masked_swa(hparams.n_swa, hparams.swa_type, ubatch->pos[i0], ubatch->pos[i1])) {
                        continue; // skip masked tokens for SWA
                    }

                    // TODO: reimplement this like in llama_kv_cache_unified
                    if (hparams.use_alibi) {
                        f = -std::abs(ubatch->pos[i0] - ubatch->pos[i1]);
                    } else {
                        f = 0.0f;
                    }
                }
                data[h*(n_kv*n_tokens) + i1*n_kv + i0] = f;
            }
        }
    }
}
````

Code after:
        std::fill(data, data + ggml_nelements(self_kq_mask), -INFINITY);

        fill_mask(data, 0, LLAMA_SWA_TYPE_NONE);

        if (debug) {
            print_mask(data, n_tokens, n_kv, 0, LLAMA_SWA_TYPE_NONE);
        }
    }

    if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
        GGML_ASSERT(self_kq_mask_swa);
        GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask_swa->buffer));

        float * data = (float *) self_kq_mask_swa->data;

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

        fill_mask(data, hparams.n_swa, hparams.swa_type);

        if (debug) {
            print_mask(data, n_tokens, n_kv, hparams.n_swa, hparams.swa_type);
        }
    }
}

void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) {
    mctx->set_input_k_idxs(self_k_idxs, ubatch);
    mctx->set_input_v_idxs(self_v_idxs, ubatch);

    mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
}

bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) {
    const auto * mctx = static_cast<const llama_kv_cache_context *>(params.mctx);

    this->mctx = mctx;

    bool res = true;

    res &= self_k_idxs->ne[0] == params.ubatch.n_tokens;
  //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there

    res &= self_kq_mask->ne[0] == mctx->get_n_kv();
    res &= self_kq_mask->ne[1] == GGML_PAD(params.ubatch.n_tokens, GGML_KQ_MASK_PAD);

    return res;
}

void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {
    mctx->get_base()->set_input_k_idxs(self_k_idxs, ubatch);
    mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch);