Method _forward_cpu

src/infer.cpp:1265–1317

}

void Model::_forward_cpu(InferenceState& s, int token, int pos, InferenceMode mode) {
  const Config& c = *config;

  // copy the token embedding into `x`
  PROFILE(_copy_embedding(s, token));

  // When decoding past the context length, keep the first few tokens in the KV cache
  // untouched as "attention sinks" while replacing the rest in ring order.
  // See StreamingLLM (https://arxiv.org/pdf/2309.17453) for more.
  int original_max_position = c.rs_original_max_position_embeddings;
  int kv_sink = pos >= original_max_position ? KV_SINKS : 0;
  int kv_pos = kv_sink + (pos - kv_sink) % (original_max_position - kv_sink);
  int kv_len = pos >= original_max_position ? original_max_position : pos + 1;

  // forward all layers in order
  for (auto b : blocks) {
    b->block(s, pos, kv_sink, kv_pos, kv_len);
  }

  if (mode == InferenceMode::HYDRATE_KV_CACHE) {
    // only hydrate the KV cache and don't compute output logits
    return;
  }

  // final layer norm
  switch (c.norm_type) {
    case LayerNormType::RMSNorm: {
      rmsnorm(s.x(), s.x(), static_cast<float*>(rms_final_weight->data), c.dim, c.norm_eps);
      break;
    }
  }

  // classifier into logits
  {
    PROFILE_BLOCK(lm_head);
    switch (c.weight_quant) {
      case Quant::F32:
      case Quant::F16: {
        matmul_unscaled(s.logits(), s.x(), *wcls);
        break;
      }
      case Quant::F8E5M2:
      case Quant::Q2_K:
      case Quant::Q3_K: {
        matmul(s.logits(), s.x(), *wcls, c.block_size.data(), scls, s.aqb());
        break;
      }
      default: {
        assert(false && "unsupported weight quantization");
      }
    }
  }
}