Felipe97/llama-cpp-compiled
01.1k
1#include "models.h"2 3#include "llama-impl.h"4#include "llama-kv-cache.h"5#include "llama-kv-cache-iswa.h"6 7void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {8 9 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);10 ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, false);11 hparams.f_final_logit_softcapping = 0.0f;12 ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);13 14 // drafts for M-RoPE targets carry degenerate sections [n_rot/2, 0, 0, 0]15 ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);16 17 ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false);18 ml.get_key(LLM_KV_DFLASH_CONV_KERNEL_SIZE, hparams.dflash_conv_kernel_size, false);19 ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE, hparams.dflash_conv_group_size, false);20 ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK, hparams.dflash_selector_rank, false);21 ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K, hparams.dflash_selector_top_k, false);22 23 if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) {24 throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata");25 }26 27 hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd;28 29 std::string layers;30 const char * sep = "";31 for (const auto id : target_layer_ids) {32 layers += sep;33 layers += std::to_string(id);34 sep = ", ";35 }36 LLAMA_LOG_INFO("%s: DFlash extract_layers = [%s]\n", __func__, layers.c_str());37 38 // DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)39 ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);40 if (hparams.dsv4_hc_mult > 0) {41 ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);42 ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);43 ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);44 ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);45 ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale);46 ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm);47 ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);48 ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all);49 if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, 0)) {50 hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;51 }52 ml.get_key(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count);53 ml.get_key(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank);54 ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);55 ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);56 ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false);57 58 GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero59 60 if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) {61 throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring");62 }63 for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {64 if (hparams.dsv4_compress_ratios[il] != 0) {65 throw std::runtime_error("DSpark DSV4 draft expects uncompressed attention on all stages");66 }67 }68 69 GGML_ASSERT(hparams.n_swa > 0);70 hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;71 hparams.set_swa_pattern(0);72 for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {73 hparams.is_swa_impl[il] = true;74 }75 hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;76 hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;77 78 type = LLM_TYPE_UNKNOWN;79 return;80 }81 82 // optional interleaved sliding-window attention with per-layer pattern array.83 // DFlash has a single rope, so the SWA rope == main rope.84 if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {85 hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;86 ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl);87 hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;88 hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;89 }90 91 type = LLM_TYPE_UNKNOWN;92}93 94void llama_model_dflash::load_arch_tensors(llama_model_loader &) {95 LLAMA_LOAD_LOCALS;96 97 const int64_t n_embd_inp = hparams.n_embd_inp_enc();98 99 tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);100 101 // reduced draft vocab (optional): d2t maps draft rows to target token ids102 int64_t n_vocab_draft = n_vocab;103 const struct ggml_tensor * d2t_meta = ml->get_tensor_meta("d2t");104 if (d2t_meta) {105 n_vocab_draft = d2t_meta->ne[0];106 d2t = create_tensor(tn(LLM_TENSOR_D2T), { n_vocab_draft }, 0);107 LLAMA_LOG_INFO("%s: DFlash using d2t mapping (draft_vocab_size = %lld)\n", __func__, (long long) n_vocab_draft);108 }109 110 // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head111 //112 // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4)113 // need their own conversion path and graph tweaks114 const struct ggml_tensor * markov_meta = ml->get_tensor_meta("markov_w1.weight");115 if (markov_meta) {116 const int64_t dspark_markov_rank = markov_meta->ne[0];117 118 dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);119 dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0);120 dspark_markov_w2_s = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "scale"), { 1 }, TENSOR_NOT_REQUIRED);121 122 dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED);123 dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED);124 125 LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);126 }127 128 const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight");129 if (selector_meta) {130 const int64_t rank = hparams.dflash_selector_rank;131 if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 ||132 hparams.dflash_conv_kernel_size <= 0 || hparams.dflash_conv_group_size <= 0) {133 throw std::runtime_error("DFlash2 model is missing conv/selector metadata");134 }135 if (n_embd % hparams.dflash_conv_group_size != 0) {136 throw std::runtime_error("DFlash2 hidden size must be divisible by conv_group_size");137 }138 if (n_embd < hparams.dflash_selector_top_k * (hparams.dflash_selector_top_k + 1)) {139 throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice");140 }141 142 dflash_selector_prev = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV, "weight"), { rank, n_vocab }, 0);143 dflash_selector_next = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT, "weight"), { rank, n_vocab }, 0);144 dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0);145 146 LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__,147 hparams.dflash_conv_kernel_size, hparams.dflash_conv_group_size,148 hparams.dflash_selector_rank, hparams.dflash_selector_top_k);149 }150 151 fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);152 fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED);153 output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)154 output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm155 156 // optional: reduced-vocab drafts ship their own lm head, full-vocab drafts can share the target's via ctx_other157 // a draft with its own embeddings + head references no target tensors and can run on devices the target does not use (e.g. -devd with a tensor-split target)158 output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED);159 160 if (hparams.dsv4_hc_mult > 0) {161 const int64_t q_lora_rank = hparams.n_lora_q;162 const int64_t n_ff_exp = hparams.n_ff_exp();163 const int64_t n_expert_shared = hparams.n_expert_shared;164 const int64_t n_embd_head = hparams.n_embd_head_k();165 const int64_t o_groups = hparams.dsv4_o_group_count;166 const int64_t o_lora_rank = hparams.dsv4_o_lora_rank;167 const int64_t hc_mult = hparams.dsv4_hc_mult;168 const int64_t hc_dim = hc_mult * n_embd;169 const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;170 171 hc_head_fn = create_tensor(tn(LLM_TENSOR_HC_HEAD_FN, "weight"), {hc_dim, hc_mult}, 0);172 hc_head_base = create_tensor(tn(LLM_TENSOR_HC_HEAD_BASE, "weight"), {hc_mult}, 0);173 hc_head_scale = create_tensor(tn(LLM_TENSOR_HC_HEAD_SCALE, "weight"), {1}, 0);174 175 for (int i = 0; i < n_layer; ++i) {176 auto & layer = layers[i];177 178 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);179 layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0);180 layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0);181 layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, 0);182 layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head}, 0);183 layer.wkv = create_tensor(tn(LLM_TENSOR_ATTN_KV, "weight", i), {n_embd, n_embd_head}, 0);184 layer.attn_kv_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_NORM, "weight", i), {n_embd_head}, 0);185 layer.wo_a = create_tensor(tn(LLM_TENSOR_ATTN_OUT_A, "weight", i), {n_head * n_embd_head / o_groups, o_lora_rank, o_groups}, TENSOR_ALLOW_RESHAPE);186 layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT_B, "weight", i), {o_groups * o_lora_rank, n_embd}, 0);187 188 layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);189 layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);190 layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);191 layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0);192 layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);193 layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);194 195 layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);196 layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0);197 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);198 199 layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);200 layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);201 layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0);202 203 layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);204 layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd }, 0);205 layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);206 }207 return;208 }209 210 for (int i = 0; i < n_layer; ++i) {211 auto & layer = layers[i];212 213 layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);214 215 layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head }, 0);216 layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0);217 layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, 0);218 layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0);219 220 layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0);221 layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0);222 223 // optional per-head attention sinks (e.g. Nemotron DSpark)224 layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), { n_head }, TENSOR_NOT_REQUIRED);225 226 layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0);227 layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);228 layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);229 layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0);230 231 if (selector_meta) {232 const int64_t kernel = hparams.dflash_conv_kernel_size;233 const int64_t groups = n_embd / hparams.dflash_conv_group_size;234 const int64_t projected = 2 * kernel * groups;235 layer.dflash_attn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);236 layer.dflash_attn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);237 layer.dflash_ffn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);238 layer.dflash_ffn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);239 }240 }241}242 243template <>244ggml_tensor * llama_model_dflash::graph<true>::build_inp_embd_enc() const {245 const int64_t n_embd_inp = hparams.n_embd_inp_enc();246 auto inp_target = std::make_unique<llm_graph_input_embd>(n_embd_inp);247 248 inp_target->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);249 ggml_set_input(inp_target->embd);250 251 ggml_tensor * cur = inp_target->embd;252 cb(cur, "inp_embd", -1);253 254 res->add_input(std::move(inp_target));255 256 return cur;257}258 259// DFlash Encoder: processes target model features through feature fusion layer260template <>261llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {262 ggml_tensor * cur = build_inp_embd_enc();263 264 cur = build_lora_mm(model.fc, cur, model.fc_s);265 cb(cur, "fc_out", -1);266 267 cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);268 cb(cur, "enc_norm_out", -1);269 270 ggml_set_output(cur);271 res->t_h_nextn = cur;272 273 ggml_build_forward_expand(gf, cur);274}275 276// DSpark (DFlash + Markov & Confidence head): Markov bias on the draft logits, chained per block position277static void build_dspark_markov_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {278 ggml_context * ctx0 = g.ctx0;279 auto & res = g.res;280 281 ggml_tensor * w1 = model.dspark_markov_w1;282 ggml_tensor * w2 = model.dspark_markov_w2;283 GGML_ASSERT(w1 && w2 && "DSpark markov weights not loaded");284 285 // confidence head is optional286 const bool has_conf = model.dspark_conf_proj != nullptr;287 288 ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens]289 const int64_t n_vocab = base->ne[0];290 const int64_t n_tok = base->ne[1];291 292 const auto it = model.gguf_kv.find("dflash.block_size");293 GGML_ASSERT(it != model.gguf_kv.end() && "DSpark draft requires 'dflash.block_size' in GGUF metadata");294 const int64_t block_size = std::stoi(it->second);295 GGML_ASSERT(block_size > 0);296 297 // bonus anchor (SpecForge exports): slot 0 is a bonus token, not a prediction slot298 const auto it_anchor = model.gguf_kv.find("dflash.sample_from_anchor");299 const bool sample_from_anchor = it_anchor == model.gguf_kv.end() || it_anchor->second == "true";300 const int64_t i_draft_beg = sample_from_anchor ? 0 : 1;301 302 const int64_t n_blocks = g.ubatch.n_seqs_unq;303 GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks");304 // runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size305 const int64_t block_drafts = n_tok / n_blocks;306 if (block_drafts > block_size) {307 return;308 }309 310 // anchor (committed last) token of every block: token 0 of each block, i.e. a strided view311 const size_t token_stride = (size_t) block_drafts * tokens->nb[0];312 const size_t base_stride = (size_t) block_drafts * base->nb[1];313 314 ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0);315 prev = ggml_cont_1d(ctx0, prev, n_blocks);316 317 ggml_tensor * cat = nullptr;318 ggml_tensor * cat_conf = nullptr;319 320 if (!sample_from_anchor) {321 // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column322 cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));323 if (has_conf) {324 cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0)));325 }326 }327 328 // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final329 // token pick, not the Markov conditioning path330 for (int64_t i = i_draft_beg; i < block_drafts; ++i) {331 ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks]332 ggml_tensor * bias = g.build_lora_mm(w2, w1_prev, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks]333 if (model.d2t) {334 // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others)335 const int64_t n_draft_vocab = bias->ne[0];336 ggml_tensor * full = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_blocks), 0.0f);337 bias = ggml_set_rows(ctx0, full,338 ggml_reshape_3d(ctx0, bias, 1, n_draft_vocab, n_blocks),339 ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1));340 bias = ggml_reshape_2d(ctx0, bias, n_vocab, n_blocks);341 }342 343 // position i of every block: strided view [n_vocab, n_blocks]344 ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]);345 ggml_tensor * col = ggml_add(ctx0, base_i, bias);346 347 cat = cat ? ggml_concat(ctx0, cat, col, 1) : col;348 349 if (has_conf) {350 // confidence head input: predicts per-position acceptance351 ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok]352 // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks]353 ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,354 (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);355 ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0);356 ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);357 if (model.dspark_conf_proj_b) {358 conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);359 }360 conf = ggml_sigmoid(ctx0, conf);361 362 cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;363 }364 365 if (i + 1 < block_drafts) {366 prev = ggml_argmax(ctx0, col);367 }368 }369 370 // cat is position-major; restore ubatch block-major order371 ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, block_drafts);372 out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks]373 out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok);374 375 if (has_conf) {376 ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts);377 conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3));378 conf = ggml_reshape_2d(ctx0, conf, 1, n_tok);379 380 // note: broadcast the [1, n_tok] confidences to n_embd-wide rows to be able to reuse `llama_get_embeddings_nextn`381 conf = ggml_repeat(ctx0, conf, res->t_embd);382 res->t_h_nextn = conf;383 ggml_build_forward_expand(g.gf, conf);384 }385 386 res->t_logits = out;387 ggml_build_forward_expand(g.gf, out);388}389 390static ggml_tensor * build_dflash2_conv(391 llm_graph_context & g,392 ggml_tensor * hidden,393 ggml_tensor * dynamic,394 ggml_tensor * base,395 int side) {396 const auto & hparams = g.hparams;397 const int64_t hidden_size = hidden->ne[0];398 const int64_t n_tokens = hidden->ne[1];399 const int64_t n_blocks = g.ubatch.n_seqs_unq;400 const int64_t kernel_size = hparams.dflash_conv_kernel_size;401 const int64_t group_size = hparams.dflash_conv_group_size;402 const int64_t n_groups = hidden_size / group_size;403 404 GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);405 GGML_ASSERT(dynamic && base && side >= 0 && side < 2);406 407 const int64_t block_size = n_tokens / n_blocks;408 ggml_context * ctx0 = g.ctx0;409 // ggml_cont copies even when the tensor is already contiguous410 if (!ggml_is_contiguous(hidden) || hidden->ne[1] != n_tokens) {411 hidden = ggml_cont_2d(ctx0, hidden, hidden_size, n_tokens);412 }413 if (!ggml_is_contiguous(dynamic) || dynamic->ne[1] != n_tokens) {414 dynamic = ggml_cont_2d(ctx0, dynamic, dynamic->ne[0], n_tokens);415 }416 ggml_tensor * blocks = ggml_reshape_3d(ctx0, hidden, hidden_size, block_size, n_blocks);417 ggml_tensor * coeffs = ggml_reshape_4d(ctx0, dynamic, n_groups, kernel_size, 2, n_tokens);418 ggml_tensor * coeffs_side = ggml_view_3d(ctx0, coeffs, n_groups, kernel_size, n_tokens,419 coeffs->nb[1], coeffs->nb[3], side * coeffs->nb[2]);420 421 ggml_tensor * coeff_all = ggml_cont(ctx0, coeffs_side);422 coeff_all = ggml_reshape_4d(ctx0, coeff_all, 1, n_groups, kernel_size, n_tokens);423 coeff_all = ggml_repeat_4d(ctx0, coeff_all, group_size, n_groups, kernel_size, n_tokens);424 425 ggml_tensor * base_side = ggml_reshape_4d(ctx0,426 ggml_view_1d(ctx0, base, hidden_size * kernel_size, side * base->nb[2]),427 group_size, n_groups, kernel_size, 1);428 429 ggml_tensor * weight_all = ggml_add(ctx0, coeff_all, base_side);430 431 ggml_tensor * result = nullptr;432 for (int64_t tap = 0; tap < kernel_size; ++tap) {433 ggml_tensor * values = blocks;434 if (tap > 0) {435 ggml_tensor * zeros = ggml_fill(ctx0,436 ggml_new_tensor_3d(ctx0, hidden->type, hidden_size, std::min(tap, block_size), n_blocks), 0.0f);437 if (tap < block_size) {438 ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks,439 blocks->nb[1], blocks->nb[2], 0);440 values = ggml_concat(ctx0, zeros, previous, 1);441 } else {442 values = zeros;443 }444 }445 values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens);446 447 ggml_tensor * weight = ggml_reshape_2d(ctx0,448 ggml_cont(ctx0, ggml_view_4d(ctx0, weight_all, group_size, n_groups, 1, n_tokens,449 weight_all->nb[1], weight_all->nb[2], weight_all->nb[3], tap * weight_all->nb[2])),450 hidden_size, n_tokens);451 452 ggml_tensor * term = ggml_mul(ctx0, weight, values);453 result = result ? ggml_add(ctx0, result, term) : term;454 }455 return result;456}457 458// DFlash2 selector: top-k candidates per block position plus the pairwise459// transition scores, packed into the nextn output slot for the CPU-side walk.460static void build_dflash2_selector(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {461 ggml_context * ctx0 = g.ctx0;462 auto & res = g.res;463 464 const auto & hparams = g.hparams;465 const int64_t n_tokens = g.n_tokens;466 const int64_t n_embd = g.n_embd;467 468 const int64_t top_k = hparams.dflash_selector_top_k;469 const int64_t rank = hparams.dflash_selector_rank;470 const int64_t n_blocks = g.ubatch.n_seqs_unq;471 GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);472 GGML_ASSERT(res->t_logits->ne[1] == n_tokens);473 if (!tokens) {474 return;475 }476 477 const int64_t tokens_per_block = n_tokens / n_blocks;478 const int64_t block_size = std::min<int64_t>(tokens_per_block, hparams.dflash_block_size);479 const int64_t row_used = top_k + top_k * top_k;480 481 ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k);482 ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens);483 ggml_tensor * unary = ggml_reshape_2d(ctx0,484 ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens);485 ggml_tensor * gate = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd);486 487 // Everything below indexes [.., tokens_per_block, n_blocks]: the block488 // position varies fastest, sequences are the outer dimension.489 ggml_tensor * cand_blk = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks);490 ggml_tensor * unary_blk = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks);491 ggml_tensor * gate_blk = ggml_reshape_3d(ctx0, gate, rank, tokens_per_block, n_blocks);492 493 // a position's score reads only the candidate sets at pos-1 and pos, so a run494 // of positions has no internal dependency and scores in one batched matmul495 auto score_run = [&](int64_t beg_pos, int64_t n_pos, ggml_tensor * pred_ids) {496 ggml_tensor * cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, n_pos, n_blocks,497 cand_blk->nb[1], cand_blk->nb[2], beg_pos * cand_blk->nb[1]));498 ggml_tensor * unary_run = ggml_cont(ctx0, ggml_view_3d(ctx0, unary_blk, top_k, n_pos, n_blocks,499 unary_blk->nb[1], unary_blk->nb[2], beg_pos * unary_blk->nb[1]));500 ggml_tensor * gate_run = ggml_cont(ctx0, ggml_view_3d(ctx0, gate_blk, rank, n_pos, n_blocks,501 gate_blk->nb[1], gate_blk->nb[2], beg_pos * gate_blk->nb[1]));502 503 const int64_t n_pred = pred_ids->ne[0] / (n_pos * n_blocks);504 505 ggml_tensor * successor = ggml_reshape_4d(ctx0,506 ggml_get_rows(ctx0, model.dflash_selector_next, ggml_reshape_1d(ctx0, cand_run, top_k * n_pos * n_blocks)),507 rank, top_k, n_pos, n_blocks);508 ggml_tensor * predecessor = ggml_reshape_4d(ctx0,509 ggml_get_rows(ctx0, model.dflash_selector_prev, pred_ids),510 rank, n_pred, n_pos, n_blocks);511 512 ggml_tensor * gate_bcast = ggml_reshape_4d(ctx0, gate_run, rank, 1, n_pos, n_blocks);513 ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, gate_bcast, predecessor));514 ggml_tensor * score = ggml_mul_mat(ctx0, successor, cond);515 if (n_pred == 1) {516 score = ggml_repeat_4d(ctx0, score, top_k, top_k, n_pos, n_blocks);517 }518 ggml_tensor * unary_bcast = ggml_reshape_4d(ctx0, unary_run, top_k, 1, n_pos, n_blocks);519 score = ggml_add(ctx0, score, ggml_repeat(ctx0, unary_bcast, score));520 521 ggml_tensor * row = ggml_concat(ctx0,522 ggml_cast(ctx0, cand_run, GGML_TYPE_F32),523 ggml_reshape_3d(ctx0, score, top_k * top_k, n_pos, n_blocks), 0);524 return ggml_pad(ctx0, row, n_embd - row_used, 0, 0, 0);525 };526 527 ggml_tensor * packed = ggml_fill(ctx0,528 ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f);529 530 if (block_size > 1) {531 // Position 1 alone: its predecessor is the anchor token, one id per532 // sequence rather than a candidate set.533 ggml_tensor * anchor_ids = ggml_cont_1d(ctx0,534 ggml_view_2d(ctx0, tokens, 1, n_blocks, tokens_per_block * tokens->nb[0], 0), n_blocks);535 packed = ggml_concat(ctx0, packed, score_run(1, 1, anchor_ids), 1);536 }537 if (block_size > 2) {538 ggml_tensor * prev_ids = ggml_reshape_1d(ctx0,539 ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, block_size - 2, n_blocks,540 cand_blk->nb[1], cand_blk->nb[2], cand_blk->nb[1])),541 top_k * (block_size - 2) * n_blocks);542 packed = ggml_concat(ctx0, packed, score_run(2, block_size - 2, prev_ids), 1);543 }544 545 packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks);546 g.cb(packed, "dflash2_lattice", -1);547 res->t_h_nextn = packed;548 ggml_build_forward_expand(g.gf, packed);549}550 551// DFlash decoder, dual-mode by batch type:552// * embd batch -> fused target features: project + inject K/V into the cache.553// * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens554template <>555llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {556 const int64_t n_embd_inp = hparams.n_embd_inp_enc();557 const int64_t n_embd_head = hparams.n_embd_head_v();558 559 GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());560 561 ggml_tensor * inp_pos = build_inp_pos();562 563 // optional iSWA: pick the matching attention input564 const bool use_iswa = hparams.swa_type != LLAMA_SWA_TYPE_NONE;565 566 llm_graph_input_attn_kv * inp_attn = nullptr;567 llm_graph_input_attn_kv_iswa * inp_attn_iswa = nullptr;568 if (use_iswa) {569 inp_attn_iswa = build_attn_inp_kv_iswa();570 } else {571 inp_attn = build_attn_inp_kv();572 }573 574 const float kq_scale = 1.0f/sqrtf(float(n_embd_head));575 576 // drafts for M-RoPE targets use degenerate sections (temporal dim only)577 int sections[4];578 std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);579 580 auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos) {581 return rope_type == GGML_ROPE_TYPE_MROPE582 ? ggml_rope_multi(ctx0, cur, pos, nullptr,583 n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,584 ext_factor, attn_factor, beta_fast, beta_slow)585 : ggml_rope_ext(ctx0, cur, pos, nullptr,586 n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,587 ext_factor, attn_factor, beta_fast, beta_slow);588 };589 590 // KV cache injection591 if (ubatch.embd) {592 auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);593 594 inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);595 ggml_set_input(inp->embd);596 597 ggml_tensor * inp_target = inp->embd;598 cb(inp_target, "inp_target_features", -1);599 600 res->add_input(std::move(inp));601 602 // fuse the target features through the encoder603 ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);604 inp_g = build_norm(inp_g, model.output_norm_enc, NULL, LLM_NORM_RMS, -1);605 cb(inp_g, "inp_g_embeddings", -1);606 607 for (int il = 0; il < n_layer; ++il) {608 const auto & layer = model.layers[il];609 610 ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g, layer.wk_s);611 ggml_tensor * Vcur = build_lora_mm(layer.wv, inp_g, layer.wv_s);612 613 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);614 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);615 616 Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);617 Kcur = build_rope(Kcur, inp_pos);618 cb(Kcur, "Kcur_injected", il);619 cb(Vcur, "Vcur_injected", il);620 621 if (use_iswa) {622 // route each layer's K/V to its sub-cache: SWA layers -> sliding cache, full -> dense623 const bool is_swa = hparams.is_swa(il);624 const auto * kv = is_swa ? inp_attn_iswa->mctx->get_swa() : inp_attn_iswa->mctx->get_base();625 ggml_tensor * k_idxs = is_swa ? inp_attn_iswa->get_k_idxs_swa() : inp_attn_iswa->get_k_idxs();626 ggml_tensor * v_idxs = is_swa ? inp_attn_iswa->get_v_idxs_swa() : inp_attn_iswa->get_v_idxs();627 // rotate K/V into the cache's rotated space628 ggml_tensor * k_rot = is_swa ? inp_attn_iswa->self_k_rot_swa : inp_attn_iswa->self_k_rot;629 ggml_tensor * v_rot = is_swa ? inp_attn_iswa->self_v_rot_swa : inp_attn_iswa->self_v_rot;630 if (k_rot) {631 Kcur = llama_mul_mat_hadamard(ctx0, Kcur, k_rot);632 }633 if (v_rot) {634 Vcur = llama_mul_mat_hadamard(ctx0, Vcur, v_rot);635 }636 ggml_build_forward_expand(gf, kv->cpy_k(ctx0, Kcur, k_idxs, il));637 ggml_build_forward_expand(gf, kv->cpy_v(ctx0, Vcur, v_idxs, il));638 } else {639 // rotate K/V into the cache's rotated space640 if (inp_attn->self_k_rot) {641 Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp_attn->self_k_rot);642 }643 if (inp_attn->self_v_rot) {644 Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp_attn->self_v_rot);645 }646 ggml_build_forward_expand(gf, inp_attn->mctx->cpy_k(ctx0, Kcur, inp_attn->get_k_idxs(), il));647 ggml_build_forward_expand(gf, inp_attn->mctx->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));648 }649 }650 651 res->t_embd = inp_g;652 653 ggml_build_forward_expand(gf, inp_g);654 return;655 }656 657 // tok_embd from the target model (shared via ctx_other)658 auto * tok_embd = model.tok_embd;659 if (tok_embd == nullptr) {660 GGML_ASSERT(cparams.ctx_other != nullptr);661 const auto * model_other = llama_get_model(cparams.ctx_other);662 663 GGML_ASSERT(model_other->tok_embd != nullptr && "DFlash decoder requires the target model's token embeddings");664 tok_embd = model_other->tok_embd;665 }666 667 auto inp = std::make_unique<llm_graph_input_embd>(n_embd);668 669 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);670 ggml_set_input(inp->tokens);671 res->t_inp_tokens = inp->tokens;672 673 ggml_tensor * inp_tokens = inp->tokens;674 675 ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);676 cb(inpL, "inp_noise_embd", -1);677 678 res->add_input(std::move(inp));679 680 for (int il = 0; il < n_layer; ++il) {681 const auto & layer = model.layers[il];682 683 ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il);684 cb(noise_norm, "noise_norm", il);685 686 ggml_tensor * attn_dynamic = nullptr;687 if (layer.dflash_attn_conv_proj) {688 attn_dynamic = build_lora_mm(layer.dflash_attn_conv_proj, noise_norm);689 noise_norm = build_dflash2_conv(*this, noise_norm, attn_dynamic, layer.dflash_attn_conv_base, 0);690 cb(noise_norm, "attn_conv_in", il);691 }692 693 ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm, layer.wq_s);694 ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm, layer.wk_s);695 ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm, layer.wv_s);696 697 Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);698 Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);699 Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);700 701 Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il);702 Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);703 704 Qcur = build_rope(Qcur, inp_pos);705 Kcur = build_rope(Kcur, inp_pos);706 cb(Qcur, "Qcur", il);707 cb(Kcur, "Kcur", il);708 cb(Vcur, "Vcur", il);709 710 // cache-aware, non-causal attention711 ggml_tensor * cur = use_iswa712 ? build_attn(inp_attn_iswa, layer.wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il)713 : build_attn(inp_attn, layer.wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il);714 715 if (attn_dynamic) {716 cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);717 cb(cur, "attn_conv_out", il);718 }719 720 ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);721 cb(ffn_inp, "ffn_inp", il);722 723 cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);724 cb(cur, "ffn_norm", il);725 726 ggml_tensor * ffn_dynamic = nullptr;727 if (layer.dflash_ffn_conv_proj) {728 ffn_dynamic = build_lora_mm(layer.dflash_ffn_conv_proj, cur);729 cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 0);730 cb(cur, "ffn_conv_in", il);731 }732 733 cur = build_ffn(cur,734 layer.ffn_up, NULL, layer.ffn_up_s,735 layer.ffn_gate, NULL, layer.ffn_gate_s,736 layer.ffn_down, NULL, layer.ffn_down_s,737 NULL,738 LLM_FFN_SILU, LLM_FFN_PAR, il);739 cb(cur, "ffn_out", il);740 741 if (ffn_dynamic) {742 cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 1);743 cb(cur, "ffn_conv_out", il);744 }745 746 cur = ggml_add(ctx0, cur, ffn_inp);747 cb(cur, "l_out", il);748 749 inpL = cur;750 }751 752 ggml_tensor * cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1);753 cb(cur, "result_norm", -1);754 755 res->t_embd = cur;756 757 // lm_head from the target model (shared via ctx_other)758 auto * output = model.output;759 auto * output_s = model.output_s;760 if (output == nullptr) {761 GGML_ASSERT(cparams.ctx_other != nullptr);762 const auto * model_other = llama_get_model(cparams.ctx_other);763 GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection");764 output = model_other->output;765 output_s = model_other->output_s;766 }767 768 cur = build_lora_mm(output, cur, output_s);769 770 // DFlash2 feeds these logits to the selector, so they need the target's output771 // transforms; DFlash1 and DSpark read them through the sampler instead772 if (model.dflash_selector_hidden) {773 if (hparams.f_logit_scale != 0.0f) {774 cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);775 }776 if (hparams.f_final_logit_softcapping > 0.0f) {777 cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);778 cur = ggml_tanh(ctx0, cur);779 cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);780 }781 }782 783 // reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t784 if (model.d2t) {785 const int64_t n_draft_vocab = cur->ne[0];786 const int64_t n_outputs = cur->ne[1];787 const int64_t n_vocab = (int64_t) model.vocab.n_tokens();788 789 GGML_ASSERT(model.d2t->type == GGML_TYPE_I64);790 GGML_ASSERT(model.d2t->ne[0] == n_draft_vocab);791 792 ggml_tensor * logits = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_outputs), -INFINITY);793 cur = ggml_set_rows(ctx0, logits,794 ggml_reshape_3d(ctx0, cur, 1, n_draft_vocab, n_outputs),795 ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1));796 cur = ggml_reshape_2d(ctx0, cur, n_vocab, n_outputs);797 }798 cb(cur, "result_output", -1);799 res->t_logits = cur;800 801 ggml_build_forward_expand(gf, cur);802 803 // DSpark: bias the draft logits with the Markov head804 if (model.dspark_markov_w1) {805 build_dspark_markov_head(*this, model, inp_tokens);806 }807 808 if (model.dflash_selector_hidden) {809 build_dflash2_selector(*this, model, inp_tokens);810 }811}812 813// DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above):814// * embd batch -> project main_x through each stage's wkv and inject K into the ring cache815// * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads816llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_graph_params & params) :817 llama_model_deepseek4::graph(params) {818 const int64_t n_embd_inp = hparams.n_embd_inp_enc();819 const int64_t n_embd_head = hparams.n_embd_head_k();820 const int64_t n_embd_head_rope = hparams.n_rot();821 const int64_t n_embd_head_nope = n_embd_head - n_embd_head_rope;822 823 ggml_tensor * inp_pos = build_inp_pos();824 825 llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();826 827 // KV cache injection: fused target features from the encoder828 if (ubatch.embd) {829 auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);830 831 inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);832 ggml_set_input(inp->embd);833 834 ggml_tensor * inp_target = inp->embd;835 cb(inp_target, "inp_target_features", -1);836 837 res->add_input(std::move(inp));838 839 // fuse the target features through the encoder840 ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);841 inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);842 cb(inp_g, "inp_g_embeddings", -1);843 844 for (int il = 0; il < n_layer; ++il) {845 const auto & layer = model.layers[il];846 847 // main-track KV: kv_norm(wkv(main_x)) with rope on the trailing dims, same848 // rope parameters as the uncompressed layers in build_attention_impl849 ggml_tensor * kv = build_lora_mm(layer.wkv, inp_g);850 kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il);851 kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens);852 853 kv = ggml_rope_ext(ctx0, kv, inp_pos, nullptr, n_embd_head_rope, rope_type, 0,854 freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);855 kv = ggml_rope_set_offset(kv, n_embd_head_nope);856 cb(kv, "kv_injected", il);857 858 if (inp_attn->self_k_rot_swa) {859 kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);860 }861 ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));862 }863 864 res->t_embd = inp_g;865 866 ggml_build_forward_expand(gf, inp_g);867 return;868 }869 870 // tok_embd from the target model (shared via ctx_other)871 auto * tok_embd = model.tok_embd;872 if (tok_embd == nullptr) {873 GGML_ASSERT(cparams.ctx_other != nullptr);874 const auto * model_other = llama_get_model(cparams.ctx_other);875 876 GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");877 tok_embd = model_other->tok_embd;878 }879 880 auto inp = std::make_unique<llm_graph_input_embd>(n_embd);881 882 inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);883 ggml_set_input(inp->tokens);884 885 ggml_tensor * inp_tokens = inp->tokens;886 887 ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);888 cb(inpL, "inp_noise_embd", -1);889 890 res->add_input(std::move(inp));891 892 const int64_t hc = hparams.dsv4_hc_mult;893 inpL = ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens);894 inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);895 cb(inpL, "hc_init", -1);896 897 for (int il = 0; il < n_layer; ++il) {898 const auto & layer = model.layers[il];899 900 ggml_tensor * residual = inpL;901 ggml_tensor * post = nullptr;902 ggml_tensor * comb = nullptr;903 904 ggml_tensor * cur = build_hc_pre(inpL,905 layer.hc_attn_fn,906 layer.hc_attn_scale,907 layer.hc_attn_base,908 &post, &comb, il);909 cb(cur, "hc_attn_pre", il);910 911 cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);912 cb(cur, "attn_norm", il);913 914 cur = build_attention(model, inp_attn, cur, inp_pos, il);915 916 inpL = build_hc_post(cur, residual, post, comb, il);917 cb(inpL, "hc_attn_post", il);918 919 residual = inpL;920 cur = build_hc_pre(inpL,921 layer.hc_ffn_fn,922 layer.hc_ffn_scale,923 layer.hc_ffn_base,924 &post, &comb, il);925 cb(cur, "hc_ffn_pre", il);926 927 cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);928 cb(cur, "ffn_norm", il);929 930 ggml_tensor * moe_out = build_moe_ffn(cur,931 layer.ffn_gate_inp,932 layer.ffn_up_exps,933 layer.ffn_gate_exps,934 layer.ffn_down_exps,935 layer.ffn_exp_probs_b,936 n_expert, hparams.n_expert_used(),937 LLM_FFN_SILU, hparams.expert_weights_norm,938 hparams.expert_weights_scale,939 (llama_expert_gating_func_type) hparams.expert_gating_func,940 il);941 cb(moe_out, "ffn_moe_out", il);942 943 ggml_tensor * ffn_shexp = build_ffn(cur,944 layer.ffn_up_shexp, nullptr, nullptr,945 layer.ffn_gate_shexp, nullptr, nullptr,946 layer.ffn_down_shexp, nullptr, nullptr,947 nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);948 cb(ffn_shexp, "ffn_shexp", il);949 950 cur = ggml_add(ctx0, moe_out, ffn_shexp);951 cb(cur, "ffn_out", il);952 953 inpL = build_hc_post(cur, residual, post, comb, il);954 cb(inpL, "l_out", il);955 }956 957 ggml_tensor * cur = build_hc_head(inpL, model.hc_head_fn, model.hc_head_scale, model.hc_head_base);958 cb(cur, "hc_head", -1);959 960 // confidence head input: the reference scores the pre-norm collapsed hidden state961 res->t_embd = cur;962 963 cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);964 cb(cur, "result_norm", -1);965 966 // lm_head from the target model (shared via ctx_other)967 auto * output = model.output;968 auto * output_s = model.output_s;969 if (output == nullptr) {970 GGML_ASSERT(cparams.ctx_other != nullptr);971 const auto * model_other = llama_get_model(cparams.ctx_other);972 GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");973 output = model_other->output;974 output_s = model_other->output_s;975 }976 977 cur = build_lora_mm(output, cur, output_s);978 cb(cur, "result_output", -1);979 res->t_logits = cur;980 981 ggml_build_forward_expand(gf, cur);982 983 if (model.dspark_markov_w1) {984 build_dspark_markov_head(*this, model, inp_tokens);985 }986}987 988std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const llm_graph_params & params) const {989 switch (params.gtype) {990 case LLM_GRAPH_TYPE_ENCODER:991 return std::make_unique<graph<true>>(*this, params);992 case LLM_GRAPH_TYPE_DEFAULT:993 case LLM_GRAPH_TYPE_DECODER:994 if (hparams.dsv4_hc_mult > 0) {995 return std::make_unique<graph_dsv4>(*this, params);996 }997 return std::make_unique<graph<false>>(*this, params);998 default:999 GGML_ABORT("invalid graph type");1000 };1001}1002