|
| 1 | +#include "deepseek_v2_mla_attention.hpp" |
| 2 | + |
| 3 | +#include "../../global_state/global_state.hpp" |
| 4 | +#include "../../utils.hpp" |
| 5 | +#include "infinicore/ops.hpp" |
| 6 | +#include "infinicore/ops/cat.hpp" |
| 7 | +#include "infinicore/ops/pad.hpp" |
| 8 | + |
| 9 | +#include <cmath> |
| 10 | +#include <stdexcept> |
| 11 | + |
| 12 | +namespace infinilm::models::deepseek_v2 { |
| 13 | +namespace { |
| 14 | + |
| 15 | +float yarn_get_mscale(float scale, float mscale) { |
| 16 | + if (scale <= 1.0f) { |
| 17 | + return 1.0f; |
| 18 | + } |
| 19 | + return 0.1f * mscale * std::log(scale) + 1.0f; |
| 20 | +} |
| 21 | + |
| 22 | +} // namespace |
| 23 | + |
| 24 | +DeepseekV2MLAAttention::DeepseekV2MLAAttention(std::shared_ptr<infinilm::config::ModelConfig> model_config, |
| 25 | + size_t layer_idx, |
| 26 | + const infinicore::Device &device) { |
| 27 | + layer_idx_ = layer_idx; |
| 28 | + hidden_size_ = model_config->get<size_t>("hidden_size"); |
| 29 | + qk_nope_head_dim_ = model_config->get<size_t>("qk_nope_head_dim"); |
| 30 | + qk_rope_head_dim_ = model_config->get<size_t>("qk_rope_head_dim"); |
| 31 | + q_head_dim_ = qk_nope_head_dim_ + qk_rope_head_dim_; |
| 32 | + v_head_dim_ = model_config->get<size_t>("v_head_dim"); |
| 33 | + kv_lora_rank_ = model_config->get<size_t>("kv_lora_rank"); |
| 34 | + mla_head_dim_ = kv_lora_rank_ + qk_rope_head_dim_; |
| 35 | + |
| 36 | + if (model_config->get_or<size_t>("q_lora_rank", 0) != 0) { |
| 37 | + throw std::runtime_error("DeepseekV2MLAAttention: q_lora_rank is not supported yet"); |
| 38 | + } |
| 39 | + |
| 40 | + const auto &dtype{model_config->get_dtype()}; |
| 41 | + const size_t total_num_heads = model_config->get<size_t>("num_attention_heads"); |
| 42 | + const bool attention_bias = model_config->get_or<bool>("attention_bias", false); |
| 43 | + const double rms_norm_eps = model_config->get<double>("rms_norm_eps"); |
| 44 | + |
| 45 | + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); |
| 46 | + const int tp_rank = rank_info.tp_rank; |
| 47 | + const int tp_size = rank_info.tp_size; |
| 48 | + if ((total_num_heads < static_cast<size_t>(tp_size)) || (total_num_heads % static_cast<size_t>(tp_size) != 0)) { |
| 49 | + throw std::runtime_error("DeepseekV2MLAAttention: num_attention_heads must be divisible by tp_size"); |
| 50 | + } |
| 51 | + num_attention_heads_ = total_num_heads / static_cast<size_t>(tp_size); |
| 52 | + attention_backend_ = infinilm::global_state::get_infinilm_config().attention_backend; |
| 53 | + |
| 54 | + auto quantization_method = model_config->get_quantization_method(); |
| 55 | + INFINICORE_NN_MODULE_INIT(q_proj, hidden_size_, total_num_heads * q_head_dim_, quantization_method, false, dtype, device, tp_rank, tp_size); |
| 56 | + INFINICORE_NN_MODULE_INIT(kv_a_proj_with_mqa, hidden_size_, kv_lora_rank_ + qk_rope_head_dim_, attention_bias, dtype, device); |
| 57 | + INFINICORE_NN_MODULE_INIT(kv_a_layernorm, kv_lora_rank_, rms_norm_eps, dtype, device); |
| 58 | + INFINICORE_NN_MODULE_INIT(kv_b_proj, kv_lora_rank_, total_num_heads * (qk_nope_head_dim_ + v_head_dim_), quantization_method, false, dtype, device, tp_rank, tp_size); |
| 59 | + INFINICORE_NN_MODULE_INIT(o_proj, total_num_heads * v_head_dim_, hidden_size_, quantization_method, attention_bias, dtype, device, tp_rank, tp_size, rank_info.comm); |
| 60 | + |
| 61 | + const size_t max_position_embeddings = model_config->get<size_t>("max_position_embeddings"); |
| 62 | + const double rope_theta = model_config->get<double>("rope_theta"); |
| 63 | + rotary_emb_ = std::make_shared<infinicore::nn::RoPE>( |
| 64 | + qk_rope_head_dim_, qk_rope_head_dim_, max_position_embeddings, rope_theta, |
| 65 | + infinicore::nn::RoPE::Algo::GPT_J, dtype, device, nullptr); |
| 66 | + |
| 67 | + softmax_scale_ = 1.0f / std::sqrt(static_cast<float>(q_head_dim_)); |
| 68 | + auto &config_json = model_config->get_config_json(); |
| 69 | + if (config_json.contains("rope_scaling") && config_json["rope_scaling"].is_object()) { |
| 70 | + const auto &rope_scaling = config_json["rope_scaling"]; |
| 71 | + const float mscale_all_dim = rope_scaling.value("mscale_all_dim", 0.0f); |
| 72 | + if (mscale_all_dim != 0.0f) { |
| 73 | + const float scaling_factor = rope_scaling.value("factor", 1.0f); |
| 74 | + const float mscale = yarn_get_mscale(scaling_factor, mscale_all_dim); |
| 75 | + softmax_scale_ *= mscale * mscale; |
| 76 | + } |
| 77 | + } |
| 78 | + |
| 79 | + latent_attn_ = std::make_shared<infinilm::layers::attention::AttentionLayer>( |
| 80 | + num_attention_heads_, mla_head_dim_, softmax_scale_, 1, layer_idx_, |
| 81 | + kv_cache_k_scale_, kv_cache_v_scale_, attention_backend_); |
| 82 | + infinilm::layers::attention::init_kv_cache_quant_params( |
| 83 | + [this](const std::string &n, infinicore::nn::Parameter p) { this->register_parameter(n, std::move(p)); }, |
| 84 | + device, kv_cache_k_scale_, kv_cache_v_scale_); |
| 85 | +} |
| 86 | + |
| 87 | +infinicore::Tensor DeepseekV2MLAAttention::position_ids_for_rope_(const infinicore::Tensor &position_ids) const { |
| 88 | + auto pos_shape = position_ids->shape(); |
| 89 | + if (pos_shape.size() == 2) { |
| 90 | + return position_ids->narrow({{0, 0, 1}})->contiguous()->view({pos_shape[1]}); |
| 91 | + } |
| 92 | + if (pos_shape.size() == 1) { |
| 93 | + return position_ids->contiguous(); |
| 94 | + } |
| 95 | + throw std::runtime_error("DeepseekV2MLAAttention: unexpected position_ids shape"); |
| 96 | +} |
| 97 | + |
| 98 | +infinicore::Tensor DeepseekV2MLAAttention::kv_b_weight_3d_() const { |
| 99 | + return kv_b_proj_->weight()->view({num_attention_heads_, qk_nope_head_dim_ + v_head_dim_, kv_lora_rank_}); |
| 100 | +} |
| 101 | + |
| 102 | +infinicore::Tensor DeepseekV2MLAAttention::project_q_nope_to_latent_(const infinicore::Tensor &q_nope) const { |
| 103 | + const size_t ntokens = q_nope->shape()[0]; |
| 104 | + auto q_nope_by_head = q_nope->permute({1, 0, 2})->contiguous(); |
| 105 | + auto w_uk_t = kv_b_weight_3d_()->narrow({{1, 0, qk_nope_head_dim_}})->contiguous(); |
| 106 | + auto q_latent = infinicore::op::matmul(q_nope_by_head, w_uk_t); |
| 107 | + return q_latent->permute({1, 0, 2})->contiguous()->view({ntokens, num_attention_heads_, kv_lora_rank_}); |
| 108 | +} |
| 109 | + |
| 110 | +infinicore::Tensor DeepseekV2MLAAttention::project_latent_to_value_(const infinicore::Tensor &attn_output, |
| 111 | + size_t batch_size, |
| 112 | + size_t seq_len) const { |
| 113 | + const size_t ntokens = batch_size * seq_len; |
| 114 | + auto latent = attn_output->view({ntokens, num_attention_heads_, mla_head_dim_}) |
| 115 | + ->narrow({{2, 0, kv_lora_rank_}}) |
| 116 | + ->contiguous(); |
| 117 | + auto latent_by_head = latent->permute({1, 0, 2})->contiguous(); |
| 118 | + auto w_uv = kv_b_weight_3d_() |
| 119 | + ->narrow({{1, qk_nope_head_dim_, v_head_dim_}}) |
| 120 | + ->permute({0, 2, 1}) |
| 121 | + ->contiguous(); |
| 122 | + auto value = infinicore::op::matmul(latent_by_head, w_uv) |
| 123 | + ->permute({1, 0, 2}) |
| 124 | + ->contiguous() |
| 125 | + ->view({batch_size, seq_len, num_attention_heads_ * v_head_dim_}); |
| 126 | + return o_proj_->forward(value); |
| 127 | +} |
| 128 | + |
| 129 | +infinicore::Tensor DeepseekV2MLAAttention::forward(const infinicore::Tensor &positions, |
| 130 | + const infinicore::Tensor &hidden_states) const { |
| 131 | + if (::infinilm::backends::AttentionBackend::STATIC_ATTN == attention_backend_) { |
| 132 | + return forward_static_(positions, hidden_states); |
| 133 | + } |
| 134 | + return forward_paged_(positions, hidden_states); |
| 135 | +} |
| 136 | + |
| 137 | +infinicore::Tensor DeepseekV2MLAAttention::forward_static_(const infinicore::Tensor &position_ids, |
| 138 | + const infinicore::Tensor &hidden_states) const { |
| 139 | + auto shape = hidden_states->shape(); |
| 140 | + const size_t batch_size = shape[0]; |
| 141 | + const size_t seq_len = shape[1]; |
| 142 | + const size_t ntokens = batch_size * seq_len; |
| 143 | + auto hidden_states_mutable = hidden_states; |
| 144 | + |
| 145 | + auto q = q_proj_->forward(hidden_states_mutable)->view({ntokens, num_attention_heads_, q_head_dim_}); |
| 146 | + auto q_nope = q->narrow({{2, 0, qk_nope_head_dim_}})->contiguous(); |
| 147 | + auto q_pe = q->narrow({{2, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous(); |
| 148 | + |
| 149 | + auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable)->view({ntokens, kv_lora_rank_ + qk_rope_head_dim_}); |
| 150 | + auto compressed_kv = compressed->narrow({{1, 0, kv_lora_rank_}})->contiguous(); |
| 151 | + auto k_pe = compressed->narrow({{1, kv_lora_rank_, qk_rope_head_dim_}})->contiguous(); |
| 152 | + |
| 153 | + auto kv_norm = kv_a_layernorm_->forward(compressed_kv); |
| 154 | + auto pos_ids = position_ids_for_rope_(position_ids); |
| 155 | + q_pe = rotary_emb_->forward(q_pe, pos_ids, true); |
| 156 | + auto k_pe_rope = rotary_emb_->forward(k_pe->view({ntokens, 1, qk_rope_head_dim_}), pos_ids, true); |
| 157 | + |
| 158 | + auto q_latent = project_q_nope_to_latent_(q_nope); |
| 159 | + auto query_states = infinicore::op::cat({q_latent, q_pe}, 2)->view({batch_size, seq_len, num_attention_heads_, mla_head_dim_}); |
| 160 | + auto key_states = infinicore::op::cat({kv_norm->view({ntokens, 1, kv_lora_rank_}), k_pe_rope}, 2) |
| 161 | + ->view({batch_size, seq_len, 1, mla_head_dim_}); |
| 162 | + auto value_states = infinicore::op::pad(kv_norm->view({batch_size, seq_len, 1, kv_lora_rank_}), |
| 163 | + {0, static_cast<int>(qk_rope_head_dim_)}, "constant", 0.0); |
| 164 | + |
| 165 | + auto attn_output = latent_attn_->forward(query_states, key_states, value_states); |
| 166 | + return project_latent_to_value_(attn_output, batch_size, seq_len); |
| 167 | +} |
| 168 | + |
| 169 | +infinicore::Tensor DeepseekV2MLAAttention::forward_paged_(const infinicore::Tensor &position_ids, |
| 170 | + const infinicore::Tensor &hidden_states) const { |
| 171 | + auto shape = hidden_states->shape(); |
| 172 | + const size_t batch_size = shape[0]; |
| 173 | + const size_t seq_len = shape[1]; |
| 174 | + ASSERT_EQ(batch_size, 1); |
| 175 | + auto hidden_states_mutable = hidden_states; |
| 176 | + |
| 177 | + auto q = q_proj_->forward(hidden_states_mutable)->view({seq_len, num_attention_heads_, q_head_dim_}); |
| 178 | + auto q_nope = q->narrow({{2, 0, qk_nope_head_dim_}})->contiguous(); |
| 179 | + auto q_pe = q->narrow({{2, qk_nope_head_dim_, qk_rope_head_dim_}})->contiguous(); |
| 180 | + |
| 181 | + auto compressed = kv_a_proj_with_mqa_->forward(hidden_states_mutable)->view({seq_len, kv_lora_rank_ + qk_rope_head_dim_}); |
| 182 | + auto compressed_kv = compressed->narrow({{1, 0, kv_lora_rank_}})->contiguous(); |
| 183 | + auto k_pe = compressed->narrow({{1, kv_lora_rank_, qk_rope_head_dim_}})->contiguous(); |
| 184 | + |
| 185 | + auto kv_norm = kv_a_layernorm_->forward(compressed_kv); |
| 186 | + auto pos_ids = position_ids_for_rope_(position_ids); |
| 187 | + q_pe = rotary_emb_->forward(q_pe, pos_ids, true); |
| 188 | + auto k_pe_rope = rotary_emb_->forward(k_pe->view({seq_len, 1, qk_rope_head_dim_}), pos_ids, true); |
| 189 | + |
| 190 | + auto q_latent = project_q_nope_to_latent_(q_nope); |
| 191 | + auto query_states = infinicore::op::cat({q_latent, q_pe}, 2); |
| 192 | + auto key_states = infinicore::op::cat({kv_norm->view({seq_len, 1, kv_lora_rank_}), k_pe_rope}, 2); |
| 193 | + auto value_states = infinicore::op::pad(kv_norm->view({seq_len, 1, kv_lora_rank_}), |
| 194 | + {0, static_cast<int>(qk_rope_head_dim_)}, "constant", 0.0); |
| 195 | + |
| 196 | + auto attn_output = latent_attn_->forward(query_states, key_states, value_states); |
| 197 | + return project_latent_to_value_(attn_output, batch_size, seq_len); |
| 198 | +} |
| 199 | + |
| 200 | +} // namespace infinilm::models::deepseek_v2 |
0 commit comments