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_rwkv7_base::llm_build_rwkv7_base (const llama_model & model ,const llm_graph_params & params ) : 6llm_graph_context (params ), 7model (model ) {} 8 9ggml_tensor * llm_build_rwkv7_base::build_rwkv7_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_RWKV7 : 16 { 17ggml_tensor * xk = ggml_add (ctx0 ,ggml_mul (ctx0 ,sx ,layer -> channel_mix_lerp_k ),cur ); 18 19ggml_tensor * k = ggml_sqr (ctx0 ,ggml_relu (ctx0 ,build_lora_mm (layer -> channel_mix_key ,xk ))); 20 21cur = build_lora_mm (layer -> channel_mix_value ,k ); 22 } 23break ; 24default : 25GGML_ABORT ("fatal error" ); 26 } 27return cur ; 28} 29 30ggml_tensor * llm_build_rwkv7_base::build_rwkv7_time_mix (llm_graph_input_rs * inp , 31ggml_tensor * cur , 32ggml_tensor * x_prev , 33ggml_tensor *& first_layer_value , 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_embd = hparams .n_embd ; 41const auto head_size = hparams .wkv_head_size ; 42const auto head_count = n_embd /head_size ; 43const auto n_seq_tokens = ubatch .n_seq_tokens ; 44 45const auto kv_head = mctx_cur -> get_head (); 46 47const auto & layer = model .layers [il ]; 48 49bool has_gating = layer .time_mix_g1 && layer .time_mix_g2 ; 50 51ggml_tensor * sx = ggml_sub (ctx0 ,x_prev ,cur ); 52ggml_tensor * dummy = ggml_new_tensor_4d (ctx0 ,GGML_TYPE_F32 ,n_embd ,n_seq_tokens ,n_seqs ,has_gating ?6 :5 ); 53sx = ggml_repeat (ctx0 ,sx ,dummy ); 54 55ggml_tensor * xxx = ggml_add (ctx0 ,ggml_mul (ctx0 ,sx ,layer .time_mix_lerp_fused ),cur ); 56 57ggml_tensor * xr = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],0 ); 58ggml_tensor * xw = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * sizeof (float )); 59ggml_tensor * xk = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 2 * sizeof (float )); 60ggml_tensor * xv = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 3 * sizeof (float )); 61ggml_tensor * xa = ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 4 * sizeof (float )); 62ggml_tensor * xg = 63has_gating ?ggml_view_2d (ctx0 ,xxx ,n_embd ,n_tokens ,xxx -> nb [1 ],n_embd * n_tokens * 5 * sizeof (float )) : 64nullptr ; 65 66ggml_tensor * r = build_lora_mm (layer .time_mix_receptance ,xr ); 67ggml_tensor * w = ggml_add ( 68ctx0 ,ggml_mul_mat (ctx0 ,layer .time_mix_w2 ,ggml_tanh (ctx0 ,ggml_mul_mat (ctx0 ,layer .time_mix_w1 ,xw ))), 69layer .time_mix_w0 ); 70w = ggml_exp (ctx0 ,ggml_scale (ctx0 ,ggml_sigmoid (ctx0 ,w ),-0.606531 )); 71 72ggml_tensor * k = build_lora_mm (layer .time_mix_key ,xk ); 73ggml_tensor * v = build_lora_mm (layer .time_mix_value ,xv ); 74if (first_layer_value == nullptr ) { 75first_layer_value = v ; 76 }else { 77// Add the first layer value as a residual connection. 78v = ggml_add (ctx0 ,v , 79ggml_mul (ctx0 ,ggml_sub (ctx0 ,first_layer_value ,v ), 80ggml_sigmoid (ctx0 ,ggml_add (ctx0 , 81ggml_mul_mat (ctx0 ,layer .time_mix_v2 , 82ggml_mul_mat (ctx0 ,layer .time_mix_v1 ,xv )), 83layer .time_mix_v0 )))); 84 } 85ggml_tensor * g = nullptr ; 86if (layer .time_mix_g1 && layer .time_mix_g2 ) { 87g = ggml_mul_mat (ctx0 ,layer .time_mix_g2 ,ggml_sigmoid (ctx0 ,ggml_mul_mat (ctx0 ,layer .time_mix_g1 ,xg ))); 88 } 89ggml_tensor * a = ggml_sigmoid ( 90ctx0 ,ggml_add (ctx0 ,ggml_mul_mat (ctx0 ,layer .time_mix_a2 ,ggml_mul_mat (ctx0 ,layer .time_mix_a1 ,xa )), 91layer .time_mix_a0 )); 92 93ggml_tensor * kk = ggml_reshape_3d (ctx0 ,ggml_mul (ctx0 ,k ,layer .time_mix_k_k ),head_size ,head_count ,n_tokens ); 94kk = ggml_l2_norm (ctx0 ,kk ,1e-12 ); 95 96ggml_tensor * ka = ggml_mul (ctx0 ,k ,layer .time_mix_k_a ); 97k = ggml_add (ctx0 ,k ,ggml_sub (ctx0 ,ggml_mul (ctx0 ,a ,ka ),ka )); 98 99r = ggml_reshape_3d (ctx0 ,r ,head_size ,head_count ,n_tokens ); 100w = ggml_reshape_3d (ctx0 ,w ,head_size ,head_count ,n_tokens ); 101k = ggml_reshape_3d (ctx0 ,k ,head_size ,head_count ,n_tokens ); 102v = ggml_reshape_3d (ctx0 ,v ,head_size ,head_count ,n_tokens ); 103a = ggml_reshape_3d (ctx0 ,a ,head_size ,head_count ,n_tokens ); 104 105ggml_tensor * wkv_state = build_rs (inp ,mctx_cur -> get_s_l (il ),hparams .n_embd_s (),n_seqs ); 106 107ggml_tensor * wkv_output = ggml_rwkv_wkv7 (ctx0 ,r ,w ,k ,v ,ggml_neg (ctx0 ,kk ),ggml_mul (ctx0 ,kk ,a ),wkv_state ); 108cur = ggml_view_1d (ctx0 ,wkv_output ,n_embd * n_tokens ,0 ); 109wkv_state = ggml_view_1d (ctx0 ,wkv_output ,n_embd * head_size * n_seqs ,n_embd * n_tokens * sizeof (float )); 110 111ggml_build_forward_expand ( 112gf ,ggml_cpy (ctx0 ,wkv_state , 113ggml_view_1d (ctx0 ,mctx_cur -> get_s_l (il ),hparams .n_embd_s ()* n_seqs , 114hparams .n_embd_s ()* kv_head * ggml_element_size (mctx_cur -> get_s_l (il ))))); 115 116if (layer .time_mix_ln && layer .time_mix_ln_b ) { 117// group norm with head_count groups 118cur = ggml_reshape_3d (ctx0 ,cur ,n_embd /head_count ,head_count ,n_tokens ); 119cur = ggml_norm (ctx0 ,cur ,64e-5f ); 120 121// Convert back to regular vectors. 122cur = ggml_reshape_2d (ctx0 ,cur ,n_embd ,n_tokens ); 123cur = ggml_add (ctx0 ,ggml_mul (ctx0 ,cur ,layer .time_mix_ln ),layer .time_mix_ln_b ); 124 }else { 125cur = ggml_reshape_2d (ctx0 ,cur ,n_embd ,n_tokens ); 126 } 127ggml_tensor * rk = ggml_sum_rows ( 128ctx0 ,ggml_mul (ctx0 ,ggml_mul (ctx0 ,k ,r ),ggml_reshape_2d (ctx0 ,layer .time_mix_r_k ,head_size ,head_count ))); 129cur = ggml_add (ctx0 ,cur ,ggml_reshape_2d (ctx0 ,ggml_mul (ctx0 ,v ,rk ),n_embd ,n_tokens )); 130 131if (has_gating ) { 132cur = ggml_mul (ctx0 ,cur ,g ); 133 } 134cur = build_lora_mm (layer .time_mix_output ,cur ); 135 136return ggml_reshape_3d (ctx0 ,cur ,n_embd ,n_seq_tokens ,n_seqs ); 137}