char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
cc45f2ada
char/aliceai-t5-moe
1#include "models.h" 2 3#include "llama-memory-recurrent.h" 4 5llm_build_rwkv6_base::llm_build_rwkv6_base (const llama_model & model ,const llm_graph_params & params ) : 6llm_graph_context (params ), 7model (model ) {} 8 9ggml_tensor * llm_build_rwkv6_base::build_rwkv6_channel_mix (const llama_layer * layer , 10ggml_tensor * cur , 11ggml_tensor * x_prev , 12llm_arch arch )const { 13ggml_tensor * sx = ggml_sub (ctx0 ,x_prev ,cur ); 14switch (arch ) { 15case LLM_ARCH_RWKV6 : 16 { 17ggml_tensor * xk = ggml_add (ctx0 ,ggml_mul (ctx0 ,sx ,layer -> channel_mix_lerp_k ),cur ); 18ggml_tensor * xr = ggml_add (ctx0 ,ggml_mul (ctx0 ,sx ,layer -> channel_mix_lerp_r ),cur ); 19 20ggml_tensor * r = ggml_sigmoid (ctx0 ,build_lora_mm (layer -> channel_mix_receptance ,xr )); 21ggml_tensor * k = ggml_sqr (ctx0 ,ggml_relu (ctx0 ,build_lora_mm (layer -> channel_mix_key ,xk ))); 22cur = ggml_mul (ctx0 ,r ,build_lora_mm (layer -> channel_mix_value ,k )); 23 } 24break ; 25default : 26GGML_ABORT ("fatal error" ); 27 } 28return cur ; 29} 30 31ggml_tensor * llm_build_rwkv6_base::build_rwkv6_time_mix (llm_graph_input_rs * inp , 32ggml_tensor * cur , 33ggml_tensor * x_prev , 34const llama_ubatch & ubatch , 35int il )const { 36const auto * mctx_cur = static_cast < const llama_memory_recurrent_context *> (mctx ); 37 38const auto n_tokens = ubatch .n_tokens ; 39const auto n_seqs = ubatch .n_seqs ; 40const auto n_seq_tokens = ubatch .n_seq_tokens ; 41const auto n_embd = hparams .n_embd ; 42const auto head_size = hparams .wkv_head_size ; 43const auto n_head = n_embd /head_size ; 44const auto n_head_kv = hparams .n_head_kv (il ); 45 46const auto kv_head = mctx_cur -> get_head (); 47 48const auto & layer = model .layers [il ]; 49 50bool is_qrwkv = layer .time_mix_first == nullptr ; 51 52ggml_tensor * sx = ggml_sub (ctx0 ,x_prev ,cur ); 53 54sx = ggml_reshape_2d (ctx0 ,sx ,n_embd ,n_tokens ); 55cur = ggml_reshape_2d (ctx0 ,cur ,n_embd ,n_tokens ); 56 57ggml_tensor * xxx = ggml_add (ctx0 ,ggml_mul (ctx0 ,sx ,layer .time_mix_lerp_x ),cur ); 58 59xxx = ggml_reshape_4d (ctx0 ,ggml_tanh (ctx0 ,ggml_mul_mat (ctx0 ,layer .time_mix_w1 ,xxx )), 60layer .time_mix_w1 -> ne [1 ] /5 ,1 ,5 ,n_tokens ); 61 62xxx = ggml_cont (ctx0 ,ggml_permute (ctx0 ,xxx ,0 ,1 ,3 ,2 )); 63 64xxx = ggml_mul_mat ( 65ctx0 ,ggml_reshape_4d (ctx0 ,layer .time_mix_w2 ,layer .time_mix_w2 -> ne [0 ],layer .time_mix_w2 -> ne [1 ],1 ,5 ),xxx ); 66 67ggml_tensor * xw ,* xk ,* xv ,* xr ,* xg ; 68if (layer .time_mix_lerp_fused ) { 69// fusing these weights makes some performance improvement 70sx = ggml_reshape_3d (ctx0 ,sx ,n_embd ,1 ,n_tokens ); 71cur = ggml_reshape_3d (ctx0 ,cur ,n_embd ,1 ,n_tokens ); 72xxx = ggml_add (ctx0 ,ggml_mul (ctx0 ,ggml_add (ctx0 ,xxx ,layer .time_mix_lerp_fused ),sx ),cur ); 73xw = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],0 ); 74xk = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * sizeof (float )); 75xv = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 2 * sizeof (float )); 76xr = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 3 * sizeof (float )); 77xg = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 4 * sizeof (float )); 78 }else { 79// for backward compatibility 80xw = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],0 ); 81xk = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * sizeof (float )); 82xv = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 2 * sizeof (float )); 83xr = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 3 * sizeof (float )); 84xg = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 4 * sizeof (float )); 85 86xw = ggml_add (ctx0 ,ggml_mul (ctx0 ,ggml_add (ctx0 ,xw ,layer .time_mix_lerp_w ),sx ),cur ); 87xk = ggml_add (ctx0 ,ggml_mul (ctx0 ,ggml_add (ctx0 ,xk ,layer .time_mix_lerp_k ),sx ),cur ); 88xv = ggml_add (ctx0 ,ggml_mul (ctx0 ,ggml_add (ctx0 ,xv ,layer .time_mix_lerp_v ),sx ),cur ); 89xr = ggml_add (ctx0 ,ggml_mul (ctx0 ,ggml_add (ctx0 ,xr ,layer .time_mix_lerp_r ),sx ),cur ); 90xg = ggml_add (ctx0 ,ggml_mul (ctx0 ,ggml_add (ctx0 ,xg ,layer .time_mix_lerp_g ),sx ),cur ); 91 } 92ggml_tensor * r = build_lora_mm (layer .time_mix_receptance ,xr ); 93ggml_tensor * k = build_lora_mm (layer .time_mix_key ,xk ); 94ggml_tensor * v = build_lora_mm (layer .time_mix_value ,xv ); 95if (layer .time_mix_receptance_b ) { 96r = ggml_add (ctx0 ,r ,layer .time_mix_receptance_b ); 97 } 98if (layer .time_mix_key_b ) { 99k = ggml_add (ctx0 ,k ,layer .time_mix_key_b ); 100 } 101if (layer .time_mix_value_b ) { 102v = ggml_add (ctx0 ,v ,layer .time_mix_value_b ); 103 } 104ggml_tensor * g = build_lora_mm (layer .time_mix_gate ,xg ); 105if (is_qrwkv ) { 106g = ggml_sigmoid (ctx0 ,g ); 107 }else { 108g = ggml_silu (ctx0 ,g ); 109 } 110if (n_head_kv != 0 && n_head_kv != n_head ) { 111GGML_ASSERT (n_head %n_head_kv == 0 ); 112k = ggml_reshape_4d (ctx0 ,k ,head_size ,1 ,n_head_kv ,n_tokens ); 113v = ggml_reshape_4d (ctx0 ,v ,head_size ,1 ,n_head_kv ,n_tokens ); 114ggml_tensor * tmp = ggml_new_tensor_4d (ctx0 ,GGML_TYPE_F32 ,head_size ,n_head /n_head_kv ,n_head_kv ,n_tokens ); 115k = ggml_repeat (ctx0 ,k ,tmp ); 116v = ggml_repeat (ctx0 ,v ,tmp ); 117 } 118k = ggml_reshape_3d (ctx0 ,k ,head_size ,n_head ,n_tokens ); 119v = ggml_reshape_3d (ctx0 ,v ,head_size ,n_head ,n_tokens ); 120r = ggml_reshape_3d (ctx0 ,r ,head_size ,n_head ,n_tokens ); 121 122ggml_tensor * w = 123ggml_mul_mat (ctx0 ,layer .time_mix_decay_w2 ,ggml_tanh (ctx0 ,ggml_mul_mat (ctx0 ,layer .time_mix_decay_w1 ,xw ))); 124 125w = ggml_add (ctx0 ,w ,layer .time_mix_decay ); 126w = ggml_exp (ctx0 ,ggml_neg (ctx0 ,ggml_exp (ctx0 ,w ))); 127w = ggml_reshape_3d (ctx0 ,w ,head_size ,n_head ,n_tokens ); 128 129if (is_qrwkv ) { 130// k = k * (1 - w) 131k = ggml_sub (ctx0 ,k ,ggml_mul (ctx0 ,k ,w )); 132 } 133ggml_tensor * wkv_state = build_rs (inp ,mctx_cur -> get_s_l (il ),hparams .n_embd_s (),n_seqs ); 134 135ggml_tensor * wkv_output ; 136if (is_qrwkv ) { 137wkv_output = ggml_gated_linear_attn (ctx0 ,k ,v ,r ,w ,wkv_state ,pow (head_size ,-0.5f )); 138 }else { 139wkv_output = ggml_rwkv_wkv6 (ctx0 ,k ,v ,r ,layer .time_mix_first ,w ,wkv_state ); 140 } 141cur = ggml_view_1d (ctx0 ,wkv_output ,n_embd * n_tokens ,0 ); 142wkv_state = ggml_view_1d (ctx0 ,wkv_output ,n_embd * head_size * n_seqs ,n_embd * n_tokens * sizeof (float )); 143 144ggml_build_forward_expand ( 145gf ,ggml_cpy (ctx0 ,wkv_state , 146ggml_view_1d (ctx0 ,mctx_cur -> get_s_l (il ),hparams .n_embd_s ()* n_seqs , 147hparams .n_embd_s ()* kv_head * ggml_element_size (mctx_cur -> get_s_l (il ))))); 148 149if (!is_qrwkv ) { 150// group norm with head_count groups 151cur = ggml_reshape_3d (ctx0 ,cur ,n_embd /n_head ,n_head ,n_tokens ); 152cur = ggml_norm (ctx0 ,cur ,64e-5f ); 153 154// Convert back to regular vectors. 155cur = ggml_reshape_2d (ctx0 ,cur ,n_embd ,n_tokens ); 156cur = ggml_add (ctx0 ,ggml_mul (ctx0 ,cur ,layer .time_mix_ln ),layer .time_mix_ln_b ); 157 }else { 158cur = ggml_reshape_2d (ctx0 ,cur ,n_embd ,n_tokens ); 159 } 160cur = ggml_mul (ctx0 ,cur ,g ); 161cur = build_lora_mm (layer .time_mix_output ,cur ); 162 163return ggml_reshape_3d (ctx0 ,cur ,n_embd ,n_seq_tokens ,n_seqs ); 164}