Class BlockMHA

src/model.h:276–363

Transformer Block - Multi-Head Attention */

/* Transformer Block - Multi-Head Attention */
struct BlockMHA : public Block {
  BlockMHA(
    int layer_i,
    const std::shared_ptr<Config> config,
    const Tensor* rms_att_weight,
    const Tensor* rms_q_a_weight,
    const Tensor* rms_kv_a_weight,
    const Tensor* rms_ffn_weight,
    const Tensor* wq,
    const Tensor* sq,
    const Tensor* wq_a,
    const Tensor* sq_a,
    const Tensor* wkv_a,
    const Tensor* skv_a,
    const Tensor* wq_b,
    const Tensor* sq_b,
    const Tensor* wkv_b,
    const Tensor* skv_b,
    const Tensor* wo,
    const Tensor* so,
    const Tensor* w1,
    const Tensor* s1,
    const Tensor* w2,
    const Tensor* s2,
    const Tensor* w3,
    const Tensor* s3,
    const Tensor* shared_w1,
    const Tensor* shared_s1,
    const Tensor* shared_w2,
    const Tensor* shared_s2,
    const Tensor* shared_w3,
    const Tensor* shared_s3,
    const Tensor* moegate,
    const Tensor* moegate_bias
  );
  ~BlockMHA() override;

  float* rms_q_a_weight() const { return _rms_q_a_weight ? static_cast<float*>(_rms_q_a_weight->data) : nullptr; }
  float* rms_kv_a_weight() const { return _rms_kv_a_weight ? static_cast<float*>(_rms_kv_a_weight->data) : nullptr; }
  std::optional<QTensor> wq() const { return _wq; }
  std::optional<QTensor> wq_a() const { return _wq_a; }
  std::optional<QTensor> wq_b() const { return _wq_b; }
  std::optional<QTensor> wkv_a() const { return _wkv_a; }
  std::optional<QTensor> wkv_b() const { return _wkv_b; }
  std::optional<QTensor> wo() const { return _wo; }
  f16_t* key_cache() const { return _key_cache; }
  f16_t* key_cache(int pos) const { return _key_cache + pos * _config->head_dim * _config->n_heads; }
  f16_t* value_cache() const { return _value_cache; }
  f16_t* value_cache(int pos) const { return _value_cache + pos * _config->v_head_dim * _config->n_heads; }

  double active_bytes(size_t pos) const override;

protected:
  void attention_impl(
    InferenceState& s, int pos, int kv_sink, int kv_pos, int kv_len
  ) const override;

private: