| 37 | #include <cstdlib> |
| 38 | #include <cmath> |
| 39 | #include <cstdio> |
| 40 | |
| 41 | namespace dflash::common { |
| 42 | |
| 43 | // Feature fusion shared by the legacy one-shot graph and the cached-KV |
| 44 | // builders: optional per-capture RMSNorm slices, fc projection, hidden_norm. |
| 45 | // Row-independent, so it is bit-identical whether run over the full window |
| 46 | // or over appended rows only. |
| 47 | static ggml_tensor * draft_fuse_features( |
| 48 | ggml_context * ctx, |
| 49 | const DraftWeights & w, |
| 50 | ggml_tensor * target_hidden_cat, |
| 51 | int n_rows, |
| 52 | bool disable_aux_hidden_norms) { |
| 53 | const float eps = DFLASH27B_RMS_EPS; |
| 54 | ggml_tensor * thc = target_hidden_cat; |
| 55 | if (!disable_aux_hidden_norms && !w.aux_hidden_norms.empty()) { |
| 56 | ggml_tensor * aux_cat = nullptr; |
| 57 | const size_t elem_sz = ggml_element_size(target_hidden_cat); |
| 58 | for (size_t i = 0; i < w.aux_hidden_norms.size(); i++) { |
| 59 | ggml_tensor * slice = ggml_view_3d(ctx, target_hidden_cat, |
| 60 | w.n_embd, n_rows, 1, |
| 61 | target_hidden_cat->nb[1], target_hidden_cat->nb[2], |
| 62 | i * (size_t)w.n_embd * elem_sz); |
| 63 | slice = ggml_rms_norm(ctx, slice, eps); |
| 64 | slice = ggml_mul(ctx, slice, w.aux_hidden_norms[i]); |
| 65 | aux_cat = aux_cat ? ggml_concat(ctx, aux_cat, slice, 0) : slice; |
| 66 | } |
| 67 | thc = aux_cat; |
| 68 | } |
| 69 | ggml_tensor * target_feat = ggml_mul_mat(ctx, w.fc, thc); |
| 70 | target_feat = ggml_rms_norm(ctx, target_feat, eps); |
| 71 | target_feat = ggml_mul (ctx, target_feat, w.hidden_norm); |
| 72 | return target_feat; |
| 73 | } |
| 74 | |
| 75 | DraftGraphOutputs build_draft_graph( |
| 76 | ggml_context * ctx, |
| 77 | const DraftWeights & w, |
| 78 | const DraftGraphInputs & in) { |
| 79 | |
| 80 | const int q_len = w.block_size; |
| 81 | const int ctx_len = in.ctx_len; |
| 82 | const int n_head = w.n_head; |
| 83 | const int n_kv = w.n_head_kv; |
| 84 | const int head_dim = w.head_dim; |
| 85 | const float eps = DFLASH27B_RMS_EPS; |
| 86 | const float rope_base = w.rope_theta; |
| 87 | |
| 88 | // ── 1. Feature fusion: target_feat = rms_norm(fc @ target_hidden_cat, hidden_norm) |
| 89 | // fc: [5*hidden, hidden] (ggml: ne[0]=5*hidden, ne[1]=hidden) |
| 90 | // target_hidden_cat: [5*hidden, ctx_len, 1] |
| 91 | // Result: [hidden, ctx_len, 1] |
| 92 | static const bool disable_aux_hidden_norms = |
| 93 | std::getenv("DFLASH_DISABLE_DRAFT_AUX_NORMS") != nullptr; |
| 94 | static const bool disable_attn_gate = |
| 95 | std::getenv("DFLASH_DISABLE_DRAFT_ATTN_GATE") != nullptr; |
| 96 | static const bool disable_swa = |
no outgoing calls