* spec : add DFlash2 support (local convolution + candidate selector) (#27342) * support DFlash2 * Add p_min in DFlash2 Assisted-by: Claude Opus 5 * Revert unnecessary changes Assisted-by: Claude Opus 5 * Revert draft sampling in rejection sampling Assisted-by: Claude Opus 5 * Refactor code structure Assisted-by: Claude Opus 5 * Delete embedding scaling Assisted-by: Claude Opus 5 * Gate output transforms on DFlash2 Assisted-by: Claude Opus 5 * Optimize Dflash 2 cost Assisted-by: Claude Opus 5 * Avoid using atoi Assisted-by: Claude Opus 5 * Modify comments Assisted-by: Claude Opus 5 * Move llama_model_dflash_selector_top_k to llama-ext.h Assisted-by: Claude Opus 5 * Formatting Assisted-by: Claude Opus 5 * Apply patch to fix the mrope bug Assisted-by: Claude Opus 5 * fix ci Assisted-by: Claude Opus 5 * Fix graph number calculation Assisted-by: Claude Opus 5 * rename hid and unary Assisted-by: Claude Opus 5 --------- Co-authored-by: Jian Chen <jianchen0311@gmail.com> Co-authored-by: Xuan-Son Nguyen <son@huggingface.co> * revert top-k.cu changes --------- Co-authored-by: Zihan Zhang <tiancaizhangdaxian@sjtu.edu.cn> Co-authored-by: Jian Chen <jianchen0311@gmail.com>
This commit is contained in:
Xuan-Son Nguyen
co-authored by
Jian Chen
Zihan Zhang
Jian Chen
Zihan Zhang
parent
58546250cf
commit
b10f9ca58c
+265
-15
@@ -7,6 +7,18 @@
|
||||
void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
|
||||
|
||||
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
|
||||
ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, false);
|
||||
hparams.f_final_logit_softcapping = 0.0f;
|
||||
ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
|
||||
|
||||
// drafts for M-RoPE targets carry degenerate sections [n_rot/2, 0, 0, 0]
|
||||
ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);
|
||||
|
||||
ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false);
|
||||
ml.get_key(LLM_KV_DFLASH_CONV_KERNEL_SIZE, hparams.dflash_conv_kernel_size, false);
|
||||
ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE, hparams.dflash_conv_group_size, false);
|
||||
ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK, hparams.dflash_selector_rank, false);
|
||||
ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K, hparams.dflash_selector_top_k, false);
|
||||
|
||||
if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) {
|
||||
throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata");
|
||||
@@ -112,6 +124,29 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
|
||||
LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);
|
||||
}
|
||||
|
||||
const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight");
|
||||
if (selector_meta) {
|
||||
const int64_t rank = hparams.dflash_selector_rank;
|
||||
if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 ||
|
||||
hparams.dflash_conv_kernel_size <= 0 || hparams.dflash_conv_group_size <= 0) {
|
||||
throw std::runtime_error("DFlash2 model is missing conv/selector metadata");
|
||||
}
|
||||
if (n_embd % hparams.dflash_conv_group_size != 0) {
|
||||
throw std::runtime_error("DFlash2 hidden size must be divisible by conv_group_size");
|
||||
}
|
||||
if (n_embd < hparams.dflash_selector_top_k * (hparams.dflash_selector_top_k + 1)) {
|
||||
throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice");
|
||||
}
|
||||
|
||||
dflash_selector_prev = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV, "weight"), { rank, n_vocab }, 0);
|
||||
dflash_selector_next = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT, "weight"), { rank, n_vocab }, 0);
|
||||
dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0);
|
||||
|
||||
LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__,
|
||||
hparams.dflash_conv_kernel_size, hparams.dflash_conv_group_size,
|
||||
hparams.dflash_selector_rank, hparams.dflash_selector_top_k);
|
||||
}
|
||||
|
||||
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
|
||||
fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED);
|
||||
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
|
||||
@@ -188,6 +223,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
|
||||
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);
|
||||
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);
|
||||
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0);
|
||||
|
||||
if (selector_meta) {
|
||||
const int64_t kernel = hparams.dflash_conv_kernel_size;
|
||||
const int64_t groups = n_embd / hparams.dflash_conv_group_size;
|
||||
const int64_t projected = 2 * kernel * groups;
|
||||
layer.dflash_attn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);
|
||||
layer.dflash_attn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);
|
||||
layer.dflash_ffn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);
|
||||
layer.dflash_ffn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -346,6 +391,167 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
|
||||
ggml_build_forward_expand(g.gf, out);
|
||||
}
|
||||
|
||||
static ggml_tensor * build_dflash2_conv(
|
||||
llm_graph_context & g,
|
||||
ggml_tensor * hidden,
|
||||
ggml_tensor * dynamic,
|
||||
ggml_tensor * base,
|
||||
int side) {
|
||||
const auto & hparams = g.hparams;
|
||||
const int64_t hidden_size = hidden->ne[0];
|
||||
const int64_t n_tokens = hidden->ne[1];
|
||||
const int64_t n_blocks = g.ubatch.n_seqs_unq;
|
||||
const int64_t kernel_size = hparams.dflash_conv_kernel_size;
|
||||
const int64_t group_size = hparams.dflash_conv_group_size;
|
||||
const int64_t n_groups = hidden_size / group_size;
|
||||
|
||||
GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
|
||||
GGML_ASSERT(dynamic && base && side >= 0 && side < 2);
|
||||
|
||||
const int64_t block_size = n_tokens / n_blocks;
|
||||
ggml_context * ctx0 = g.ctx0;
|
||||
// ggml_cont copies even when the tensor is already contiguous
|
||||
if (!ggml_is_contiguous(hidden) || hidden->ne[1] != n_tokens) {
|
||||
hidden = ggml_cont_2d(ctx0, hidden, hidden_size, n_tokens);
|
||||
}
|
||||
if (!ggml_is_contiguous(dynamic) || dynamic->ne[1] != n_tokens) {
|
||||
dynamic = ggml_cont_2d(ctx0, dynamic, dynamic->ne[0], n_tokens);
|
||||
}
|
||||
ggml_tensor * blocks = ggml_reshape_3d(ctx0, hidden, hidden_size, block_size, n_blocks);
|
||||
ggml_tensor * coeffs = ggml_reshape_4d(ctx0, dynamic, n_groups, kernel_size, 2, n_tokens);
|
||||
ggml_tensor * coeffs_side = ggml_view_3d(ctx0, coeffs, n_groups, kernel_size, n_tokens,
|
||||
coeffs->nb[1], coeffs->nb[3], side * coeffs->nb[2]);
|
||||
|
||||
ggml_tensor * coeff_all = ggml_cont(ctx0, coeffs_side);
|
||||
coeff_all = ggml_reshape_4d(ctx0, coeff_all, 1, n_groups, kernel_size, n_tokens);
|
||||
coeff_all = ggml_repeat_4d(ctx0, coeff_all, group_size, n_groups, kernel_size, n_tokens);
|
||||
|
||||
ggml_tensor * base_side = ggml_reshape_4d(ctx0,
|
||||
ggml_view_1d(ctx0, base, hidden_size * kernel_size, side * base->nb[2]),
|
||||
group_size, n_groups, kernel_size, 1);
|
||||
|
||||
ggml_tensor * weight_all = ggml_add(ctx0, coeff_all, base_side);
|
||||
|
||||
ggml_tensor * result = nullptr;
|
||||
for (int64_t tap = 0; tap < kernel_size; ++tap) {
|
||||
ggml_tensor * values = blocks;
|
||||
if (tap > 0) {
|
||||
ggml_tensor * zeros = ggml_fill(ctx0,
|
||||
ggml_new_tensor_3d(ctx0, hidden->type, hidden_size, std::min(tap, block_size), n_blocks), 0.0f);
|
||||
if (tap < block_size) {
|
||||
ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks,
|
||||
blocks->nb[1], blocks->nb[2], 0);
|
||||
values = ggml_concat(ctx0, zeros, previous, 1);
|
||||
} else {
|
||||
values = zeros;
|
||||
}
|
||||
}
|
||||
values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens);
|
||||
|
||||
ggml_tensor * weight = ggml_reshape_2d(ctx0,
|
||||
ggml_cont(ctx0, ggml_view_4d(ctx0, weight_all, group_size, n_groups, 1, n_tokens,
|
||||
weight_all->nb[1], weight_all->nb[2], weight_all->nb[3], tap * weight_all->nb[2])),
|
||||
hidden_size, n_tokens);
|
||||
|
||||
ggml_tensor * term = ggml_mul(ctx0, weight, values);
|
||||
result = result ? ggml_add(ctx0, result, term) : term;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// DFlash2 selector: top-k candidates per block position plus the pairwise
|
||||
// transition scores, packed into the nextn output slot for the CPU-side walk.
|
||||
static void build_dflash2_selector(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {
|
||||
ggml_context * ctx0 = g.ctx0;
|
||||
auto & res = g.res;
|
||||
|
||||
const auto & hparams = g.hparams;
|
||||
const int64_t n_tokens = g.n_tokens;
|
||||
const int64_t n_embd = g.n_embd;
|
||||
|
||||
const int64_t top_k = hparams.dflash_selector_top_k;
|
||||
const int64_t rank = hparams.dflash_selector_rank;
|
||||
const int64_t n_blocks = g.ubatch.n_seqs_unq;
|
||||
GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
|
||||
GGML_ASSERT(res->t_logits->ne[1] == n_tokens);
|
||||
if (!tokens) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int64_t tokens_per_block = n_tokens / n_blocks;
|
||||
const int64_t block_size = std::min<int64_t>(tokens_per_block, hparams.dflash_block_size);
|
||||
const int64_t row_used = top_k + top_k * top_k;
|
||||
|
||||
ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k);
|
||||
ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens);
|
||||
ggml_tensor * unary = ggml_reshape_2d(ctx0,
|
||||
ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens);
|
||||
ggml_tensor * gate = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd);
|
||||
|
||||
// Everything below indexes [.., tokens_per_block, n_blocks]: the block
|
||||
// position varies fastest, sequences are the outer dimension.
|
||||
ggml_tensor * cand_blk = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks);
|
||||
ggml_tensor * unary_blk = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks);
|
||||
ggml_tensor * gate_blk = ggml_reshape_3d(ctx0, gate, rank, tokens_per_block, n_blocks);
|
||||
|
||||
// a position's score reads only the candidate sets at pos-1 and pos, so a run
|
||||
// of positions has no internal dependency and scores in one batched matmul
|
||||
auto score_run = [&](int64_t beg_pos, int64_t n_pos, ggml_tensor * pred_ids) {
|
||||
ggml_tensor * cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, n_pos, n_blocks,
|
||||
cand_blk->nb[1], cand_blk->nb[2], beg_pos * cand_blk->nb[1]));
|
||||
ggml_tensor * unary_run = ggml_cont(ctx0, ggml_view_3d(ctx0, unary_blk, top_k, n_pos, n_blocks,
|
||||
unary_blk->nb[1], unary_blk->nb[2], beg_pos * unary_blk->nb[1]));
|
||||
ggml_tensor * gate_run = ggml_cont(ctx0, ggml_view_3d(ctx0, gate_blk, rank, n_pos, n_blocks,
|
||||
gate_blk->nb[1], gate_blk->nb[2], beg_pos * gate_blk->nb[1]));
|
||||
|
||||
const int64_t n_pred = pred_ids->ne[0] / (n_pos * n_blocks);
|
||||
|
||||
ggml_tensor * successor = ggml_reshape_4d(ctx0,
|
||||
ggml_get_rows(ctx0, model.dflash_selector_next, ggml_reshape_1d(ctx0, cand_run, top_k * n_pos * n_blocks)),
|
||||
rank, top_k, n_pos, n_blocks);
|
||||
ggml_tensor * predecessor = ggml_reshape_4d(ctx0,
|
||||
ggml_get_rows(ctx0, model.dflash_selector_prev, pred_ids),
|
||||
rank, n_pred, n_pos, n_blocks);
|
||||
|
||||
ggml_tensor * gate_bcast = ggml_reshape_4d(ctx0, gate_run, rank, 1, n_pos, n_blocks);
|
||||
ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, gate_bcast, predecessor));
|
||||
ggml_tensor * score = ggml_mul_mat(ctx0, successor, cond);
|
||||
if (n_pred == 1) {
|
||||
score = ggml_repeat_4d(ctx0, score, top_k, top_k, n_pos, n_blocks);
|
||||
}
|
||||
ggml_tensor * unary_bcast = ggml_reshape_4d(ctx0, unary_run, top_k, 1, n_pos, n_blocks);
|
||||
score = ggml_add(ctx0, score, ggml_repeat(ctx0, unary_bcast, score));
|
||||
|
||||
ggml_tensor * row = ggml_concat(ctx0,
|
||||
ggml_cast(ctx0, cand_run, GGML_TYPE_F32),
|
||||
ggml_reshape_3d(ctx0, score, top_k * top_k, n_pos, n_blocks), 0);
|
||||
return ggml_pad(ctx0, row, n_embd - row_used, 0, 0, 0);
|
||||
};
|
||||
|
||||
ggml_tensor * packed = ggml_fill(ctx0,
|
||||
ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f);
|
||||
|
||||
if (block_size > 1) {
|
||||
// Position 1 alone: its predecessor is the anchor token, one id per
|
||||
// sequence rather than a candidate set.
|
||||
ggml_tensor * anchor_ids = ggml_cont_1d(ctx0,
|
||||
ggml_view_2d(ctx0, tokens, 1, n_blocks, tokens_per_block * tokens->nb[0], 0), n_blocks);
|
||||
packed = ggml_concat(ctx0, packed, score_run(1, 1, anchor_ids), 1);
|
||||
}
|
||||
if (block_size > 2) {
|
||||
ggml_tensor * prev_ids = ggml_reshape_1d(ctx0,
|
||||
ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, block_size - 2, n_blocks,
|
||||
cand_blk->nb[1], cand_blk->nb[2], cand_blk->nb[1])),
|
||||
top_k * (block_size - 2) * n_blocks);
|
||||
packed = ggml_concat(ctx0, packed, score_run(2, block_size - 2, prev_ids), 1);
|
||||
}
|
||||
|
||||
packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks);
|
||||
g.cb(packed, "dflash2_lattice", -1);
|
||||
res->t_h_nextn = packed;
|
||||
ggml_build_forward_expand(g.gf, packed);
|
||||
}
|
||||
|
||||
// DFlash decoder, dual-mode by batch type:
|
||||
// * embd batch -> fused target features: project + inject K/V into the cache.
|
||||
// * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens
|
||||
@@ -370,6 +576,20 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
|
||||
const float kq_scale = 1.0f/sqrtf(float(n_embd_head));
|
||||
|
||||
// drafts for M-RoPE targets use degenerate sections (temporal dim only)
|
||||
int sections[4];
|
||||
std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections);
|
||||
|
||||
auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos) {
|
||||
return rope_type == GGML_ROPE_TYPE_MROPE
|
||||
? ggml_rope_multi(ctx0, cur, pos, nullptr,
|
||||
n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow)
|
||||
: ggml_rope_ext(ctx0, cur, pos, nullptr,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
};
|
||||
|
||||
// KV cache injection
|
||||
if (ubatch.embd) {
|
||||
auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
|
||||
@@ -392,11 +612,7 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
|
||||
|
||||
Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
|
||||
Kcur = ggml_rope_ext(
|
||||
ctx0, Kcur, inp_pos, nullptr,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow
|
||||
);
|
||||
Kcur = build_rope(Kcur, inp_pos);
|
||||
cb(Kcur, "Kcur_injected", il);
|
||||
cb(Vcur, "Vcur_injected", il);
|
||||
|
||||
@@ -450,6 +666,7 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
|
||||
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
|
||||
ggml_set_input(inp->tokens);
|
||||
res->t_inp_tokens = inp->tokens;
|
||||
|
||||
ggml_tensor * inp_tokens = inp->tokens;
|
||||
|
||||
@@ -464,6 +681,13 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(noise_norm, "noise_norm", il);
|
||||
|
||||
ggml_tensor * attn_dynamic = nullptr;
|
||||
if (layer.dflash_attn_conv_proj) {
|
||||
attn_dynamic = build_lora_mm(layer.dflash_attn_conv_proj, noise_norm);
|
||||
noise_norm = build_dflash2_conv(*this, noise_norm, attn_dynamic, layer.dflash_attn_conv_base, 0);
|
||||
cb(noise_norm, "attn_conv_in", il);
|
||||
}
|
||||
|
||||
ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm);
|
||||
ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm);
|
||||
ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm);
|
||||
@@ -475,16 +699,8 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il);
|
||||
Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
|
||||
|
||||
Qcur = ggml_rope_ext(
|
||||
ctx0, Qcur, inp_pos, nullptr,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow
|
||||
);
|
||||
Kcur = ggml_rope_ext(
|
||||
ctx0, Kcur, inp_pos, nullptr,
|
||||
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
|
||||
ext_factor, attn_factor, beta_fast, beta_slow
|
||||
);
|
||||
Qcur = build_rope(Qcur, inp_pos);
|
||||
Kcur = build_rope(Kcur, inp_pos);
|
||||
cb(Qcur, "Qcur", il);
|
||||
cb(Kcur, "Kcur", il);
|
||||
cb(Vcur, "Vcur", il);
|
||||
@@ -494,12 +710,24 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il)
|
||||
: build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
|
||||
|
||||
if (attn_dynamic) {
|
||||
cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);
|
||||
cb(cur, "attn_conv_out", il);
|
||||
}
|
||||
|
||||
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
|
||||
cb(ffn_inp, "ffn_inp", il);
|
||||
|
||||
cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
|
||||
cb(cur, "ffn_norm", il);
|
||||
|
||||
ggml_tensor * ffn_dynamic = nullptr;
|
||||
if (layer.dflash_ffn_conv_proj) {
|
||||
ffn_dynamic = build_lora_mm(layer.dflash_ffn_conv_proj, cur);
|
||||
cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 0);
|
||||
cb(cur, "ffn_conv_in", il);
|
||||
}
|
||||
|
||||
cur = build_ffn(cur,
|
||||
layer.ffn_up, NULL, layer.ffn_up_s,
|
||||
layer.ffn_gate, NULL, layer.ffn_gate_s,
|
||||
@@ -508,6 +736,11 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
LLM_FFN_SILU, LLM_FFN_PAR, il);
|
||||
cb(cur, "ffn_out", il);
|
||||
|
||||
if (ffn_dynamic) {
|
||||
cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 1);
|
||||
cb(cur, "ffn_conv_out", il);
|
||||
}
|
||||
|
||||
cur = ggml_add(ctx0, cur, ffn_inp);
|
||||
cb(cur, "l_out", il);
|
||||
|
||||
@@ -532,6 +765,19 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
|
||||
cur = build_lora_mm(output, cur, output_s);
|
||||
|
||||
// DFlash2 feeds these logits to the selector, so they need the target's output
|
||||
// transforms; DFlash1 and DSpark read them through the sampler instead
|
||||
if (model.dflash_selector_hidden) {
|
||||
if (hparams.f_logit_scale != 0.0f) {
|
||||
cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);
|
||||
}
|
||||
if (hparams.f_final_logit_softcapping > 0.0f) {
|
||||
cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
|
||||
cur = ggml_tanh(ctx0, cur);
|
||||
cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
|
||||
}
|
||||
}
|
||||
|
||||
// reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t
|
||||
if (model.d2t) {
|
||||
const int64_t n_draft_vocab = cur->ne[0];
|
||||
@@ -556,6 +802,10 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
if (model.dspark_markov_w1) {
|
||||
build_dspark_markov_head(*this, model, inp_tokens);
|
||||
}
|
||||
|
||||
if (model.dflash_selector_hidden) {
|
||||
build_dflash2_selector(*this, model, inp_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
// DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above):
|
||||
|
||||
Reference in New Issue
Block a user