graph : add kv_idxs to unified_iswa input [no ci]

Author: ggerganov

src/llama-graph.cpp

@@ -291,6 +291,11 @@ void llm_graph_input_attn_kv_unified::set_input(const llama_ubatch * ubatch) {
 }
 
 void llm_graph_input_attn_kv_unified_iswa::set_input(const llama_ubatch * ubatch) {
+    if (self_kv_idxs) {
+        kv_state->get_base()->set_input_kv_idxs(self_kv_idxs, ubatch);
+        kv_state->get_swa ()->set_input_kv_idxs(self_kv_idxs, ubatch);
+    }
+
     if (self_kq_mask) {
         kv_state->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
     }
@@ -1287,8 +1292,10 @@ ggml_tensor * llm_graph_context::build_attn(
 
     // store to KV cache
     {
-        ggml_build_forward_expand(gf, kv_state->cpy_k(ctx0, k_cur, nullptr, il));
-        ggml_build_forward_expand(gf, kv_state->cpy_v(ctx0, v_cur, nullptr, il));
+        const auto & kv_idxs = inp->get_kv_idxs();
+
+        ggml_build_forward_expand(gf, kv_state->cpy_k(ctx0, k_cur, kv_idxs, il));
+        ggml_build_forward_expand(gf, kv_state->cpy_v(ctx0, v_cur, kv_idxs, il));
     }
 
     const auto & kq_mask = is_swa ? inp->get_kq_mask_swa() : inp->get_kq_mask();
@@ -1428,6 +1435,9 @@ llm_graph_input_attn_kv_unified_iswa * llm_graph_context::build_attn_inp_kv_unif
     {
         const auto n_kv = kv_state->get_base()->get_n_kv();
 
+        inp->self_kv_idxs = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+        ggml_set_input(inp->self_kv_idxs);
+
         inp->self_kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD));
         //cb(inp->self_kq_mask, "KQ_mask", -1);
         ggml_set_input(inp->self_kq_mask);

src/llama-graph.h

@@ -276,9 +276,12 @@ class llm_graph_input_attn_kv_unified_iswa : public llm_graph_input_i {
 
     void set_input(const llama_ubatch * ubatch) override;
 
+    ggml_tensor * get_kv_idxs()     const { return self_kv_idxs; }
     ggml_tensor * get_kq_mask()     const { return self_kq_mask_cnv; }
     ggml_tensor * get_kq_mask_swa() const { return self_kq_mask_swa_cnv; }
 
+    ggml_tensor * self_kv_idxs = nullptr; // I32 [n_batch]
+
     ggml_tensor * self_kq_mask         = nullptr; // F32 [n_kv, n_batch]
     ggml_tensor * self_kq_mask_cnv     = nullptr; //     [n_kv, n_batch]
     ggml_tensor * self_kq_mask_swa     = nullptr; // F32 [n_kv, n_batch]