Method BlockMLA

src/model.cpp:518–620

BlockMLA::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
) : Block(layer_i, config, rms_att_weight, rms_ffn_weight, w1, s1, w2, s2, w3, s3, shared_w1, shared_s1, shared_w2, shared_s2, shared_w3, shared_s3, moegate, moegate_bias) {

  bool need_block_scales = _config->weight_quant == Quant::F8E5M2;
  int b0 = config->block_size[0];
  int b1 = config->block_size[1];

  _rms_q_a_weight = check_tensor(
    rms_q_a_weight, Quant::F32, {config->q_lora_rank, 0, 0, 0}, __LINE__
  );
  _rms_kv_a_weight = check_tensor(
    rms_kv_a_weight, Quant::F32, {config->kv_lora_rank, 0, 0, 0}, __LINE__ // Only norm latent part
  );

  _wq_a = check_tensor(
    wq_a, config->weight_quant, {config->q_lora_rank, config->dim, 0, 0}, __LINE__
  );
  _wkv_a = check_tensor(
    wkv_a, config->weight_quant, {config->kv_lora_rank + config->qk_rope_head_dim, config->dim, 0, 0}, __LINE__
  );
  _wc = check_tensor(
    wc, config->weight_quant, {config->n_heads * config->kv_lora_rank, config->q_lora_rank, 0, 0}, __LINE__
  );
  _wq_rope_b = check_tensor(
    wq_rope_b, config->weight_quant, {config->n_heads * config->qk_rope_head_dim, config->q_lora_rank, 0, 0}, __LINE__
  );