calculet-npu-research-archive/source/baseline-expanded/calrt-sdk-0.7.6/usr/local/include/calrt/calrt_kv.h

#pragma once

#include <array>
#include <cstddef>
#include <cstdint>
#include <deque>
#include <memory>
#include <unordered_map>
#include <vector>

#include "calrt_calbin.h"
#include "calrt_utils.h"
#include "calrt_vdevice.h"

namespace calrt {

typedef int32_t KvPos;
typedef int32_t KvSeqId;

struct KvCellAddr_s {
    int64_t src_offset; // offset from base_addr
    int64_t dst_offset; // offset from base_addr
};

enum class KvBatchMode_e : int32_t {
    ONLY_VALID = 0,
    ALL = 1, // 0:only mv valid batch ; 1:mv 16 batch ;recommond threhold 5 batch
    AUTO
};

enum class KvUpdateAlgo_e : int32_t { V_CACHE_MOVE = 0, K_CACHE_ROPE = 1 };

enum class KvDataType_e : int32_t { BF16 = 0, S8 = 1, INVALID };

#define D0 16
// kv_cache memory layout:
//  batch=1: K&V cache [layer,head,1,          max_seq_len/d0,head_dim/d2,d0,d2]
//  batch>1: k_cache   [layer,head,batch,      max_seq_len/d0,head_dim/d2,d0,d2]
//           v_cavhe   [layer,head,max_seq_len,batch/d0,      head_dim/d2,d0,d2]
//  k_cache and v_cache have completely continous memory layout
struct kv_update_config_p16_s {
    int64_t base_addr;     // k or v_cache base addr,to compute offset
    int64_t k_head_stride; // k_cache: size of each head, reduce mul [max_batch,max_seq/d0,head_dim/d2,d0,d2]*byte_size
    int64_t v_max_seq_len_stride; // v_cache: size of each max_seq_len,
                                  // reduce mul [max_batch/d0,head_dim/d2,d0,d2]*byte_size
    int64_t v_head_stride;        // v_cache: size of each head,
                                  // reduce mul [max_seq_len,max_batch/d0,head_dim/d2,d0,d2]*byte_size
    KvUpdateAlgo_e kv_update_algo;
    KvBatchMode_e data_mv_mode;

    int32_t model_batch_num; // to specify one batch or mul batch memory of v_cache,1-onebatch;16-16batch
    int32_t n_layer;         // all layer
    int32_t n_head;          // all head
    int32_t n_batch;         // only valid batch
    int32_t n_tokens_len;    // valid n_seq_len

    int32_t head_dim_div_d2; // assert head_dim_div_d2 % 2 == 0

    KvDataType_e sin_cos_table_data_type; // support bf16  first; 0-bf16
    int16_t *sin_table;
    int16_t *cos_table;

    KvDataType_e kv_data_type; // support bf16 first for k_cache, bf16 s8 for v_cache, 0-bf16,1-s8
    KvCellAddr_s *addrs;       // [n_batch]

    kv_update_config_p16_s() = default;
    ~kv_update_config_p16_s() { free(addrs); }

    kv_update_config_p16_s(const kv_update_config_p16_s &) = delete;
    kv_update_config_p16_s &operator=(const kv_update_config_p16_s &) = delete;
    kv_update_config_p16_s(kv_update_config_p16_s &&) = default;
    kv_update_config_p16_s &operator=(kv_update_config_p16_s &&) = default;
};

struct KvHwBatch {
    KvHwBatch(int32_t batch);
    KvSeqId seq_id;
    int32_t batch_id;
    size_t seq_len;
};

struct KvSeqInfo {
    KvSeqId seq_id;
    size_t seq_len;
};

struct CalrtSeqShift {
    KvSeqId seq_id;
    KvPos dst_pos;
    KvPos src_pos;
    KvPos len;
};

class CALRT_API KvManager {
  public:
    KvManager(calrt::VirtualDevice* vdev, calrt::Calbin* calbin);

    bool canShift();

    // check if we have avaliable hardware resource
    bool canAllocate(size_t seq_num);
    void Allocate(KvSeqInfo &seq_info);

    void Free(KvSeqId seq_id);

    void RemoveTokensAtEnd(KvSeqId seq_id, size_t n_tokens);

    bool isExist(KvSeqId seq_id);
    // keep track of kv state
    bool Apply(KvSeqInfo &seq_info);

    // if batch_mode == ALL, only seqs[0] will be used
    void DoShift(KvBatchMode_e batch_mode, std::vector<CalrtSeqShift> &seqs);

    KvPos GetSeqLen(KvSeqId seq_id) const;

    void Clear();

    KvPos seqPosMin(KvSeqId seq_id) const;
    KvPos seqPosMax(KvSeqId seq_id) const;

  private:
    VirtualDevice* vdev;
    Calbin* calbin;
    CalbinLLM_s llm_spec;

    std::deque<std::unique_ptr<KvHwBatch>> free_hw_batch;

    std::unordered_map<KvSeqId, std::unique_ptr<KvHwBatch>> seq_to_hw_batch;

    std::vector<float> inv_freq;
    std::unordered_map<int, std::vector<int16_t>> sin_cache; // bf16
    std::unordered_map<int, std::vector<int16_t>> cos_cache;

    template <float (*tri)(float)> int16_t *get_table(std::unordered_map<int, std::vector<int16_t>> &cache, int offset);

    kv_update_config_p16_s build_kv_config(KvUpdateAlgo_e k_or_v, KvBatchMode_e batch_mode, int32_t max_batch,
                                           int32_t valid_batch, std::vector<CalrtSeqShift> &seqs);
};

// std::unique_ptr<KvManager> kv_manager_init_from_spec(const CalbinLLM_s &calbin_llm);

static int16_t float_to_bfloat16(float f) {
    uint32_t float_bits;
    std::memcpy(&float_bits, &f, sizeof(float)); // Copy float bits to uint32_t

    // Bfloat16 is essentially the upper 16 bits of the float
    // with the lower 16 bits of the mantissa truncated.
    int16_t bfloat16_bits = static_cast<int16_t>(float_bits >> 16);

    return bfloat16_bits;
}

template <float (*tri)(float)>
int16_t *KvManager::get_table(std::unordered_map<int, std::vector<int16_t>> &cache, int offset) {
    auto it = cache.find(offset);

    if (it != cache.end()) {
        return it->second.data();
    }

    std::vector<int16_t> table(inv_freq.size() * 2);

    for (size_t i = 0; i < inv_freq.size(); ++i) {
        float theta = static_cast<float>(-offset) * inv_freq[i]; // use negative pos because we use this for kv shift
        float val = tri(theta);

        int16_t bf16 = float_to_bfloat16(val);

        table[i] = bf16;
        table[inv_freq.size() + i] = bf16;
    }

    auto [ins_it, _] = cache.emplace(offset, std::move(table));
    return ins_it->second.data();
}

} // namespace calrt