Class BlockMLA

src/model.h:366–453

Transformer Block - Multi-Latent Attention */

/* Transformer Block - Multi-Latent Attention */
struct BlockMLA : public Block {
  BlockMLA(
    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_a,
    const Tensor* sq_a,
    const Tensor* wkv_a,
    const Tensor* skv_a,
    const Tensor* wo,
    const Tensor* so,
    const Tensor* wc,
    const Tensor* sc,
    const Tensor* wq_rope_b,
    const Tensor* sq_rope_b,
    const Tensor* wv_b,
    const Tensor* sv_b,
    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
  );
  ~BlockMLA() 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_a() const { return _wq_a; }
  std::optional<QTensor> wkv_a() const { return _wkv_a; }
  std::optional<QTensor> wo() const { return _wo; }
  std::optional<QTensor> wc() const { return _wc; }
  std::optional<QTensor> wq_rope_b() const { return _wq_rope_b; }
  std::optional<QTensor> wv_b() const { return _wv_b; }
  f16_t* kv_nope_cache() const { return _kv_nope_cache; }
  f16_t* kv_nope_cache(int pos) const { return _kv_nope_cache + pos * _config->kv_lora_rank; }
  f16_t* kv_rope_cache() const { return _kv_rope_cache; }
  f16_t* kv_rope_cache(int pos) const { return _kv_rope_cache + pos * _config->qk_rope_head_dim; }

  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:
  template <typename T>