add llm_graph_input_dsv4

Author: am17an

src/CMakeLists.txt

@@ -25,6 +25,7 @@ add_library(llama
             llama-kv-cache.cpp
             llama-kv-cache-iswa.cpp
             llama-kv-cache-dsa.cpp
+            llama-kv-cache-dsv4.cpp
             llama-memory.cpp
             llama-memory-hybrid.cpp
             llama-memory-hybrid-iswa.cpp

src/llama-graph.cpp

@@ -8,6 +8,7 @@
 #include "llama-kv-cache.h"
 #include "llama-kv-cache-iswa.h"
 #include "llama-kv-cache-dsa.h"
+#include "llama-kv-cache-dsv4.h"
 #include "llama-memory-hybrid.h"
 #include "llama-memory-hybrid-iswa.h"
 #include "llama-memory-recurrent.h"
@@ -17,6 +18,7 @@
 #include <cstring>
 #include <numeric>
 #include <sstream>
+#include <string>
 #include <unordered_set>
 
 // dedup helpers
@@ -620,6 +622,223 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) {
     return res;
 }
 
+static void dsv4_set_i64(ggml_tensor * dst, const std::vector<int64_t> & src) {
+    if (!dst || !dst->buffer) {
+        return;
+    }
+
+    GGML_ASSERT(dst->ne[0] == (int64_t) src.size());
+    ggml_backend_tensor_set(dst, src.data(), 0, src.size()*ggml_element_size(dst));
+}
+
+static void dsv4_set_i32(ggml_tensor * dst, const std::vector<int32_t> & src) {
+    if (!dst || !dst->buffer) {
+        return;
+    }
+
+    GGML_ASSERT(dst->ne[0] == (int64_t) src.size());
+    ggml_backend_tensor_set(dst, src.data(), 0, src.size()*ggml_element_size(dst));
+}
+
+static void dsv4_set_kq_mask(
+        ggml_tensor * dst,
+        const llama_kv_cache_dsv4_context::comp_plan & plan,
+        uint32_t n_tokens) {
+    if (!dst || !dst->buffer) {
+        return;
+    }
+
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->ne[0] == plan.n_kv);
+    GGML_ASSERT(dst->ne[1] == (int64_t) n_tokens);
+    GGML_ASSERT(dst->ne[2] == 1);
+    GGML_ASSERT(dst->ne[3] == 1);
+    GGML_ASSERT((int64_t) plan.n_visible.size() == dst->ne[1]);
+    GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
+
+    float * data = (float *) dst->data;
+
+    for (int64_t i = 0; i < dst->ne[1]; ++i) {
+        const int32_t n_visible = plan.n_visible[i];
+
+        for (int64_t j = 0; j < dst->ne[0]; ++j) {
+            data[i*dst->ne[0] + j] = j < n_visible ? 0.0f : -INFINITY;
+        }
+    }
+}
+
+static std::string dsv4_plan_positions(const std::vector<int32_t> & values) {
+    std::ostringstream ss;
+    ss << "[";
+    for (size_t i = 0; i < values.size(); ++i) {
+        if (i > 0) {
+            ss << ", ";
+        }
+        ss << values[i];
+    }
+    ss << "]";
+    return ss.str();
+}
+
+static bool dsv4_compress_debug() {
+    static const bool debug = []() {
+        const char * env = getenv("LLAMA_DSV4_COMPRESS_DEBUG");
+        return env && atoi(env) > 0;
+    }();
+
+    return debug;
+}
+
+static void dsv4_set_comp_inputs(
+        const llm_graph_input_dsv4::comp_input & inp,
+        const llama_kv_cache_dsv4_context::comp_plan & plan,
+        const char * name,
+        bool debug,
+        uint32_t n_tokens) {
+    dsv4_set_i64(inp.write_idxs, plan.write_idxs);
+    dsv4_set_i32(inp.write_pos,  plan.write_pos);
+    dsv4_set_i32(inp.write_end,  plan.write_end);
+    dsv4_set_i32(inp.pending_end, plan.pending_end);
+    dsv4_set_i32(inp.state_idxs, plan.state_idxs);
+    dsv4_set_i32(inp.state_pos, plan.state_pos);
+    dsv4_set_i32(inp.state_read_idxs, plan.state_read_idxs);
+    dsv4_set_i64(inp.state_write_idxs, plan.state_write_idxs);
+    dsv4_set_i32(inp.state_write_pos, plan.state_write_pos);
+    dsv4_set_i32(inp.state_write_end, plan.state_write_end);
+    dsv4_set_i32(inp.n_visible,  plan.n_visible);
+    dsv4_set_kq_mask(inp.kq_mask, plan, n_tokens);
+
+    if (debug || dsv4_compress_debug()) {
+        LLAMA_LOG_INFO("%s: %s ratio=%u, n_tokens=%u, write_end=%s, state_write_end=%s, pending_end=%s\n",
+                __func__, name, plan.ratio, n_tokens,
+                dsv4_plan_positions(plan.write_end).c_str(),
+                dsv4_plan_positions(plan.state_write_end).c_str(),
+                dsv4_plan_positions(plan.pending_end).c_str());
+    }
+}
+
+static bool dsv4_can_reuse_tensor_1d(ggml_tensor * t, int64_t ne0) {
+    return (t == nullptr && ne0 == 0) || (t != nullptr && t->ne[0] == ne0);
+}
+
+static bool dsv4_can_reuse_kq_mask(
+        ggml_tensor * t,
+        const llama_kv_cache_dsv4_context::comp_plan & plan,
+        uint32_t n_tokens) {
+    if (plan.n_kv == 0) {
+        return t == nullptr;
+    }
+
+    return t != nullptr &&
+           t->ne[0] == plan.n_kv &&
+           t->ne[1] == (int64_t) n_tokens &&
+           t->ne[2] == 1 &&
+           t->ne[3] == 1;
+}
+
+static bool dsv4_can_reuse_comp_input(
+        const llm_graph_input_dsv4::comp_input & inp,
+        const llama_kv_cache_dsv4_context::comp_plan & plan,
+        uint32_t n_tokens) {
+    const int64_t n_write = plan.write_idxs.size();
+
+    bool res = true;
+    res &= dsv4_can_reuse_tensor_1d(inp.write_idxs, n_write);
+    res &= dsv4_can_reuse_tensor_1d(inp.write_pos,  n_write);
+    res &= dsv4_can_reuse_tensor_1d(inp.write_end,  n_write);
+    res &= dsv4_can_reuse_tensor_1d(inp.pending_end, plan.pending_end.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_idxs, plan.state_idxs.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_pos, plan.state_pos.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_read_idxs, plan.state_read_idxs.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_write_idxs, plan.state_write_idxs.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_write_pos, plan.state_write_pos.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.state_write_end, plan.state_write_end.size());
+    res &= dsv4_can_reuse_tensor_1d(inp.n_visible,  plan.n_visible.size());
+    res &= dsv4_can_reuse_kq_mask(inp.kq_mask, plan, n_tokens);
+
+    return res;
+}
+
+static ggml_tensor * dsv4_build_input_1d(
+        ggml_context * ctx,
+        ggml_type type,
+        int64_t ne0,
+        const std::string & name) {
+    if (ne0 == 0) {
+        return nullptr;
+    }
+
+    ggml_tensor * res = ggml_new_tensor_1d(ctx, type, ne0);
+    ggml_set_input(res);
+    ggml_set_name(res, name.c_str());
+
+    return res;
+}
+
+static void dsv4_build_comp_inputs(
+        ggml_context * ctx,
+        llm_graph_input_dsv4::comp_input & inp,
+        const llama_kv_cache_dsv4_context::comp_plan & plan,
+        const char * name) {
+    const int64_t n_write = plan.write_idxs.size();
+
+    inp.write_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I64, n_write, std::string("dsv4_") + name + "_write_idxs");
+    inp.write_pos  = dsv4_build_input_1d(ctx, GGML_TYPE_I32, n_write, std::string("dsv4_") + name + "_write_pos");
+    inp.write_end  = dsv4_build_input_1d(ctx, GGML_TYPE_I32, n_write, std::string("dsv4_") + name + "_write_end");
+    inp.pending_end = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.pending_end.size(), std::string("dsv4_") + name + "_pending_end");
+    inp.state_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_idxs.size(), std::string("dsv4_") + name + "_state_idxs");
+    inp.state_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_pos.size(), std::string("dsv4_") + name + "_state_pos");
+    inp.state_read_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_read_idxs.size(), std::string("dsv4_") + name + "_state_read_idxs");
+    inp.state_write_idxs = dsv4_build_input_1d(ctx, GGML_TYPE_I64, plan.state_write_idxs.size(), std::string("dsv4_") + name + "_state_write_idxs");
+    inp.state_write_pos = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_write_pos.size(), std::string("dsv4_") + name + "_state_write_pos");
+    inp.state_write_end = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.state_write_end.size(), std::string("dsv4_") + name + "_state_write_end");
+    inp.n_visible  = dsv4_build_input_1d(ctx, GGML_TYPE_I32, plan.n_visible.size(), std::string("dsv4_") + name + "_n_visible");
+
+    if (plan.n_kv > 0) {
+        inp.kq_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, plan.n_kv, plan.n_visible.size(), 1, 1);
+        ggml_set_input(inp.kq_mask);
+        ggml_set_name(inp.kq_mask, (std::string("dsv4_") + name + "_kq_mask").c_str());
+    }
+}
+
+void llm_graph_input_dsv4::set_input(const llama_ubatch * ubatch) {
+    inp_raw->mctx = mctx->get_raw();
+    inp_raw->set_input(ubatch);
+
+    dsv4_set_comp_inputs(inp_csa, mctx->get_csa_plan(), "csa", debug > 0, ubatch->n_tokens);
+    dsv4_set_comp_inputs(inp_hca, mctx->get_hca_plan(), "hca", debug > 0, ubatch->n_tokens);
+    dsv4_set_comp_inputs(inp_lid, mctx->get_lid_plan(), "lid", debug > 0, ubatch->n_tokens);
+
+    if (inp_lid.k_rot && inp_lid.k_rot->buffer) {
+        mctx->get_lid()->set_input_k_rot(inp_lid.k_rot);
+    }
+}
+
+bool llm_graph_input_dsv4::can_reuse(const llm_graph_params & params) {
+    const auto * mctx = static_cast<const llama_kv_cache_dsv4_context *>(params.mctx);
+
+    this->mctx = mctx;
+    inp_raw->mctx = mctx->get_raw();
+
+    bool res = true;
+
+    if (inp_raw->self_k_idxs && inp_raw->self_k_idxs->buffer) {
+        res &= inp_raw->self_k_idxs->ne[0] == params.ubatch.n_tokens;
+        res &= can_reuse_kq_mask(inp_raw->self_kq_mask, mctx->get_raw()->get_base(), params.ubatch, params.cparams);
+    }
+
+    if (inp_raw->self_k_idxs_swa && inp_raw->self_k_idxs_swa->buffer) {
+        res &= inp_raw->self_k_idxs_swa->ne[0] == params.ubatch.n_tokens;
+        res &= can_reuse_kq_mask(inp_raw->self_kq_mask_swa, mctx->get_raw()->get_swa(), params.ubatch, params.cparams);
+    }
+
+    res &= dsv4_can_reuse_comp_input(inp_csa, mctx->get_csa_plan(), params.ubatch.n_tokens);
+    res &= dsv4_can_reuse_comp_input(inp_hca, mctx->get_hca_plan(), params.ubatch.n_tokens);
+    res &= dsv4_can_reuse_comp_input(inp_lid, mctx->get_lid_plan(), params.ubatch.n_tokens);
+
+    return res;
+}
+
 void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) {
     GGML_ASSERT(cross_kq_mask);
 
@@ -2731,6 +2950,46 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const
     return (llm_graph_input_attn_kv_iswa *) res->add_input(std::move(inp));
 }
 
+llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
+    const auto * mctx_cur = static_cast<const llama_kv_cache_dsv4_context *>(mctx);
+    const auto * raw_ctx  = mctx_cur->get_raw();
+
+    auto inp_raw = std::make_unique<llm_graph_input_attn_kv_iswa>(hparams, cparams, raw_ctx);
+
+    {
+        inp_raw->self_k_idxs = raw_ctx->get_base()->build_input_k_idxs(ctx0, ubatch);
+        inp_raw->self_v_idxs = raw_ctx->get_base()->build_input_v_idxs(ctx0, ubatch);
+
+        inp_raw->self_kq_mask = build_attn_inp_kq_mask(ctx0, raw_ctx->get_base(), ubatch, cparams);
+        inp_raw->self_kq_mask_cnv = inp_raw->self_kq_mask;
+    }
+
+    {
+        GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "DSV4 expects SWA raw cache");
+
+        inp_raw->self_k_idxs_swa = raw_ctx->get_swa()->build_input_k_idxs(ctx0, ubatch);
+        inp_raw->self_v_idxs_swa = raw_ctx->get_swa()->build_input_v_idxs(ctx0, ubatch);
+
+        inp_raw->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, raw_ctx->get_swa(), ubatch, cparams);
+        inp_raw->self_kq_mask_swa_cnv = inp_raw->self_kq_mask_swa;
+    }
+
+    inp_raw->self_k_rot = raw_ctx->get_base()->build_input_k_rot(ctx0);
+    inp_raw->self_v_rot = raw_ctx->get_base()->build_input_v_rot(ctx0);
+
+    inp_raw->self_k_rot_swa = raw_ctx->get_swa()->build_input_k_rot(ctx0);
+    inp_raw->self_v_rot_swa = raw_ctx->get_swa()->build_input_v_rot(ctx0);
+
+    auto inp = std::make_unique<llm_graph_input_dsv4>(cparams, std::move(inp_raw), mctx_cur);
+
+    dsv4_build_comp_inputs(ctx0, inp->inp_csa, mctx_cur->get_csa_plan(), "csa");
+    dsv4_build_comp_inputs(ctx0, inp->inp_hca, mctx_cur->get_hca_plan(), "hca");
+    dsv4_build_comp_inputs(ctx0, inp->inp_lid, mctx_cur->get_lid_plan(), "lid");
+    inp->inp_lid.k_rot = mctx_cur->get_lid()->build_input_k_rot(ctx0);
+
+    return (llm_graph_input_dsv4 *) res->add_input(std::move(inp));
+}
+
 ggml_tensor * llm_graph_context::build_rs(
         ggml_tensor * s,
         ggml_tensor * state_copy_main,

src/llama-graph.h

@@ -23,6 +23,7 @@ struct llama_memory_context_i;
 
 class llama_kv_cache_context;
 class llama_kv_cache_dsa_context;
+class llama_kv_cache_dsv4_context;
 class llama_kv_cache_iswa_context;
 class llama_memory_recurrent_context;
 class llama_memory_hybrid_context;
@@ -458,6 +459,57 @@ class llm_graph_input_attn_kv_iswa : public llm_graph_input_i {
     const llama_kv_cache_iswa_context * mctx;
 };
 
+class llm_graph_input_dsv4 : public llm_graph_input_i {
+public:
+    struct comp_input {
+        ggml_tensor * write_idxs = nullptr; // I64 [n_write]
+        ggml_tensor * write_pos  = nullptr; // I32 [n_write]
+        ggml_tensor * write_end  = nullptr; // I32 [n_write]
+        ggml_tensor * pending_end = nullptr; // I32 [n_pending]
+
+        ggml_tensor * state_idxs       = nullptr; // I32 [n_state]
+        ggml_tensor * state_pos        = nullptr; // I32 [n_state]
+        ggml_tensor * state_read_idxs  = nullptr; // I32 [ratio*n_state_write]
+        ggml_tensor * state_write_idxs = nullptr; // I64 [n_state_write]
+        ggml_tensor * state_write_pos  = nullptr; // I32 [n_state_write]
+        ggml_tensor * state_write_end  = nullptr; // I32 [n_state_write]
+
+        ggml_tensor * n_visible  = nullptr; // I32 [n_batch]
+        ggml_tensor * kq_mask    = nullptr; // F32 [n_kv, n_batch]
+
+        ggml_tensor * k_rot      = nullptr;
+    };
+
+    llm_graph_input_dsv4(
+            const llama_cparams & cparams,
+            std::unique_ptr<llm_graph_input_attn_kv_iswa> inp_raw,
+            const llama_kv_cache_dsv4_context * mctx) :
+        inp_raw(std::move(inp_raw)),
+        cparams(cparams),
+        mctx(mctx) {
+    }
+    ~llm_graph_input_dsv4() = default;
+
+    void set_input(const llama_ubatch * ubatch) override;
+
+    bool can_reuse(const llm_graph_params & params) override;
+
+    llm_graph_input_attn_kv_iswa * get_raw() const { return inp_raw.get(); }
+    const comp_input & get_csa() const { return inp_csa; }
+    const comp_input & get_hca() const { return inp_hca; }
+    const comp_input & get_lid() const { return inp_lid; }
+
+    std::unique_ptr<llm_graph_input_attn_kv_iswa> inp_raw;
+
+    comp_input inp_csa;
+    comp_input inp_hca;
+    comp_input inp_lid;
+
+    const llama_cparams cparams;
+
+    const llama_kv_cache_dsv4_context * mctx;
+};
+
 class llm_graph_input_attn_cross : public llm_graph_input_i {
 public:
     llm_graph_input_attn_cross(const llama_cross * cross) : cross(cross) {}
@@ -1033,6 +1085,8 @@ struct llm_graph_context {
 
     llm_graph_input_attn_kv_iswa * build_attn_inp_kv_iswa() const;
 
+    llm_graph_input_dsv4 * build_inp_dsv4() const;
+
     // note: if k_cur or v_cur are not provided, they will not be stored in the memory
     ggml_tensor * build_attn(
             llm_graph_input_attn_kv_iswa * inp,