Method _block_cpu

src/infer.cpp:811–932

// - Block KV cache is hydrated.
template <typename T>
void Block::_block_cpu(
  InferenceState& s,  // inference state
  int pos,            // index of the current token in the sequence
  int kv_sink,        // number of sink tokens currently in the KV cache
  int kv_pos,         // index of the current token in the kv cache, must be in [0..kv_len) since kv cache is a ring buffer
  int kv_len          // number of tokens in the kv cache that we will attend over
) const {
  const Config& c = *_config;

  // Attention pre-norm
  switch (c.norm_type) {
    case LayerNormType::RMSNorm: {
      rmsnorm(s.xb(), s.x(), rms_att_weight(), c.dim, c.norm_eps);
      break;
    }
  }

  // Attention output into `hb`
  attention_impl(s, pos, kv_sink, kv_pos, kv_len);

  // Residual back into `x`
  for (int i = 0; i < c.dim; ++i) {
    s.x()[i] += s.hb()[i];
  }

  // FFN pre-norm
  switch (c.norm_type) {
    case LayerNormType::RMSNorm: {
      rmsnorm(s.xb(), s.x(), rms_ffn_weight(), c.dim, c.norm_eps);
      break;
    }
  }

  if (c.n_routed_experts > 0 && moegate() != std::nullopt) {
    PROFILE_BLOCK(ffn_moe);
    // Block is a sparse MoE FFN layer
    PROFILE(matmul_unscaled(s.moe_weights(), s.xb(), *moegate()));
    moe_gate(
      s.active_experts_weights(), moegate_bias(), s.active_experts(), s.moe_weights(),
      c.n_routed_experts, c.n_active_routed, c.norm_topk_prob, c.routed_scaling_factor,
      c.scoring_func, c.topk_method, c.n_group, c.topk_group
    );
    for (int k = 0; k < c.n_active_routed; ++k) {
      int expert_index = s.active_experts()[k];
      // mix self.w2(F.silu(self.w1(x)) * self.w3(x))
      // Note this is a feedforward with a GLU, not a simple MLP.
      PROFILE(matmul_expert(s.hb(), s.xb(), *w1(), expert_index, c.block_size.data(), _s1, s.aqb()));
      PROFILE(matmul_expert(s.hb2(), s.xb(), *w3(), expert_index, c.block_size.data(), _s3, s.aqb()));
      switch (c.act) {
        case ActivationType::GELU: {
          for (int i = 0; i < c.moe_intermediate_size; ++i) {
            s.hb()[i] = gelu(s.hb()[i]) * s.hb2()[i];
          }
          break;
        }
        case ActivationType::SILU: {
          for (int i = 0; i < c.moe_intermediate_size; ++i) {
            s.hb()[i] = silu(s.hb()[i]) * s.hb2()[i];