Method InferenceState

src/model.cpp:677–726

}

InferenceState::InferenceState(const std::shared_ptr<Config> config): 
  _config(config) {
  assert(config);
  _x = new float[config->dim]();
  _xb = new float[config->dim]();
  _xb2 = new float[std::max({
    config->dim, 
    config->n_heads * config->v_head_dim, 
    config->n_heads * config->kv_lora_rank
  })]();
  _hb = new float[std::max({
    config->dim, 
    config->hidden_dim
  })]();
  _hb2 = new float[config->hidden_dim]();
  if (config->q_lora_rank > 0) {
    _q_a = new float[config->q_lora_rank]();
  }
  _q = new float[config->n_heads * config->head_dim]();
  _kv_a = new float[config->kv_lora_rank + config->qk_rope_head_dim]();
  _kv_b = new float[config->n_heads * (config->head_dim-config->qk_rope_head_dim+config->v_head_dim)]();
  _ropebuf = new float[config->n_heads * config->qk_rope_head_dim]();
  _k = new float[config->n_heads * config->head_dim]();
  _v = new float[config->n_heads * config->v_head_dim]();
  _att = new float[config->n_heads * config->max_seq_len]();
  _logits = new float[config->vocab_size]();
  _logit_indices = new int[config->vocab_size]();
  for (int i = 0; i < config->vocab_size; i++){
    _logit_indices[i] = i;
  }
  if (config->use_mla) {
    _q_c = new float[config->n_heads * config->kv_lora_rank]();
    _q_rope = new float[config->n_heads * config->qk_rope_head_dim]();
  }
  if (config->n_routed_experts > 0) {
    _moe_weights = new float[config->n_routed_experts]();
    _active_experts = new int[config->n_active_routed]();
    _active_experts_weights = new float[config->n_active_routed]();
  }
  // TODO: consider dynamically resizing based on inputs
  size_t aqb_nitems = std::max({
    config->dim,
    config->moe_intermediate_size,
    config->n_heads * config->v_head_dim,
    config->n_heads * config->kv_lora_rank,
    config->hidden_dim
  });
  size_t aqb_nblocks = aqb_nitems / QK_K;
  _aqb = new uint8_t[aqb_nblocks * sizeof(block_q8_K)]();
}

InferenceState::~InferenceState() {
  if (_device == Device::CPU) {