fit: Fix memory allocation for MTP layers (#26605)

This commit is contained in:
Sergey Malinin
2026-08-05 13:29:45 +02:00
committed by GitHub
parent 9303cdd8d3
commit 9a688e51e6
2 changed files with 19 additions and 1 deletions
+15
View File
@@ -2890,6 +2890,21 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid,
int64_t n_embd_, int64_t n_embd_q_, int64_t n_embd_k_, int64_t n_embd_v_,
int flags) {
const int64_t n_embd_qkv = n_embd_q_ + n_embd_k_ + n_embd_v_;
if (flags & TENSOR_SKIP) {
const int skip = TENSOR_NOT_REQUIRED | TENSOR_SKIP;
create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", bid), {n_embd_, n_embd_qkv}, skip | TENSOR_SKIP_IF_VIRTUAL);
create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, skip | TENSOR_SKIP_IF_VIRTUAL);
create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", bid), {n_embd_, n_embd_q_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", bid), {n_embd_, n_embd_k_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", bid), {n_embd_, n_embd_v_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_Q, "bias", bid), {n_embd_q_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_K, "bias", bid), {n_embd_k_}, skip);
create_tensor(tn(LLM_TENSOR_ATTN_V, "bias", bid), {n_embd_v_}, skip);
return;
}
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", bid), {n_embd_, n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);
if (layer.wqkv) {
layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL);