| 7 | namespace mllm::nn { |
| 8 | |
| 9 | DynamicCache::DynamicCache(int32_t layer_nums, int32_t q_heads, int32_t kv_heads, int32_t kv_dims, bool use_fa2) |
| 10 | : layer_nums_(layer_nums), q_heads_(q_heads), kv_heads_(kv_heads), kv_dims_(kv_dims), use_fa2_(use_fa2) { |
| 11 | for (int i = 0; i < layer_nums_; ++i) { |
| 12 | k_cache_.emplace_back(Tensor::nil()); |
| 13 | v_cache_.emplace_back(Tensor::nil()); |
| 14 | } |
| 15 | } |
| 16 | |
| 17 | std::array<Tensor, 2> DynamicCache::updateKVCache(int32_t layer_idx, Tensor k, Tensor v) { |
| 18 | // The input should be [B, H, S, D] |
nothing calls this directly
no outgoing calls
no test coverage detected