Method _attention_impl

src/infer.cpp:935–1049

template<typename T>
void BlockMHA::_attention_impl(
  InferenceState& s, int pos, int kv_sink, int kv_pos, int kv_len
) const {
  PROFILE_BLOCK(attn_mha);
  const Config& c = *_config;

  // qkv matmuls for this position
  if (c.q_lora_rank > 0) {
    PROFILE(matmul(s.q_a(), s.xb(), *wq_a(), c.block_size.data(), _sq_a, s.aqb()));
    switch (c.norm_type) {
      case LayerNormType::RMSNorm: {
        rmsnorm(s.q_a(), s.q_a(), this->rms_q_a_weight(), c.q_lora_rank, c.norm_eps);
        break;
      }
    }
    PROFILE(matmul(s.q(), s.q_a(), *wq_b(), c.block_size.data(), _sq_b, s.aqb()));
  } else {
    PROFILE(matmul(s.q(), s.xb(), *wq(), c.block_size.data(), _sq, s.aqb()));
  }
  PROFILE(matmul(s.kv_a(), s.xb(), *wkv_a(), c.block_size.data(), _skv_a, s.aqb()));

  // Apply RoPE positional encoding
  int q_pe_offset = c.head_dim - c.qk_rope_head_dim;
  bool is_v3 = c.has_moegate_bias;
  for (int h = 0; h < c.n_heads; h++) {
    if (is_v3) {
      rope_v3(s.q(h) + q_pe_offset, c.qk_rope_head_dim, c.qk_rope_head_dim, pos, c.rope_theta);
    } else {
      rope(s.ropebuf(), s.q(h) + q_pe_offset, c.qk_rope_head_dim, c.qk_rope_head_dim, pos, c.rope_theta);
    }
  }
  int kv_pe_offset = c.kv_lora_rank;
  float* k_rope = s.kv_a() + kv_pe_offset;
  if (is_v3) {
    rope_v3(k_rope, c.qk_rope_head_dim, c.qk_rope_head_dim, pos, c.rope_theta);
  } else {
    rope(s.ropebuf(), k_rope, c.qk_rope_head_dim, c.qk_rope_head_dim, pos, c.rope_theta);
  }
  // rms norm to non-pe chunk of kv_a
  rmsnorm(s.kv_a(), s.kv_a(), this->rms_kv_a_weight(), c.kv_lora_rank, c.norm_eps);
  // un-compress the latent kv via multiplication with wkv_b
  int qk_nope_head_dim = c.head_dim - c.qk_rope_head_dim;
  PROFILE(matmul(s.kv_b(), s.kv_a(), *wkv_b(), c.block_size.data(), _skv_b, s.aqb()));
  // concatenate kv_b and k_rope in each head to build key heads
  for (int h = 0; h < c.n_heads; h++) {
    for (int i = 0; i < qk_nope_head_dim; i++) {
      s.k(h)[i] = s.kv_b(h)[i];
    }
    for (int i = 0; i < c.qk_rope_head_dim; i++) {
      s.k(h)[qk_nope_head_dim + i] = k_rope[i];
    }
  }
  // transfer value heads from kv_b
  for (int h = 0; h < c.n_heads; h++) {
    for (int i = 0; i < c.v_head_dim; i++) {
      s.v(h)[i] = s.kv_b(h)[qk_nope_head_dim + i];
    }
  }