Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3#include "llama-kv-cache.h"4#include "llama-kv-cache-dsa.h"5 6#include <cmath>7 8// iHC (independent Hyper-Connections) helpers. Same layout as the DeepSeek-V4 HC, but without9// the comb/sinkhorn term: hc_fn makes only 2*hc coefficients (pre + post). The streams mix10// through the pre-reduce / post-distribute round trip instead.11 12static size_t hy_v4_elem_offset(const ggml_tensor * t, int64_t i) {13 return ggml_row_size(t->type, i);14}15 16static ggml_tensor * hy_v4_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) {17 return ggml_view_1d(ctx, t, ne0, hy_v4_elem_offset(t, i0));18}19 20static ggml_tensor * hy_v4_view_2d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t i0) {21 return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], hy_v4_elem_offset(t, i0));22}23 24void llama_model_hy_v4::load_arch_hparams(llama_model_loader & ml) {25 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);26 ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);27 ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);28 ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);29 ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);30 ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);31 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);32 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);33 ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);34 ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);35 ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);36 37 // routed-expert SwiGLU logits clamp (shared/dense experts are NOT clamped, so38 // swiglu_clamp_shexp is intentionally left at its 0 default)39 ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false);40 41 ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult);42 ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);43 ml.get_key(LLM_KV_HYPER_CONNECTION_MAGNITUDE, hparams.hc_magnitude);44 45 // DSA is absent on the all-full_attention checkpoints, so indexer_top_k stays 0 there46 ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head, false);47 ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size, false);48 ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false);49 50 if (hparams.indexer_top_k > 0) {51 // the reference plumbs rms_norm_eps into the indexer k_norm LayerNorm, and build_norm52 // reads f_norm_eps for LLM_NORM53 hparams.f_norm_eps = hparams.f_norm_rms_eps;54 55 if (hparams.indexer_n_head == 0 || hparams.indexer_head_size <= hparams.n_rot()) {56 throw std::runtime_error("hy_v4: bad indexer head count / key length");57 }58 59 ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false);60 if (!hparams.is_indexer_full(0)) {61 throw std::runtime_error("hy_v4: layer 0 must own an indexer, nothing precedes it to share");62 }63 }64 65 GGML_ASSERT(hparams.is_mla());66 67 type = LLM_TYPE_UNKNOWN;68}69 70void llama_model_hy_v4::load_arch_tensors(llama_model_loader &) {71 LLAMA_LOAD_LOCALS;72 73 const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();74 const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();75 const int64_t n_embd_head_qk_rope = hparams.n_rot();76 const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope;77 GGML_ASSERT(n_embd_head_qk_nope >= 1);78 79 const int64_t q_lora_rank = hparams.n_lora_q;80 const int64_t kv_lora_rank = hparams.n_lora_kv;81 const int64_t n_ff_exp = hparams.n_ff_exp();82 const int64_t n_expert_shared = hparams.n_expert_shared;83 const int64_t hc = hparams.dsv4_hc_mult;84 85 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);86 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);87 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);88 89 // global iHC head (collapses hc streams before the final norm)90 hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc * n_embd, hc}, 0);91 hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc}, 0);92 hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);93 94 for (int i = 0; i < n_layer; ++i) {95 auto & layer = layers[i];96 97 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);98 layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);99 100 layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);101 layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);102 layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k_mla}, 0);103 layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, 0);104 layer.attn_kv_a_norm= create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM,"weight", i), {kv_lora_rank}, 0);105 layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, 0);106 layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, 0);107 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, 0);108 layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head * n_embd_head_v_mla}, 0);109 110 // only "full" indexer layers ship weights; "shared" layers reuse their top-k111 if (hparams.indexer_top_k > 0 && hparams.is_indexer_full(i)) {112 const int64_t n_indexer_head = hparams.indexer_n_head;113 const int64_t n_embd_indexer = hparams.indexer_head_size;114 115 layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_indexer_head * n_embd_indexer}, 0);116 layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, n_embd_indexer}, 0);117 layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_embd_indexer}, 0);118 layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {n_embd_indexer}, 0);119 layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_indexer_head}, 0);120 }121 122 layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc * n_embd, 2 * hc}, 0);123 layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {2 * hc}, 0);124 layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {2}, 0);125 layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc * n_embd, 2 * hc}, 0);126 layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {2 * hc}, 0);127 layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {2}, 0);128 129 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);130 131 if (i < (int) hparams.n_layer_dense_lead) {132 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);133 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);134 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);135 } else {136 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);137 layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, TENSOR_NOT_REQUIRED);138 139 if (n_expert == 0) {140 throw std::runtime_error("n_expert must be > 0");141 }142 if (n_expert_used == 0) {143 throw std::runtime_error("n_expert_used must be > 0");144 }145 146 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);147 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);148 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);149 150 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);151 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd}, 0);152 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);153 }154 }155}156 157std::unique_ptr<llm_graph_context> llama_model_hy_v4::build_arch_graph(const llm_graph_params & params) const {158 return std::make_unique<graph>(*this, params);159}160 161// reduce hc streams x[:,i,:] weighted by w[i,:] -> [n_embd, n_tokens]162// reference runs this in fp32 (inside the float() / autocast(fp32) context)163static ggml_tensor * hy_v4_hc_reduce(ggml_context * ctx0, ggml_tensor * x, ggml_tensor * w, int64_t hc, int64_t n_embd, int64_t nt, ggml_type out_type) {164 ggml_tensor * x_f32 = ggml_cast(ctx0, x, GGML_TYPE_F32);165 ggml_tensor * result = nullptr;166 for (int64_t ih = 0; ih < hc; ++ih) {167 ggml_tensor * xh = ggml_view_2d(ctx0, x_f32, n_embd, nt, x_f32->nb[2], ih * x_f32->nb[1]);168 ggml_tensor * wh = ggml_view_2d(ctx0, w, 1, nt, w->nb[1], ih * w->nb[0]);169 ggml_tensor * cur = ggml_mul(ctx0, xh, wh);170 result = result ? ggml_add(ctx0, result, cur) : cur;171 }172 return ggml_cast(ctx0, result, out_type);173}174 175ggml_tensor * llama_model_hy_v4::graph::build_hc_pre(176 ggml_tensor * x,177 ggml_tensor * hc_fn,178 ggml_tensor * hc_scale,179 ggml_tensor * hc_base,180 ggml_tensor ** post,181 int il) const {182 const int64_t hc = hparams.dsv4_hc_mult;183 const int64_t nt = x->ne[2];184 GGML_ASSERT(x->ne[0] == n_embd && x->ne[1] == hc);185 186 ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc * n_embd, nt);187 ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps);188 ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); // [2*hc, nt]189 cb(mixes, "hc_mixes", il);190 191 ggml_tensor * scale_pre = hy_v4_view_1d(ctx0, hc_scale, 1, 0);192 ggml_tensor * scale_post = hy_v4_view_1d(ctx0, hc_scale, 1, 1);193 ggml_tensor * base_pre = hy_v4_view_1d(ctx0, hc_base, hc, 0);194 ggml_tensor * base_post = hy_v4_view_1d(ctx0, hc_base, hc, hc);195 196 // pre = sigmoid(mixes[:hc]*scale_pre + base_pre) + eps197 ggml_tensor * pre = hy_v4_view_2d(ctx0, mixes, hc, nt, 0);198 pre = ggml_mul(ctx0, pre, scale_pre);199 pre = ggml_add(ctx0, pre, base_pre);200 pre = ggml_sigmoid(ctx0, pre);201 pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);202 cb(pre, "hc_pre", il);203 204 // post = magnitude*sigmoid(mixes[hc:2hc]*scale_post + base_post) + eps205 ggml_tensor * po = hy_v4_view_2d(ctx0, mixes, hc, nt, hc);206 po = ggml_mul(ctx0, po, scale_post);207 po = ggml_add(ctx0, po, base_post);208 po = ggml_sigmoid(ctx0, po);209 po = ggml_scale(ctx0, po, hparams.hc_magnitude);210 po = ggml_scale_bias(ctx0, po, 1.0f, hparams.dsv4_hc_eps);211 *post = po;212 cb(po, "hc_post_gate", il);213 214 return hy_v4_hc_reduce(ctx0, x, pre, hc, n_embd, nt, x->type);215}216 217ggml_tensor * llama_model_hy_v4::graph::build_hc_post(218 ggml_tensor * x,219 ggml_tensor * residual,220 ggml_tensor * post,221 int il) const {222 GGML_UNUSED(il);223 const int64_t hc = hparams.dsv4_hc_mult;224 const int64_t nt = x->ne[1];225 GGML_ASSERT(x->ne[0] == n_embd);226 GGML_ASSERT(residual->ne[1] == hc);227 228 // reference HC post runs entirely in fp32 to avoid bf16 rounding accumulation229 // across 78 layers: post.float() * x.float() + residual.float() -> .to(dtype)230 ggml_tensor * x_f32 = ggml_cast(ctx0, x, GGML_TYPE_F32);231 ggml_tensor * post_f32 = ggml_cast(ctx0, post, GGML_TYPE_F32);232 ggml_tensor * res_f32 = ggml_cast(ctx0, residual, GGML_TYPE_F32);233 234 ggml_tensor * out = nullptr;235 for (int64_t i = 0; i < hc; ++i) {236 ggml_tensor * res_i = ggml_view_2d(ctx0, res_f32, n_embd, nt, res_f32->nb[2], i * res_f32->nb[1]);237 ggml_tensor * post_i = ggml_view_2d(ctx0, post_f32, 1, nt, post_f32->nb[1], i * post_f32->nb[0]);238 ggml_tensor * cur = ggml_add(ctx0, res_i, ggml_mul(ctx0, x_f32, post_i));239 cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt);240 out = out ? ggml_concat(ctx0, out, cur, 1) : cur;241 }242 243 // cast back to the original type (bf16)244 out = ggml_cast(ctx0, out, residual->type);245 return out; // [n_embd, hc, nt]246}247 248ggml_tensor * llama_model_hy_v4::graph::build_hc_head(249 ggml_tensor * x,250 ggml_tensor * hc_fn,251 ggml_tensor * hc_scale,252 ggml_tensor * hc_base) const {253 const int64_t hc = hparams.dsv4_hc_mult;254 const int64_t nt = x->ne[2];255 256 ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc * n_embd, nt);257 ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps);258 ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); // [hc, nt]259 cb(mixes, "hc_head_mixes", -1);260 261 ggml_tensor * pre = ggml_mul(ctx0, mixes, hc_scale);262 pre = ggml_add(ctx0, pre, hc_base);263 pre = ggml_sigmoid(ctx0, pre);264 pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);265 cb(pre, "hc_head_pre", -1);266 267 return hy_v4_hc_reduce(ctx0, x, pre, hc, n_embd, nt, x->type);268}269 270ggml_tensor * llama_model_hy_v4::graph::build_attention(271 const llama_model & model,272 llm_graph_input_attn_k * inp_attn,273 ggml_tensor * cur,274 ggml_tensor * inp_pos,275 float kq_scale,276 int il) const {277 const auto & layer = model.layers[il];278 279 const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();280 const int64_t n_embd_head_qk_rope = hparams.n_rot();281 const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;282 const uint32_t kv_lora_rank = hparams.n_lora_kv;283 284 ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur);285 q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);286 q = ggml_mul_mat(ctx0, layer.wq_b, q);287 288 ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens,289 ggml_row_size(q->type, n_embd_head_k), ggml_row_size(q->type, n_embd_head_k) * n_head, 0);290 ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens,291 ggml_row_size(q->type, n_embd_head_k), ggml_row_size(q->type, n_embd_head_k) * n_head,292 ggml_row_size(q->type, n_embd_head_qk_nope));293 294 ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);295 ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,296 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);297 ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,298 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),299 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),300 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));301 302 q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,303 ext_factor, attn_factor, beta_fast, beta_slow);304 cb(q_pe, "q_pe", il);305 k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,306 ext_factor, attn_factor, beta_fast, beta_slow);307 cb(k_pe, "k_pe", il);308 309 kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);310 cb(kv_cmpr, "kv_cmpr", il);311 312 // MLA absorption: q_nope @ wk_b -> compressed space313 q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);314 ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);315 q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);316 317 // note: rope must go first for in-place context shifting in build_rope_shift()318 ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);319 320 kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);321 ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);322 ggml_tensor * Vcur = kv_cmpr;323 324 // MLA-as-MQA; wo applied manually below so the gated-MLA gate can sit before o_proj325 ggml_tensor * attn = build_attn(inp_attn,326 nullptr, nullptr, nullptr,327 Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, layer.wv_b, kq_scale, il);328 cb(attn, "attn_kqv", il); // [n_head * n_embd_head_v, n_tokens]329 330 // gated MLA: elementwise sigmoid gate on the decompressed attention output331 ggml_tensor * gate = ggml_mul_mat(ctx0, layer.wqkv_gate, cur);332 gate = ggml_sigmoid(ctx0, gate);333 attn = ggml_mul(ctx0, attn, gate);334 cb(attn, "attn_gated", il);335 336 ggml_tensor * out = build_lora_mm(layer.wo, attn);337 cb(out, "attn_out", il);338 339 return out;340}341 342ggml_tensor * llama_model_hy_v4::graph::build_indexer_top_k(343 const llama_model & model,344 llm_graph_input_attn_k_dsa * inp_attn_dsa,345 ggml_tensor * cur,346 ggml_tensor * qr,347 ggml_tensor * inp_pos,348 int il) const {349 const auto & layer = model.layers[il];350 351 const int64_t n_indexer_head = hparams.indexer_n_head;352 const int64_t n_embd_indexer = hparams.indexer_head_size;353 const int64_t n_embd_indexer_rope = hparams.n_rot();354 const int64_t n_embd_indexer_nope = n_embd_indexer - n_embd_indexer_rope;355 356 // nope rows come first, so rope only the last n_embd_indexer_rope rows, same as the MLA path357 ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr);358 359 iq = ggml_reshape_3d(ctx0, iq, n_embd_indexer, n_indexer_head, n_tokens);360 361 iq = ggml_rope_ext(ctx0, iq, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base,362 freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);363 iq = ggml_rope_set_offset(iq, n_embd_indexer_nope);364 cb(iq, "indexer_q", il);365 366 ggml_tensor * ik = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur);367 368 ik = build_norm(ik, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il);369 370 ik = ggml_reshape_3d(ctx0, ik, n_embd_indexer, 1, n_tokens);371 372 ik = ggml_rope_ext(ctx0, ik, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base,373 freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);374 ik = ggml_rope_set_offset(ik, n_embd_indexer_nope);375 cb(ik, "indexer_k", il);376 377 // the reference applies a Hadamard rotation here, but it only helps its FP8 kernels.378 // it is orthogonal, so it does not change q.k and we can skip it.379 380 const auto * mctx_lid = inp_attn_dsa->mctx->get_lid();381 const auto & k_idxs_lid = inp_attn_dsa->get_k_idxs_lid();382 ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, ik, k_idxs_lid, il));383 384 ggml_tensor * iw = ggml_mul_mat(ctx0, layer.indexer_proj, cur);385 386 ik = mctx_lid->get_k(ctx0, il);387 388 const auto n_stream = ik->ne[3];389 iq = ggml_view_4d(ctx0, iq, iq->ne[0], iq->ne[1], iq->ne[2]/n_stream, n_stream,390 iq->nb[1], iq->nb[2], iq->nb[3]/n_stream, 0);391 iw = ggml_view_4d(ctx0, iw, iw->ne[0], iw->ne[1]/n_stream, iw->ne[2], n_stream,392 iw->nb[1], iw->nb[2]/n_stream, iw->nb[3]/n_stream, 0);393 394 // fold both reference scale factors into the weights before the big score tensor395 iw = ggml_scale(ctx0, iw, 1.0f / sqrtf(float(n_embd_indexer * n_indexer_head)));396 397 ggml_tensor * score = nullptr;398 if (cparams.fused_lid) {399 score = ggml_lightning_indexer(ctx0, iq, ik, iw, inp_attn_dsa->get_kq_mask_lid());400 cb(score, "indexer_score", il);401 res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il});402 } else {403 iq = ggml_permute(ctx0, iq, 0, 2, 1, 3);404 ik = ggml_permute(ctx0, ik, 0, 2, 1, 3);405 406 score = ggml_mul_mat(ctx0, ik, iq);407 score = ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3));408 score = ggml_relu(ctx0, score);409 score = ggml_mul(ctx0, score, iw);410 score = ggml_sum_rows(ctx0, score);411 score = ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3));412 score = ggml_add(ctx0, score, inp_attn_dsa->get_kq_mask_lid());413 cb(score, "indexer_score", il);414 }415 416 const uint32_t n_top_k = score->ne[0] < (int64_t) hparams.indexer_top_k ? score->ne[0] : hparams.indexer_top_k;417 418 return ggml_cont(ctx0, ggml_top_k(ctx0, score, n_top_k));419}420 421ggml_tensor * llama_model_hy_v4::graph::build_attention_dsa(422 const llama_model & model,423 llm_graph_input_attn_k_dsa * inp_attn_dsa,424 ggml_tensor * cur,425 ggml_tensor * inp_pos,426 ggml_tensor ** last_top_k,427 float kq_scale,428 int il) const {429 const auto & layer = model.layers[il];430 431 const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();432 const int64_t n_embd_head_qk_rope = hparams.n_rot();433 const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;434 const uint32_t kv_lora_rank = hparams.n_lora_kv;435 436 ggml_tensor * qr = ggml_mul_mat(ctx0, layer.wq_a, cur);437 qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);438 439 if (hparams.is_indexer_full(il)) {440 *last_top_k = build_indexer_top_k(model, inp_attn_dsa, cur, qr, inp_pos, il);441 cb(*last_top_k, "top_k", il);442 }443 GGML_ASSERT(*last_top_k != nullptr);444 445 ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr);446 447 ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens,448 ggml_row_size(q->type, n_embd_head_k), ggml_row_size(q->type, n_embd_head_k) * n_head, 0);449 ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens,450 ggml_row_size(q->type, n_embd_head_k), ggml_row_size(q->type, n_embd_head_k) * n_head,451 ggml_row_size(q->type, n_embd_head_qk_nope));452 453 ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);454 ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,455 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);456 ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,457 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),458 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),459 ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));460 461 q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,462 ext_factor, attn_factor, beta_fast, beta_slow);463 cb(q_pe, "q_pe", il);464 k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,465 ext_factor, attn_factor, beta_fast, beta_slow);466 cb(k_pe, "k_pe", il);467 468 kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);469 cb(kv_cmpr, "kv_cmpr", il);470 471 q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);472 ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);473 q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);474 475 ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);476 477 kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);478 ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);479 ggml_tensor * Vcur = kv_cmpr;480 481 ggml_tensor * attn = build_attn(inp_attn_dsa,482 nullptr, nullptr, nullptr,483 Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, layer.wv_b, *last_top_k, kq_scale, il);484 cb(attn, "attn_kqv", il);485 486 ggml_tensor * gate = ggml_mul_mat(ctx0, layer.wqkv_gate, cur);487 gate = ggml_sigmoid(ctx0, gate);488 attn = ggml_mul(ctx0, attn, gate);489 cb(attn, "attn_gated", il);490 491 ggml_tensor * out = build_lora_mm(layer.wo, attn);492 cb(out, "attn_out", il);493 494 return out;495}496 497llama_model_hy_v4::graph::graph(const llama_model & model, const llm_graph_params & params) :498 llm_graph_context(params) {499 const int64_t hc = hparams.dsv4_hc_mult;500 const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();501 const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));502 503 ggml_tensor * cur;504 505 const bool is_dsa = hparams.indexer_top_k > 0;506 507 ggml_tensor * inp = build_inp_embd(model.tok_embd);508 ggml_tensor * inp_pos = build_inp_pos();509 llm_graph_input_attn_k * inp_attn = is_dsa ? nullptr : build_attn_inp_k();510 llm_graph_input_attn_k_dsa * inp_attn_dsa = is_dsa ? build_attn_inp_k_dsa() : nullptr;511 ggml_tensor * inp_out_ids = build_inp_out_ids();512 513 // top-k of the last "full" indexer layer, reused by the following "shared" layers514 ggml_tensor * last_top_k = nullptr;515 516 // expand the single embedding into hc parallel residual streams517 ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens);518 inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);519 cb(inpL, "hc_init", -1);520 521 for (int il = 0; il < n_layer; ++il) {522 ggml_tensor * residual = inpL;523 ggml_tensor * post = nullptr;524 525 cur = build_hc_pre(inpL, model.layers[il].hc_attn_fn, model.layers[il].hc_attn_scale,526 model.layers[il].hc_attn_base, &post, il);527 cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il);528 cb(cur, "attn_norm", il);529 530 cur = is_dsa531 ? build_attention_dsa(model, inp_attn_dsa, cur, inp_pos, &last_top_k, kq_scale, il)532 : build_attention(model, inp_attn, cur, inp_pos, kq_scale, il);533 534 inpL = build_hc_post(cur, residual, post, il);535 cb(inpL, "hc_attn_out", il);536 537 residual = inpL;538 cur = build_hc_pre(inpL, model.layers[il].hc_ffn_fn, model.layers[il].hc_ffn_scale,539 model.layers[il].hc_ffn_base, &post, il);540 cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il);541 cb(cur, "ffn_norm", il);542 543 const auto & layer = model.layers[il];544 if ((uint32_t) il < hparams.n_layer_dense_lead) {545 cur = build_ffn(cur,546 layer.ffn_up, NULL, NULL,547 layer.ffn_gate, NULL, NULL,548 layer.ffn_down, NULL, NULL,549 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);550 cb(cur, "ffn_out", il);551 } else {552 ggml_tensor * moe_out = build_moe_ffn(cur,553 layer.ffn_gate_inp,554 layer.ffn_up_exps,555 layer.ffn_gate_exps,556 layer.ffn_down_exps,557 layer.ffn_exp_probs_b,558 n_expert, n_expert_used,559 LLM_FFN_SILU, hparams.expert_weights_norm,560 hparams.expert_weights_scale,561 (llama_expert_gating_func_type) hparams.expert_gating_func,562 il,563 nullptr,564 nullptr);565 cb(moe_out, "ffn_moe_out", il);566 567 ggml_tensor * ffn_shexp = build_ffn(cur,568 layer.ffn_up_shexp, NULL, NULL,569 layer.ffn_gate_shexp, NULL, NULL,570 layer.ffn_down_shexp, NULL, NULL,571 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);572 cb(ffn_shexp, "ffn_shexp", il);573 574 cur = ggml_add(ctx0, moe_out, ffn_shexp);575 cb(cur, "ffn_out", il);576 }577 578 inpL = build_hc_post(cur, residual, post, il);579 cb(inpL, "l_out", il);580 }581 582 // prune to the requested output rows once, after all HC streams are done583 if (inp_out_ids) {584 ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd * hc, n_tokens);585 flat = ggml_get_rows(ctx0, flat, inp_out_ids);586 inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs);587 }588 589 cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);590 cb(cur, "hc_head", -1);591 592 cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);593 cb(cur, "result_norm", -1);594 res->t_embd = cur;595 596 cur = ggml_mul_mat(ctx0, model.output, cur);597 cb(cur, "result_output", -1);598 res->t_logits = cur;599 600 ggml_build_forward_expand(gf, cur);601}602