#include "llama-dsv4.h"

#include "llama-context.h"
#include "llama-model.h"
#include "llama-impl.h"

#include "ggml.h"
#include "ggml-backend.h"

#include <algorithm>
#include <cstring>
#include <map>
#include <stdexcept>
#include <type_traits>
#include <unordered_set>

static bool dsv4_cache_type_supported(ggml_type type) {
    return type == GGML_TYPE_F16 || type == GGML_TYPE_BF16 || type == GGML_TYPE_Q8_0;
}

static bool dsv4_validate_cache_type(ggml_type type, int64_t width, const char * name) {
    if (!dsv4_cache_type_supported(type)) {
        LLAMA_LOG_ERROR("%s: unsupported DSV4 %s cache type %s\n", __func__, name, ggml_type_name(type));
        return false;
    }
    if (ggml_is_quantized(type) && width % ggml_blck_size(type) != 0) {
        LLAMA_LOG_ERROR("%s: DSV4 %s cache width %d is not aligned to %d elements for %s\n",
                __func__, name, (int)width, (int)ggml_blck_size(type), ggml_type_name(type));
        return false;
    }
    return true;
}

static ggml_backend_buffer_type_t llama_dsv4_layer_buft(const llama_context & lctx, int32_t il) {
    if (il >= 0 && il < (int32_t) lctx.model.buft_layer.size() && lctx.model.buft_layer[il].buft != nullptr) {
        return lctx.model.buft_layer[il].buft;
    }

    if (il >= 0 && il < (int32_t) lctx.model.layers.size()) {
        const ggml_tensor * ref = lctx.model.layers[il].attn_comp_wkv;
        if (ref == nullptr) {
            ref = lctx.model.layers[il].wq_a;
        }
        if (ref != nullptr && ref->buffer != nullptr) {
            return ggml_backend_buffer_get_type(ref->buffer);
        }
    }

    return llama_default_buffer_type_cpu(true);
}

static uint32_t dsv4_comp_size(uint32_t kv_size, uint32_t ratio) {
    return std::max<uint32_t>(1, (kv_size + ratio - 1)/ratio);
}

static bool dsv4_validate_csa_lid_visibility(
        const llama_context & lctx,
        uint32_t csa_kv_size,
        uint32_t lid_kv_size) {
    const auto & csa_plan = lctx.dsv4.csa_plan;
    const auto & lid_plan = lctx.dsv4.lid_plan;
    const auto & csa_ctx = lctx.dsv4.csa_ctx;
    const auto & lid_ctx = lctx.dsv4.lid_ctx;

    if (csa_kv_size != lid_kv_size ||
            csa_plan.n_stream != lid_plan.n_stream ||
            csa_plan.n_kv != lid_plan.n_kv ||
            csa_plan.n_visible != lid_plan.n_visible ||
            csa_ctx.graph_n_stream != lid_ctx.graph_n_stream ||
            csa_ctx.n_kv != lid_ctx.n_kv ||
            csa_ctx.sinfo.strm != lid_ctx.sinfo.strm ||
            csa_ctx.sinfo.idxs != lid_ctx.sinfo.idxs ||
            csa_ctx.sinfo.s0 != lid_ctx.sinfo.s0 ||
            csa_ctx.sinfo.s1 != lid_ctx.sinfo.s1) {
        LLAMA_LOG_ERROR("%s: DSV4 CSA/LID visibility contracts differ\n", __func__);
        return false;
    }

    return true;
}

static void dsv4_batch_shape(
        const llama_batch & batch,
        uint32_t & n_seqs,
        uint32_t & n_seq_tokens) {
    n_seqs = 1;
    n_seq_tokens = (uint32_t) std::max(1, batch.n_tokens);

    if (batch.n_tokens <= 0 || batch.n_seq_id == nullptr || batch.seq_id == nullptr) {
        return;
    }

    std::map<llama_seq_id, uint32_t> counts;
    for (int32_t i = 0; i < batch.n_tokens; ++i) {
        if (batch.n_seq_id[i] != 1 || batch.seq_id[i] == nullptr) {
            return;
        }

        counts[batch.seq_id[i][0]]++;
    }

    if (counts.empty()) {
        return;
    }

    const uint32_t seq_tokens = counts.begin()->second;
    for (const auto & [_, count] : counts) {
        if (count != seq_tokens) {
            return;
        }
    }

    n_seqs = (uint32_t) counts.size();
    n_seq_tokens = std::max<uint32_t>(1, seq_tokens);
}

static bool dsv4_batch_has_coupled(const llama_batch & batch) {
    if (batch.n_tokens <= 0 || batch.n_seq_id == nullptr) {
        return false;
    }

    for (int32_t i = 0; i < batch.n_tokens; ++i) {
        if (batch.n_seq_id[i] > 1) {
            return true;
        }
    }

    return false;
}

static bool dsv4_token_has_seq(const llama_batch & batch, int32_t i, llama_seq_id seq_id) {
    if (batch.n_seq_id == nullptr || batch.seq_id == nullptr || batch.seq_id[i] == nullptr) {
        return seq_id == 0;
    }

    for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
        if (batch.seq_id[i][s] == seq_id) {
            return true;
        }
    }

    return false;
}

static std::vector<llama_seq_id> dsv4_batch_unique_seq_ids(const llama_batch & batch) {
    std::vector<llama_seq_id> seq_ids;
    std::unordered_set<llama_seq_id> seen;

    if (batch.n_tokens <= 0 || batch.n_seq_id == nullptr || batch.seq_id == nullptr) {
        seq_ids.push_back(0);
        return seq_ids;
    }

    for (int32_t i = 0; i < batch.n_tokens; ++i) {
        if (batch.n_seq_id[i] <= 0 || batch.seq_id[i] == nullptr) {
            continue;
        }

        for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
            const llama_seq_id seq_id = batch.seq_id[i][s];
            if (seen.insert(seq_id).second) {
                seq_ids.push_back(seq_id);
            }
        }
    }

    if (seq_ids.empty()) {
        seq_ids.push_back(0);
    }

    return seq_ids;
}

static int64_t dsv4_stream_offset(uint32_t n_stream, llama_seq_id seq_id, uint32_t size) {
    if (n_stream <= 1) {
        return 0;
    }

    if (seq_id < 0 || (uint32_t) seq_id >= n_stream) {
        LLAMA_LOG_ERROR("%s: DSV4 seq_id %d is outside stream range %u\n", __func__, seq_id, n_stream);
        return -1;
    }

    return (int64_t) seq_id*size;
}

static int64_t dsv4_comp_graph_n_stream(const llama_batch & batch, uint32_t n_stream) {
    if (n_stream <= 1) {
        return 1;
    }

    const std::vector<llama_seq_id> seq_ids = dsv4_batch_unique_seq_ids(batch);
    if (seq_ids.size() <= 1 || dsv4_batch_has_coupled(batch)) {
        return 1;
    }

    return (int64_t) seq_ids.size();
}

static std::vector<llama_seq_id> dsv4_build_stream_seq_ids(
        const llama_batch & batch,
        uint32_t n_stream) {
    if (n_stream <= 1) {
        return { 0 };
    }

    const std::vector<llama_seq_id> seq_ids = dsv4_batch_unique_seq_ids(batch);
    if (seq_ids.size() <= 1 || dsv4_batch_has_coupled(batch)) {
        return { seq_ids.empty() ? 0 : seq_ids.front() };
    }

    return seq_ids;
}

static llama_context::dsv4_runtime::slot_info dsv4_build_comp_sinfo(
        const llama_batch & batch,
        uint32_t n_stream) {
    llama_context::dsv4_runtime::slot_info sinfo;

    const std::vector<llama_seq_id> seq_ids = dsv4_build_stream_seq_ids(batch, n_stream);
    const int64_t graph_n_stream = (int64_t) seq_ids.size();
    bool have_stream = false;

    sinfo.s0 = INT_MAX;
    sinfo.s1 = 0;
    sinfo.resize((size_t) std::max<int64_t>(1, graph_n_stream));
    for (int64_t s = 0; s < graph_n_stream; ++s) {
        const llama_seq_id seq_id = seq_ids[(size_t) s];
        const int64_t strm = dsv4_stream_offset(n_stream, seq_id, 1);
        if (strm < 0) {
            continue;
        }
        sinfo.strm[(size_t) s] = (llama_seq_id) strm;
        sinfo.idxs[(size_t) s].assign(1, 0);
        sinfo.s0 = std::min<int32_t>(sinfo.s0, (int32_t) strm);
        sinfo.s1 = std::max<int32_t>(sinfo.s1, (int32_t) strm);
        have_stream = true;
    }

    if (!have_stream) {
        sinfo.resize(1);
        sinfo.strm[0] = 0;
        sinfo.idxs[0].assign(1, 0);
        sinfo.s0 = 0;
        sinfo.s1 = 0;
    }

    if (n_stream > 1 && sinfo.s1 - sinfo.s0 + 1 != (int32_t) sinfo.n_stream()) {
        LLAMA_LOG_ERROR("%s: DSV4 compressed streams are not contiguous in batch\n", __func__);
    }

    return sinfo;
}

static llama_context::dsv4_runtime::slot_info dsv4_build_raw_read_sinfo(
        const llama_context::dsv4_runtime::slot_info & sinfo_write,
        const llama_batch & batch,
        uint32_t n_stream) {
    if (!dsv4_batch_has_coupled(batch)) {
        return sinfo_write;
    }

    const llama_seq_id seq_id =
            (batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.n_tokens > 0 && batch.n_seq_id[0] > 0 && batch.seq_id[0] != nullptr)
            ? batch.seq_id[0][0]
            : 0;
    const int64_t strm = dsv4_stream_offset(n_stream, seq_id, 1);
    if (strm < 0) {
        return {};
    }

    size_t i_stream = 0;
    for (; i_stream < sinfo_write.n_stream(); ++i_stream) {
        if ((int64_t) sinfo_write.strm[i_stream] == strm) {
            break;
        }
    }
    if (i_stream == sinfo_write.n_stream()) {
        LLAMA_LOG_ERROR("%s: DSV4 raw write stream not found for coupled read\n", __func__);
        return {};
    }

    llama_context::dsv4_runtime::slot_info sinfo;
    sinfo.resize(1);
    sinfo.strm[0] = sinfo_write.strm[i_stream];
    sinfo.idxs[0] = sinfo_write.idxs[i_stream];
    sinfo.s0 = (int32_t) strm;
    sinfo.s1 = sinfo.s0;

    return sinfo;
}

static bool dsv4_validate_batch_seq_ids(
        const llama_context & lctx,
        const llama_batch & batch) {
    if (batch.n_tokens <= 0 || batch.n_seq_id == nullptr || batch.seq_id == nullptr) {
        return true;
    }

    const uint32_t n_stream = std::max<uint32_t>(1, lctx.cparams.n_seq_max);
    for (int32_t i = 0; i < batch.n_tokens; ++i) {
        if (batch.n_seq_id[i] <= 0 || batch.seq_id[i] == nullptr) {
            LLAMA_LOG_ERROR("%s: DSV4 token %d is missing seq_id ownership\n", __func__, i);
            return false;
        }

        for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
            const llama_seq_id seq_id = batch.seq_id[i][s];
            if (seq_id < 0 || (uint32_t) seq_id >= n_stream) {
                LLAMA_LOG_ERROR("%s: DSV4 token %d seq_id %d is outside n_seq_max=%u\n",
                        __func__, i, seq_id, n_stream);
                return false;
            }
        }
    }

    return true;
}

static bool dsv4_build_raw_context(
        const llama_context & lctx,
        const llama_batch & batch,
        llama_context::dsv4_runtime::raw_context & raw) {
    raw = {};
    const uint32_t n_stream = std::max<uint32_t>(1, lctx.cparams.n_seq_max);
    const std::vector<llama_seq_id> write_seq_ids = dsv4_build_stream_seq_ids(batch, n_stream);
    raw.sinfo_write = dsv4_build_comp_sinfo(batch, n_stream);
    raw.sinfo_read  = dsv4_build_raw_read_sinfo(raw.sinfo_write, batch, n_stream);
    raw.graph_n_stream = (int64_t) raw.sinfo_write.n_stream();
    std::vector<llama_seq_id> read_seq_ids = write_seq_ids;

    if (dsv4_batch_has_coupled(batch)) {
        const llama_seq_id coupled_seq_id =
                (batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.n_tokens > 0 && batch.n_seq_id[0] > 0 && batch.seq_id[0] != nullptr)
                ? batch.seq_id[0][0]
                : 0;
        read_seq_ids.assign(1, coupled_seq_id);
    }

    if (batch.n_tokens <= 0) {
        return true;
    }

    const llama_kv_cache & kv = lctx.kv_self;
    if (kv.head + batch.n_tokens > (int32_t) kv.size) {
        LLAMA_LOG_ERROR("%s: DSV4 raw write slots [%d, %d) are outside kv cache size %u\n",
                __func__, kv.head, kv.head + batch.n_tokens, kv.size);
        return false;
    }

    raw.write_counts.push_back(batch.n_tokens);
    for (int32_t i = 0; i < batch.n_tokens; ++i) {
        const int32_t slot = kv.head + i;
        const llama_kv_cell & cell = kv.cells[(size_t) slot];

        if (batch.pos != nullptr && cell.pos != batch.pos[i]) {
            LLAMA_LOG_ERROR("%s: DSV4 raw write slot %d pos mismatch: cell=%d batch=%d\n",
                    __func__, slot, cell.pos, batch.pos[i]);
            return false;
        }

        raw.write_src_idxs.push_back(i);
        raw.write_dst_idxs.push_back(slot);
    }

    raw.n_kv = 0;

    for (size_t s = 0; s < raw.sinfo_read.n_stream(); ++s) {
        const llama_seq_id seq_id = read_seq_ids[s];
        raw.sinfo_read.idxs[s].clear();
        int32_t count = 0;
        for (uint32_t slot = 0; slot < kv.size; ++slot) {
            const llama_kv_cell & cell = kv.cells[slot];
            if (cell.is_empty() || cell.pos < 0) {
                continue;
            }
            if (!cell.has_seq_id(seq_id)) {
                continue;
            }
            raw.sinfo_read.idxs[s].push_back(slot);
            raw.read_dst_idxs.push_back((int32_t) slot);
            ++count;
        }
        raw.read_counts.push_back(count);
        raw.n_kv = std::max<int64_t>(raw.n_kv, count);
    }

    if (raw.read_counts.empty()) {
        raw.read_counts.push_back(0);
    }

    for (size_t s = 0; s < raw.sinfo_write.n_stream(); ++s) {
        const llama_seq_id seq_id = write_seq_ids[s];
        raw.sinfo_write.idxs[s].clear();
        for (int32_t i = 0; i < batch.n_tokens; ++i) {
            if (!dsv4_token_has_seq(batch, i, seq_id)) {
                continue;
            }
            raw.sinfo_write.idxs[s].push_back((uint32_t) (kv.head + i));
        }
    }

    if (raw.sinfo_write.n_stream() > 1) {
        std::vector<int32_t> write_src_idxs;
        std::vector<int32_t> write_dst_idxs;
        const size_t rows_per_stream = raw.sinfo_write.size();
        for (size_t s = 0; s < raw.sinfo_write.n_stream(); ++s) {
            if (raw.sinfo_write.idxs[s].size() != rows_per_stream) {
                LLAMA_LOG_ERROR("%s: DSV4 packed batch has unequal raw-write rows per stream\n", __func__);
                return false;
            }

            for (int32_t i = 0; i < batch.n_tokens; ++i) {
                if (dsv4_token_has_seq(batch, i, write_seq_ids[s])) {
                    write_src_idxs.push_back(i);
                }
            }

            for (uint32_t slot : raw.sinfo_write.idxs[s]) {
                write_dst_idxs.push_back((int32_t) slot);
            }
        }

        raw.write_src_idxs = std::move(write_src_idxs);
        raw.write_dst_idxs = std::move(write_dst_idxs);
    }

    // The graph exposes a rectangular raw-key view. Repeat the last valid row
    // for shorter streams; the corresponding mask entries remain -INFINITY.
    // This preserves the logical visibility while allowing one get_rows op to
    // serve all streams.
    if (raw.n_kv > 0) {
        raw.read_dst_idxs.clear();
        const size_t read_rows = GGML_PAD((size_t) raw.n_kv, 256u);
        for (size_t s = 0; s < raw.sinfo_read.n_stream(); ++s) {
            const auto & rows = raw.sinfo_read.idxs[s];
            for (uint32_t slot : rows) {
                raw.read_dst_idxs.push_back((int32_t) slot);
            }

            const int32_t pad = rows.empty() ? 0 : (int32_t) rows.back();
            for (size_t i = rows.size(); i < read_rows; ++i) {
                raw.read_dst_idxs.push_back(pad);
            }
        }
    }

    return true;
}

static llama_context::dsv4_runtime::comp_context dsv4_build_comp_context(
        const llama_batch & batch,
        uint32_t n_stream,
        int64_t n_kv) {
    llama_context::dsv4_runtime::comp_context ctx;
    ctx.sinfo = dsv4_build_comp_sinfo(batch, n_stream);
    ctx.graph_n_stream = dsv4_comp_graph_n_stream(batch, n_stream);
    ctx.n_kv = n_kv;
    return ctx;
}

static llama_context::dsv4_runtime::comp_plan dsv4_build_reserve_comp_plan(
        const llama_batch & batch,
        uint32_t ratio,
        bool overlap,
        uint32_t state_size,
        uint32_t kv_size,
        uint32_t n_stream) {
    llama_context::dsv4_runtime::comp_plan plan;
    plan.n_visible.resize((size_t) batch.n_tokens, (int32_t) kv_size);
    plan.n_stream = dsv4_comp_graph_n_stream(batch, n_stream);
    plan.n_kv = kv_size;

    if (batch.n_tokens == 0) {
        return plan;
    }

    uint32_t n_seqs = 1;
    uint32_t n_seq_tokens = 1;
    dsv4_batch_shape(batch, n_seqs, n_seq_tokens);

    plan.n_visible.assign((size_t) batch.n_tokens, 0);

    const uint64_t n_blocks_u64 = (uint64_t) n_seqs*((n_seq_tokens + ratio - 1)/ratio);
    const size_t n_blocks = (size_t) std::max<uint64_t>(1, n_blocks_u64);
    GGML_ASSERT((uint64_t) n_blocks == std::max<uint64_t>(1, n_blocks_u64));
    const uint64_t state_rows = (uint64_t) state_size*(uint64_t) n_stream;
    const size_t n_persist = (size_t) std::min<uint64_t>((uint64_t) batch.n_tokens, state_rows);

    plan.state_pos.resize((size_t) batch.n_tokens);
    plan.state_persist_src_idxs.resize(n_persist);
    plan.state_persist_dst_idxs.resize(n_persist);
    plan.state_read_idxs.resize((overlap ? 2u : 1u)*ratio*n_blocks);
    plan.state_write_idxs.resize(n_blocks);
    plan.state_write_pos.resize(n_blocks);

    return plan;
}

static uint32_t dsv4_cache_kv_size(const std::vector<ggml_tensor *> & tensors) {
    for (ggml_tensor * tensor : tensors) {
        if (tensor != nullptr) {
            return (uint32_t) tensor->ne[1];
        }
    }

    return 0;
}

static uint32_t dsv4_cache_state_size(const std::vector<ggml_tensor *> & tensors) {
    for (ggml_tensor * tensor : tensors) {
        if (tensor != nullptr) {
            return (uint32_t) tensor->ne[1];
        }
    }

    return 0;
}

static bool dsv4_validate_comp_plan(
        const char * tag,
        const llama_batch & batch,
        const llama_context::dsv4_runtime::comp_plan & plan,
        uint32_t ratio,
        bool overlap,
        uint32_t state_size,
        uint32_t kv_size,
        uint32_t n_stream) {
    const int64_t max_state_read_idx = (int64_t) state_size*n_stream + batch.n_tokens + (overlap ? 0 : -1);

    if (plan.n_visible.size() != (size_t) std::max(0, batch.n_tokens)) {
        LLAMA_LOG_ERROR("%s: DSV4 %s plan n_visible size mismatch: got=%zu expected=%d\n",
                __func__, tag, plan.n_visible.size(), std::max(0, batch.n_tokens));
        return false;
    }

    if (plan.state_pos.size() > (size_t) std::max(0, batch.n_tokens)) {
        LLAMA_LOG_ERROR("%s: DSV4 %s plan has too many state_pos rows: %zu > %d\n",
                __func__, tag, plan.state_pos.size(), std::max(0, batch.n_tokens));
        return false;
    }

    if (plan.state_persist_src_idxs.size() != plan.state_persist_dst_idxs.size()) {
        LLAMA_LOG_ERROR("%s: DSV4 %s persist idx size mismatch: src=%zu dst=%zu\n",
                __func__, tag, plan.state_persist_src_idxs.size(), plan.state_persist_dst_idxs.size());
        return false;
    }

    if (plan.state_write_idxs.size() != plan.state_write_pos.size()) {
        LLAMA_LOG_ERROR("%s: DSV4 %s write idx size mismatch: idxs=%zu pos=%zu\n",
                __func__, tag, plan.state_write_idxs.size(), plan.state_write_pos.size());
        return false;
    }

    for (size_t i = 0; i < plan.n_visible.size(); ++i) {
        const int32_t n_visible = plan.n_visible[i];
        if (n_visible < 0 || (uint32_t) n_visible > kv_size) {
            LLAMA_LOG_ERROR("%s: DSV4 %s n_visible[%zu]=%d exceeds kv_size=%u\n",
                    __func__, tag, i, n_visible, kv_size);
            return false;
        }
    }

    for (size_t i = 0; i < plan.state_pos.size(); ++i) {
        const int64_t pos = plan.state_pos[i];
        if (pos < 0 || pos >= (int64_t) ratio) {
            LLAMA_LOG_ERROR("%s: DSV4 %s state_pos[%zu]=%lld outside ratio=%u\n",
                    __func__, tag, i, (long long) pos, ratio);
            return false;
        }
    }

    for (size_t i = 0; i < plan.state_persist_src_idxs.size(); ++i) {
        const int64_t src = plan.state_persist_src_idxs[i];
        const int64_t dst = plan.state_persist_dst_idxs[i];
        if (src < 0 || src >= batch.n_tokens) {
            LLAMA_LOG_ERROR("%s: DSV4 %s persist src[%zu]=%lld outside current batch rows=%d\n",
                    __func__, tag, i, (long long) src, batch.n_tokens);
            return false;
        }
        if (dst < 0 || (uint32_t) dst >= state_size*n_stream) {
            LLAMA_LOG_ERROR("%s: DSV4 %s persist dst[%zu]=%lld outside state_size*n_stream=%u\n",
                    __func__, tag, i, (long long) dst, state_size*n_stream);
            return false;
        }
    }

    for (size_t i = 0; i < plan.state_read_idxs.size(); ++i) {
        const int64_t idx = plan.state_read_idxs[i];
        if (idx < 0 || idx > max_state_read_idx) {
            LLAMA_LOG_ERROR("%s: DSV4 %s read idx[%zu]=%lld outside max source row=%lld\n",
                    __func__, tag, i, (long long) idx, (long long) max_state_read_idx);
            return false;
        }
    }

    for (size_t i = 0; i < plan.state_write_idxs.size(); ++i) {
        const int64_t idx = plan.state_write_idxs[i];
        if (idx < 0 || (uint32_t) idx >= kv_size*n_stream) {
            LLAMA_LOG_ERROR("%s: DSV4 %s write idx[%zu]=%lld outside kv_size*n_stream=%u\n",
                    __func__, tag, i, (long long) idx, kv_size*n_stream);
            return false;
        }
    }

    if (plan.n_kv == 0 || (uint32_t) plan.n_kv > kv_size) {
        LLAMA_LOG_ERROR("%s: DSV4 %s plan n_kv=%lld outside kv_size=%u\n",
                __func__, tag, (long long) plan.n_kv, kv_size);
        return false;
    }

    return true;
}

static llama_context::dsv4_runtime::comp_plan dsv4_build_comp_plan(
        const llama_batch & batch,
        uint32_t ratio,
        bool overlap,
        uint32_t state_size,
        uint32_t kv_size,
        uint32_t n_stream) {
    llama_context::dsv4_runtime::comp_plan plan;
    plan.n_visible.resize((size_t) batch.n_tokens);
    plan.n_stream = dsv4_comp_graph_n_stream(batch, n_stream);

    if (n_stream <= 1 && dsv4_batch_unique_seq_ids(batch).size() > 1) {
        LLAMA_LOG_ERROR("%s: DSV4 single compressed stream cannot serve multiple sequences\n", __func__);
        return plan;
    }

    const int64_t state_rows = (int64_t) state_size*n_stream;

    struct persist_row {
        int32_t dst;
        int32_t src;
        llama_pos pos;
    };

    std::vector<persist_row> persist_rows;
    std::vector<int32_t> overlap_prev_reads;
    std::vector<int32_t> overlap_cur_reads;
    std::map<std::pair<llama_seq_id, llama_pos>, int32_t> curr_token_idx_map;

    for (int32_t i = 0; i < batch.n_tokens; ++i) {
        const llama_seq_id seq_id =
                batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.n_seq_id[i] > 0 && batch.seq_id[i] != nullptr
                ? batch.seq_id[i][0]
                : 0;
        curr_token_idx_map[std::make_pair(seq_id, batch.pos[i])] = i;
    }

    const auto state_source_idx = [&](llama_seq_id seq_id, llama_pos pos) -> int32_t {
        if (pos < 0) {
            return (int32_t) (state_rows + batch.n_tokens);
        }

        const auto it = curr_token_idx_map.find(std::make_pair(seq_id, pos));
        if (it != curr_token_idx_map.end()) {
            return (int32_t) (state_rows + it->second);
        }

        const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);
        GGML_ASSERT(stream_off >= 0);
        return (int32_t) (stream_off + pos%state_size);
    };

    for (int32_t i = 0; i < batch.n_tokens; ++i) {
        const llama_pos pos = batch.pos[i];
        if (pos < 0) {
            continue;
        }

        plan.state_pos.push_back((int32_t) (pos%ratio));

        const int64_t n_visible = (int64_t) (pos + 1)/ratio;
        plan.n_visible[(size_t) i] = (int32_t) n_visible;
        plan.n_kv = std::max(plan.n_kv, n_visible);

        const int32_t n_token_seqs =
                batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.seq_id[i] != nullptr
                ? batch.n_seq_id[i]
                : 1;
        for (int32_t s = 0; s < n_token_seqs; ++s) {
            const llama_seq_id seq_id =
                    batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.seq_id[i] != nullptr
                    ? batch.seq_id[i][s]
                    : 0;
            const int64_t stream_off = dsv4_stream_offset(n_stream, seq_id, state_size);
            const int32_t state_idx = (int32_t) (stream_off + pos%state_size);
            const auto it = std::find_if(persist_rows.begin(), persist_rows.end(), [state_idx](const persist_row & row) {
                return row.dst == state_idx;
            });
            if (it == persist_rows.end()) {
                persist_rows.push_back({ state_idx, i, pos });
            } else if (pos > it->pos) {
                it->src = i;
                it->pos = pos;
            }

            if ((pos + 1) % ratio != 0) {
                continue;
            }

            const llama_pos source_start = pos + 1 - ratio;
            const int64_t cache_off = dsv4_stream_offset(n_stream, seq_id, kv_size);
            plan.state_write_idxs.push_back(cache_off + pos/ratio);
            plan.state_write_pos.push_back((int32_t) source_start);

            if (overlap) {
                const llama_pos prev_start = source_start - ratio;
                for (uint32_t j = 0; j < ratio; ++j) {
                    overlap_prev_reads.push_back(state_source_idx(seq_id, prev_start + j));
                }
                for (uint32_t j = 0; j < ratio; ++j) {
                    overlap_cur_reads.push_back(state_source_idx(seq_id, source_start + j));
                }
            } else {
                for (uint32_t j = 0; j < ratio; ++j) {
                    plan.state_read_idxs.push_back(state_source_idx(seq_id, source_start + j));
                }
            }
        }
    }

    if (ratio == llama_context::dsv4_runtime::CSA_RATIO && plan.state_write_idxs.empty() && !plan.state_pos.empty()) {
        const llama_seq_id seq_id0 =
                batch.n_seq_id != nullptr && batch.seq_id != nullptr && batch.n_seq_id[0] > 0 && batch.seq_id[0] != nullptr
                ? batch.seq_id[0][0]
                : 0;
        const uint32_t source_idx = (uint32_t) state_source_idx(seq_id0, batch.pos[0]);
        const int64_t cache_off = std::max<int64_t>(0, dsv4_stream_offset(n_stream, seq_id0, kv_size));
        plan.state_write_idxs.push_back(cache_off + (int64_t) kv_size - 1);
        plan.state_write_pos.push_back(0);

        if (overlap) {
            for (uint32_t j = 0; j < ratio; ++j) {
                overlap_prev_reads.push_back(source_idx);
                overlap_cur_reads.push_back(source_idx);
            }
        } else {
            for (uint32_t j = 0; j < ratio; ++j) {
                plan.state_read_idxs.push_back(source_idx);
            }
        }
    }

    if (overlap) {
        plan.state_read_idxs.reserve(overlap_prev_reads.size() + overlap_cur_reads.size());
        plan.state_read_idxs.insert(plan.state_read_idxs.end(), overlap_prev_reads.begin(), overlap_prev_reads.end());
        plan.state_read_idxs.insert(plan.state_read_idxs.end(), overlap_cur_reads.begin(), overlap_cur_reads.end());
    }

    plan.n_kv = GGML_PAD(plan.n_kv, 256u);

    std::sort(persist_rows.begin(), persist_rows.end(), [](const persist_row & a, const persist_row & b) {
        return a.dst < b.dst;
    });

    for (const persist_row & row : persist_rows) {
        plan.state_persist_src_idxs.push_back(row.src);
        plan.state_persist_dst_idxs.push_back(row.dst);
    }

    if (plan.n_kv == 0) {
        plan.n_kv = GGML_PAD(1, 256u);
    }

    return plan;
}

template<typename T>
static void dsv4_set_input_tensor(ggml_tensor * tensor, const std::vector<T> & values) {
    if (tensor == nullptr || tensor->buffer == nullptr || values.empty()) {
        return;
    }
    ggml_backend_tensor_set(tensor, values.data(), 0, values.size()*sizeof(T));
}

static void dsv4_set_mask_tensor(
        ggml_tensor * tensor,
        const llama_context::dsv4_runtime::comp_plan & plan,
        int32_t n_tokens) {
    if (tensor == nullptr) {
        return;
    }

    if (tensor->buffer == nullptr) {
        return;
    }

    const int64_t width = tensor->ne[0];
    const int64_t height = tensor->ne[1];
    auto type = tensor->type;
    GGML_ASSERT(type == GGML_TYPE_F16 || type == GGML_TYPE_F32);

    //printf("%s: preparing mask %s of type %s with %ld x %ld entries\n", __func__, tensor->name, ggml_type_name(type), tensor->ne[0], tensor->ne[1]);
    if (type == GGML_TYPE_F16) {
        auto h_inf = ggml_fp32_to_fp16(-INFINITY);
        auto h_zero = ggml_fp32_to_fp16(0.0f);
        std::vector<ggml_fp16_t> storage((size_t) width*height, h_inf);
        for (int32_t i = 0; i < n_tokens; ++i) {
            const int32_t n_visible = i < (int32_t) plan.n_visible.size() ? plan.n_visible[(size_t) i] : 0;
            //if (i == 0) printf("    n_visible = %d\n", n_visible);
            for (int32_t j = 0; j < n_visible && j < width; ++j) {
                storage[(size_t) i*width + j] = h_zero;
            }
        }
        ggml_backend_tensor_set(tensor, storage.data(), 0, storage.size()*sizeof(ggml_fp16_t));
    } else {
        std::vector<float> storage((size_t) width*height, -INFINITY);
        for (int32_t i = 0; i < n_tokens; ++i) {
            const int32_t n_visible = i < (int32_t) plan.n_visible.size() ? plan.n_visible[(size_t) i] : 0;
            for (int32_t j = 0; j < n_visible && j < width; ++j) {
                storage[(size_t) i*width + j] = 0.0f;
            }
        }
        ggml_backend_tensor_set(tensor, storage.data(), 0, storage.size()*sizeof(float));
    }
}

bool llama_context::ensure_dsv4_cache_tensors() {
    const int32_t n_layer = model.hparams.n_layer;
    const int64_t n_embd_head = model.hparams.n_embd_head_k(0);
    const int64_t n_indexer_head = model.hparams.indexer_head_size;
    const uint32_t n_stream = std::max<uint32_t>(1, cparams.n_seq_max);
    const uint32_t csa_kv = GGML_PAD(dsv4_comp_size(cparams.n_ctx, dsv4_runtime::CSA_RATIO), 256u);
    const uint32_t hca_kv = GGML_PAD(dsv4_comp_size(cparams.n_ctx, dsv4_runtime::HCA_RATIO), 256u);

    if (!dsv4_validate_cache_type(kv_self.type_k, n_embd_head, "raw/CSA/HCA") ||
        !dsv4_validate_cache_type(cparams.idx_type_k, n_indexer_head, "LID")) {
        return false;
    }

    if (dsv4.cache.cache_ctx != nullptr &&
        (int32_t) dsv4.cache.csa_k.size() == n_layer &&
        dsv4.cache.n_stream == n_stream) {
        return true;
    }

    free_dsv4_cache_tensors();

    ggml_init_params params = {
        /*.mem_size   =*/ (size_t) (16 * std::max(1, n_layer)) * ggml_tensor_overhead(),
        /*.mem_buffer =*/ nullptr,
        /*.no_alloc   =*/ true,
    };

    dsv4.cache.cache_ctx = ggml_init(params);
    if (dsv4.cache.cache_ctx == nullptr) {
        LLAMA_LOG_ERROR("%s: failed to allocate DSV4 cache context\n", __func__);
        return false;
    }

    auto & cache = dsv4.cache;
    cache.n_stream = n_stream;
    cache.csa_k.resize((size_t) n_layer, nullptr);
    cache.hca_k.resize((size_t) n_layer, nullptr);
    cache.lid_k.resize((size_t) n_layer, nullptr);
    cache.csa_state_kv.resize((size_t) n_layer, nullptr);
    cache.csa_state_score.resize((size_t) n_layer, nullptr);
    cache.hca_state_kv.resize((size_t) n_layer, nullptr);
    cache.hca_state_score.resize((size_t) n_layer, nullptr);
    cache.lid_state_kv.resize((size_t) n_layer, nullptr);
    cache.lid_state_score.resize((size_t) n_layer, nullptr);

    auto alloc_tensor = [&](ggml_tensor * tensor, ggml_backend_buffer_type_t buft) -> bool {
        const size_t tensor_bytes = ggml_backend_buft_get_alloc_size(buft, tensor);
        ggml_backend_buffer_t buf = ggml_backend_buft_alloc_buffer(buft, tensor_bytes);
        if (buf == nullptr) {
            return false;
        }
        ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_COMPUTE);
        ggml_backend_tensor_alloc(buf, tensor, ggml_backend_buffer_get_base(buf));
        ggml_backend_buffer_clear(buf, 0);
        cache.cache_bufs.push_back(buf);
        return true;
    };

    for (int32_t il = 0; il < n_layer; ++il) {
        const uint32_t ratio = model.hparams.dsv4_compress_ratios[(size_t) il];
        ggml_backend_buffer_type_t buft = llama_dsv4_layer_buft(*this, il);

        if (ratio == dsv4_runtime::CSA_RATIO) {
            cache.csa_k[(size_t) il] = ggml_new_tensor_3d(cache.cache_ctx, kv_self.type_k, n_embd_head, csa_kv*n_stream, 1);
            cache.lid_k[(size_t) il] = ggml_new_tensor_3d(cache.cache_ctx, cparams.idx_type_k, n_indexer_head, csa_kv*n_stream, 1);
            cache.csa_state_kv[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, 2*n_embd_head, 2*dsv4_runtime::CSA_RATIO*n_stream);
            cache.csa_state_score[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, 2*n_embd_head, 2*dsv4_runtime::CSA_RATIO*n_stream);
            cache.lid_state_kv[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, 2*n_indexer_head, 2*dsv4_runtime::CSA_RATIO*n_stream);
            cache.lid_state_score[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, 2*n_indexer_head, 2*dsv4_runtime::CSA_RATIO*n_stream);

            if (!alloc_tensor(cache.csa_k[(size_t) il], buft) ||
                !alloc_tensor(cache.lid_k[(size_t) il], buft) ||
                !alloc_tensor(cache.csa_state_kv[(size_t) il], buft) ||
                !alloc_tensor(cache.csa_state_score[(size_t) il], buft) ||
                !alloc_tensor(cache.lid_state_kv[(size_t) il], buft) ||
                !alloc_tensor(cache.lid_state_score[(size_t) il], buft)) {
                LLAMA_LOG_ERROR("%s: failed to allocate DSV4 CSA/LID buffers for layer %d\n", __func__, il);
                free_dsv4_cache_tensors();
                return false;
            }
        } else if (ratio == dsv4_runtime::HCA_RATIO) {
            cache.hca_k[(size_t) il] = ggml_new_tensor_3d(cache.cache_ctx, kv_self.type_k, n_embd_head, hca_kv*n_stream, 1);
            cache.hca_state_kv[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, n_embd_head, dsv4_runtime::HCA_RATIO*n_stream);
            cache.hca_state_score[(size_t) il] = ggml_new_tensor_2d(cache.cache_ctx, GGML_TYPE_F32, n_embd_head, dsv4_runtime::HCA_RATIO*n_stream);

            if (!alloc_tensor(cache.hca_k[(size_t) il], buft) ||
                !alloc_tensor(cache.hca_state_kv[(size_t) il], buft) ||
                !alloc_tensor(cache.hca_state_score[(size_t) il], buft)) {
                LLAMA_LOG_ERROR("%s: failed to allocate DSV4 HCA buffers for layer %d\n", __func__, il);
                free_dsv4_cache_tensors();
                return false;
            }
        }
    }

    auto bytes = [](const auto & tensors) {
        size_t total = 0;
        for (const ggml_tensor * tensor : tensors) {
            if (tensor != nullptr) {
                total += ggml_nbytes(tensor);
            }
        }
        return total;
    };

    const size_t csa_k_bytes = bytes(cache.csa_k);
    const size_t hca_k_bytes = bytes(cache.hca_k);
    const size_t lid_k_bytes = bytes(cache.lid_k);
    const size_t csa_state_bytes = bytes(cache.csa_state_kv) + bytes(cache.csa_state_score);
    const size_t hca_state_bytes = bytes(cache.hca_state_kv) + bytes(cache.hca_state_score);
    const size_t lid_state_bytes = bytes(cache.lid_state_kv) + bytes(cache.lid_state_score);

    LLAMA_LOG_INFO("%s: DSV4 cache: CSA K=%7.2f MiB (%s), HCA K=%7.2f MiB (%s), LID K=%7.2f MiB (%s), states=%7.2f MiB, total=%7.2f MiB, streams=%u\n",
            __func__,
            (float) csa_k_bytes / (1024.0f * 1024.0f), ggml_type_name(kv_self.type_k),
            (float) hca_k_bytes / (1024.0f * 1024.0f), ggml_type_name(kv_self.type_k),
            (float) lid_k_bytes / (1024.0f * 1024.0f), ggml_type_name(cparams.idx_type_k),
            (float) (csa_state_bytes + hca_state_bytes + lid_state_bytes) / (1024.0f * 1024.0f),
            (float) (csa_k_bytes + hca_k_bytes + lid_k_bytes + csa_state_bytes + hca_state_bytes + lid_state_bytes) / (1024.0f * 1024.0f),
            n_stream);

    return true;
}

void llama_context::free_dsv4_cache_tensors() {
    auto release_vector = [](auto & v) {
        using vec_type = std::decay_t<decltype(v)>;
        vec_type().swap(v);
    };

    for (ggml_backend_buffer_t buf : dsv4.cache.cache_bufs) {
        if (buf != nullptr) {
            ggml_backend_buffer_free(buf);
        }
    }
    release_vector(dsv4.cache.cache_bufs);
    release_vector(dsv4.cache.csa_k);
    release_vector(dsv4.cache.hca_k);
    release_vector(dsv4.cache.lid_k);
    release_vector(dsv4.cache.csa_state_kv);
    release_vector(dsv4.cache.csa_state_score);
    release_vector(dsv4.cache.hca_state_kv);
    release_vector(dsv4.cache.hca_state_score);
    release_vector(dsv4.cache.lid_state_kv);
    release_vector(dsv4.cache.lid_state_score);
    dsv4.cache.n_stream = 1;
    if (dsv4.cache.cache_ctx != nullptr) {
        ggml_free(dsv4.cache.cache_ctx);
        dsv4.cache.cache_ctx = nullptr;
    }
}

void llama_reset_dsv4_state(llama_context * ctx, int32_t seq_id) {
    if (ctx == nullptr) {
        return;
    }

    const uint32_t n_stream = std::max<uint32_t>(1, ctx->dsv4.cache.n_stream);
    if (seq_id >= (llama_seq_id) n_stream) {
        LLAMA_LOG_ERROR("%s: DSV4 seq_id %d is outside stream range %u\n", __func__, seq_id, n_stream);
        return;
    }

    if (seq_id < 0) {
        for (ggml_backend_buffer_t buf : ctx->dsv4.cache.cache_bufs) {
            ggml_backend_buffer_clear(buf, 0);
        }
        return;
    }

    auto clear_tensor = [seq_id, n_stream](ggml_tensor * tensor) {
        if (tensor == nullptr) {
            return;
        }

        GGML_ASSERT(tensor->ne[1] % n_stream == 0);
        const size_t row_bytes = tensor->nb[1];
        const size_t rows_per_stream = (size_t) tensor->ne[1] / n_stream;
        const size_t offset = (size_t) seq_id * rows_per_stream * row_bytes;
        const size_t bytes = rows_per_stream * row_bytes;
        std::vector<uint8_t> zeros(bytes, 0);
        ggml_backend_tensor_set(tensor, zeros.data(), offset, bytes);
    };

    for (ggml_tensor * tensor : ctx->dsv4.cache.csa_k) clear_tensor(tensor);
    for (ggml_tensor * tensor : ctx->dsv4.cache.hca_k) clear_tensor(tensor);
    for (ggml_tensor * tensor : ctx->dsv4.cache.lid_k) clear_tensor(tensor);
    for (ggml_tensor * tensor : ctx->dsv4.cache.csa_state_kv) clear_tensor(tensor);
    for (ggml_tensor * tensor : ctx->dsv4.cache.csa_state_score) clear_tensor(tensor);
    for (ggml_tensor * tensor : ctx->dsv4.cache.hca_state_kv) clear_tensor(tensor);
    for (ggml_tensor * tensor : ctx->dsv4.cache.hca_state_score) clear_tensor(tensor);
    for (ggml_tensor * tensor : ctx->dsv4.cache.lid_state_kv) clear_tensor(tensor);
    for (ggml_tensor * tensor : ctx->dsv4.cache.lid_state_score) clear_tensor(tensor);
}

bool llama_prepare_dsv4_graph_inputs(llama_context & lctx, const llama_batch & batch, bool set_tensors, bool reserve_plan) {
    if (lctx.model.arch != LLM_ARCH_DEEPSEEK4) {
        return true;
    }

    if (!dsv4_validate_batch_seq_ids(lctx, batch)) {
        return false;
    }

    if (!lctx.ensure_dsv4_cache_tensors()) {
        return false;
    }

    const uint32_t cache_n_stream = std::max<uint32_t>(1, lctx.dsv4.cache.n_stream);
    const uint32_t csa_kv_size = dsv4_cache_kv_size(lctx.dsv4.cache.csa_k)/cache_n_stream;
    const uint32_t hca_kv_size = dsv4_cache_kv_size(lctx.dsv4.cache.hca_k)/cache_n_stream;
    const uint32_t lid_kv_size = dsv4_cache_kv_size(lctx.dsv4.cache.lid_k)/cache_n_stream;
    const uint32_t csa_state_size = dsv4_cache_state_size(lctx.dsv4.cache.csa_state_kv)/cache_n_stream;
    const uint32_t hca_state_size = dsv4_cache_state_size(lctx.dsv4.cache.hca_state_kv)/cache_n_stream;
    const uint32_t lid_state_size = dsv4_cache_state_size(lctx.dsv4.cache.lid_state_kv)/cache_n_stream;

    const auto build_plan = [&](uint32_t ratio, bool overlap, uint32_t state_size, uint32_t kv_size, uint32_t n_stream) {
        return reserve_plan
            ? dsv4_build_reserve_comp_plan(batch, ratio, overlap, state_size, kv_size, n_stream)
            : dsv4_build_comp_plan(batch, ratio, overlap, state_size, kv_size, n_stream);
    };

    lctx.dsv4.raw = {};
    if (!reserve_plan && !dsv4_build_raw_context(lctx, batch, lctx.dsv4.raw)) {
        return false;
    }

    //auto tim1 = ggml_time_us();
    lctx.dsv4.csa_plan = build_plan(llama_context::dsv4_runtime::CSA_RATIO, true, csa_state_size, csa_kv_size, cache_n_stream);
    lctx.dsv4.hca_plan = build_plan(llama_context::dsv4_runtime::HCA_RATIO, false, hca_state_size, hca_kv_size, cache_n_stream);
    lctx.dsv4.lid_plan = build_plan(llama_context::dsv4_runtime::CSA_RATIO, true, lid_state_size, lid_kv_size, cache_n_stream);
    lctx.dsv4.csa_ctx = dsv4_build_comp_context(batch, cache_n_stream, lctx.dsv4.csa_plan.n_kv);
    lctx.dsv4.hca_ctx = dsv4_build_comp_context(batch, cache_n_stream, lctx.dsv4.hca_plan.n_kv);
    lctx.dsv4.lid_ctx = dsv4_build_comp_context(batch, cache_n_stream, lctx.dsv4.lid_plan.n_kv);
    //auto tim2 = ggml_time_us();
    //fprintf(stderr, "%s: %ld us to buils plans\n", __func__, tim2-tim1);

    if (!dsv4_validate_comp_plan("csa", batch, lctx.dsv4.csa_plan, llama_context::dsv4_runtime::CSA_RATIO, true, csa_state_size, csa_kv_size, cache_n_stream) ||
        !dsv4_validate_comp_plan("hca", batch, lctx.dsv4.hca_plan, llama_context::dsv4_runtime::HCA_RATIO, false, hca_state_size, hca_kv_size, cache_n_stream) ||
        !dsv4_validate_comp_plan("lid", batch, lctx.dsv4.lid_plan, llama_context::dsv4_runtime::CSA_RATIO, true, lid_state_size, lid_kv_size, cache_n_stream) ||
        !dsv4_validate_csa_lid_visibility(lctx, csa_kv_size, lid_kv_size)) {
        return false;
    }

    if (!set_tensors) {
        return true;
    }

    //tim1 = ggml_time_us();

    dsv4_set_input_tensor(lctx.dsv4.inputs.raw_k_write_src_idxs, lctx.dsv4.raw.write_src_idxs);
    dsv4_set_input_tensor(lctx.dsv4.inputs.raw_k_write_idxs, lctx.dsv4.raw.write_dst_idxs);
    dsv4_set_input_tensor(lctx.dsv4.inputs.raw_k_read_idxs, lctx.dsv4.raw.read_dst_idxs);

    auto set_comp = [&](llama_context::dsv4_runtime::comp_inputs & inputs, llama_context::dsv4_runtime::comp_plan & plan, bool set_mask) {
        dsv4_set_input_tensor(inputs.state_pos, plan.state_pos);
        dsv4_set_input_tensor(inputs.state_persist_src_idxs, plan.state_persist_src_idxs);
        dsv4_set_input_tensor(inputs.state_persist_dst_idxs, plan.state_persist_dst_idxs);
        dsv4_set_input_tensor(inputs.state_read_idxs, plan.state_read_idxs);
        dsv4_set_input_tensor(inputs.state_write_idxs, plan.state_write_idxs);
        dsv4_set_input_tensor(inputs.state_write_pos, plan.state_write_pos);
        if (set_mask) {
            dsv4_set_mask_tensor(inputs.kq_mask, plan, batch.n_tokens);
        }
    };

    set_comp(lctx.dsv4.inputs.csa, lctx.dsv4.csa_plan, true);
    set_comp(lctx.dsv4.inputs.hca, lctx.dsv4.hca_plan, true);
    set_comp(lctx.dsv4.inputs.lid, lctx.dsv4.lid_plan, false);

    //tim2 = ggml_time_us();
    //fprintf(stderr, "%s: setting input tensors took %ld us\n", __func__, tim2 - tim1);

    return true;
}
