From 27462403a0cd7b9ca8943281cbe02566daefe2b3 Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Mon, 10 Aug 2026 18:25:23 -0700 Subject: [PATCH] dflash: port GIDD log-SNR conditioning and prefer drafter-own embeddings Two drafter-side deltas our fork's standalone dspark arch carried, re-applied on the upstream dflash implementation (Path A of the dspark re-port): 1. GIDD log-SNR conditioning (LogSnrEmbed): some GIDD-trained drafters ship a log_snr_embed module, a sinusoidal featurization of a per-position log-SNR value through a 2-layer SiLU MLP, added to the draft noise embedding before the backbone. Anchor position of each block at max_log_snr, masked positions at min_log_snr; the feature matrix is a pure function of quantities known at graph-build time, so it is precomputed host-side and staged through a new llm_graph_input_dspark_logsnr input. New optional metadata: dflash.log_snr_conditioning (bool) plus dflash.min_log_snr/max_log_snr (required, validated finite and max > min once the flag is set); new tensors log_snr_fc1/fc2.{weight,bias}, REQUIRED when the flag is set so a broken conversion fails loudly instead of silently running unconditioned. 2. Optional drafter-own token_embd/output: loaded as TENSOR_NOT_REQUIRED; the decoder graph already prefers model-own tensors and only borrows the target's via ctx_other when absent. Carrying the drafter's own full-precision embeddings/head protects accept rate when the target is heavily quantized, and makes the drafter loadable without ctx_other during memory fitting. Drafters without either feature convert and load exactly as before. --- src/llama-arch.cpp | 7 ++++ src/llama-arch.h | 5 +++ src/llama-graph.cpp | 11 +++++ src/llama-graph.h | 21 ++++++++++ src/llama-hparams.h | 6 +++ src/llama-model.h | 6 +++ src/models/dflash.cpp | 95 +++++++++++++++++++++++++++++++++++++++++++ 7 files changed, 151 insertions(+) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 836cfade226..ed33ef3eb0c 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -317,6 +317,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_TARGET_LAYERS, "%s.target_layers" }, { LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" }, + { LLM_KV_LOG_SNR_CONDITIONING, "%s.log_snr_conditioning" }, + { LLM_KV_MIN_LOG_SNR, "%s.min_log_snr" }, + { LLM_KV_MAX_LOG_SNR, "%s.max_log_snr" }, { LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" }, { LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" }, @@ -620,6 +623,8 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, + { LLM_TENSOR_DSPARK_LOG_SNR_FC1, "log_snr_fc1" }, + { LLM_TENSOR_DSPARK_LOG_SNR_FC2, "log_snr_fc2" }, }; // declare information about the model weight tensors: @@ -878,6 +883,8 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, {LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DSPARK_LOG_SNR_FC1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DSPARK_LOG_SNR_FC2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index 49c2a6ac399..98d983ebef0 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -363,6 +363,9 @@ enum llm_kv { LLM_KV_TARGET_LAYERS, LLM_KV_TARGET_HIDDEN_SIZE, + LLM_KV_LOG_SNR_CONDITIONING, + LLM_KV_MIN_LOG_SNR, + LLM_KV_MAX_LOG_SNR, LLM_KV_NORM_BEFORE_RESIDUAL, LLM_KV_NORM_BEFORE_FC, @@ -628,6 +631,8 @@ enum llm_tensor { LLM_TENSOR_DSPARK_MARKOV_W1, LLM_TENSOR_DSPARK_MARKOV_W2, LLM_TENSOR_DSPARK_CONF_PROJ, + LLM_TENSOR_DSPARK_LOG_SNR_FC1, + LLM_TENSOR_DSPARK_LOG_SNR_FC2, }; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index b5507afde64..43fa2062663 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -64,6 +64,17 @@ static bool can_reuse_kq_mask( // impl +void llm_graph_input_dspark_logsnr::set_input(const llama_ubatch * ubatch) { + // ignores the ubatch entirely: v_feat was precomputed at graph-build time + // from n_tokens/block_drafts/min_log_snr/max_log_snr, nothing here depends + // on the current ubatch. + GGML_UNUSED(ubatch); + if (feat && !v_feat.empty()) { + GGML_ASSERT((int64_t) v_feat.size() == ggml_nelements(feat)); + ggml_backend_tensor_set(feat, v_feat.data(), 0, ggml_nbytes(feat)); + } +} + void llm_graph_input_embd::set_input(const llama_ubatch * ubatch) { if (ubatch->token) { const int64_t n_tokens = ubatch->n_tokens; diff --git a/src/llama-graph.h b/src/llama-graph.h index 32d8d395aa4..5eebf0855a5 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -136,6 +136,27 @@ class llm_graph_input_embd : public llm_graph_input_i { const int64_t n_embd = 0; }; +// dspark GIDD log-SNR conditioning (LogSnrEmbed): the sinusoidal feature +// matrix fed into log_snr_fc1/fc2. Carries no external staged state -- the +// per-position log-SNR pattern (anchor position of each block at max_log_snr, +// mask positions at min_log_snr) and its sinusoidal featurization are a pure +// function of n_tokens/block_drafts/min_log_snr/max_log_snr, all known at +// graph-build time, so the caller precomputes the full [128, n_tokens] +// feature matrix once and this class just stages it as an input (ggml's +// no_alloc graph context means even build-time-constant data has to go +// through set_input(), same as everything else here). +class llm_graph_input_dspark_logsnr : public llm_graph_input_i { +public: + llm_graph_input_dspark_logsnr(std::vector feat) : v_feat(std::move(feat)) {} + virtual ~llm_graph_input_dspark_logsnr() = default; + + void set_input(const llama_ubatch * ubatch) override; + + ggml_tensor * feat = nullptr; // F32 [128, n_tokens] + + std::vector v_feat; +}; + // similar to llm_graph_input_embd but with an additional hidden state input class llm_graph_input_embd_h : public llm_graph_input_i { public: diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 6e8336c9874..c0ab4249982 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -195,6 +195,12 @@ struct llama_hparams { // e.g. the eagle3 encoder fuses target_layers * target_hidden features uint32_t n_embd_inp_enc_impl = 0; + // dspark GIDD log-SNR conditioning (LogSnrEmbed); off unless the drafter + // GGUF sets .log_snr_conditioning + bool dspark_log_snr_conditioning = false; + float dspark_min_log_snr = 0.0f; + float dspark_max_log_snr = 0.0f; + // output embedding dimension (0 = use n_embd) uint32_t n_embd_out_impl = 0; diff --git a/src/llama-model.h b/src/llama-model.h index 6b9e94a0a69..5f956a6654c 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -613,6 +613,12 @@ struct llama_model { struct ggml_tensor * dspark_conf_proj = nullptr; struct ggml_tensor * dspark_conf_proj_b = nullptr; + // dspark GIDD log-SNR conditioning (present only when hparams.dspark_log_snr_conditioning) + struct ggml_tensor * dspark_log_snr_fc1_w = nullptr; // [128 -> n_embd] + struct ggml_tensor * dspark_log_snr_fc1_b = nullptr; + struct ggml_tensor * dspark_log_snr_fc2_w = nullptr; // [n_embd -> n_embd] + struct ggml_tensor * dspark_log_snr_fc2_b = nullptr; + // unified vector to store target-model extracted layer ids in eagle3, dflash, etc. std::vector target_layer_ids; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index daff6e78f1c..5c98eab7965 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -14,6 +14,23 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd; + // dspark GIDD log-SNR conditioning (drafters trained with the GIDD bundle) -- + // absent on every other drafter, must default off + ml.get_key(LLM_KV_LOG_SNR_CONDITIONING, hparams.dspark_log_snr_conditioning, false); + if (hparams.dspark_log_snr_conditioning) { + // required once the flag is set: silently defaulting either bound to 0.0f + // would collapse (max_snr - min_snr) to zero in the featurizer and fill + // the conditioning input with NaNs instead of failing to load + ml.get_key(LLM_KV_MIN_LOG_SNR, hparams.dspark_min_log_snr, true); + ml.get_key(LLM_KV_MAX_LOG_SNR, hparams.dspark_max_log_snr, true); + if (!std::isfinite(hparams.dspark_min_log_snr) || !std::isfinite(hparams.dspark_max_log_snr)) { + throw std::runtime_error("dspark log-SNR conditioning: min/max_log_snr must be finite"); + } + if (!(hparams.dspark_max_log_snr > hparams.dspark_min_log_snr)) { + throw std::runtime_error("dspark log-SNR conditioning: max_log_snr must be greater than min_log_snr"); + } + } + LLAMA_LOG_INFO("%s: DFlash extract_layers = [", __func__); for (size_t i = 0; i < target_layer_ids.size(); ++i) { LLAMA_LOG_INFO("%d%s", target_layer_ids[i], i + 1 < target_layer_ids.size() ? ", " : ""); @@ -96,6 +113,27 @@ 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); } + // dspark GIDD log-SNR conditioning (LogSnrEmbed): unlike the markov head + // above, this is built into the decoder graph and changes the draft + // embedding every forward pass, so if the GGUF says log_snr_conditioning + // is on, the weights are REQUIRED -- a missing tensor here is a broken + // conversion, not something to silently degrade past (an earlier converter + // silently dropped these 4 tensors and invalidated a measurement) + if (hparams.dspark_log_snr_conditioning) { + const int64_t n_freq = 128; // matches LogSnrEmbed.NUM_FREQ_FEATURES + dspark_log_snr_fc1_w = create_tensor(tn(LLM_TENSOR_DSPARK_LOG_SNR_FC1, "weight"), { n_freq, n_embd }, 0); + dspark_log_snr_fc1_b = create_tensor(tn(LLM_TENSOR_DSPARK_LOG_SNR_FC1, "bias"), { n_embd }, 0); + dspark_log_snr_fc2_w = create_tensor(tn(LLM_TENSOR_DSPARK_LOG_SNR_FC2, "weight"), { n_embd, n_embd }, 0); + dspark_log_snr_fc2_b = create_tensor(tn(LLM_TENSOR_DSPARK_LOG_SNR_FC2, "bias"), { n_embd }, 0); + } + + // optional drafter-own token embeddings and lm_head: when present the + // decoder graph prefers them over the target's (borrowed via ctx_other). + // Load-bearing for low-bit targets: borrowing means the drafter runs the + // target's quantized tok_embd/output, which measurably costs accept rate. + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm @@ -417,6 +455,63 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra res->add_input(std::move(inp)); + // dspark GIDD log-SNR conditioning (LogSnrEmbed): added to the draft noise + // embedding BEFORE the layer loop, matching the training-side reference + // implementation. Per-position log-SNR is the fixed round-1 inference + // convention: the anchor position of each block (row 0 of every block; the + // ubatch is block-major, same layout build_dspark_markov_head relies on) + // at max_log_snr, every masked position at min_log_snr. + if (hparams.dspark_log_snr_conditioning) { + GGML_ASSERT(model.dspark_log_snr_fc1_w && model.dspark_log_snr_fc2_w && + model.dspark_log_snr_fc1_b && model.dspark_log_snr_fc2_b); + + const int64_t n_blocks = ubatch.n_seqs_unq; + GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0 && "log-SNR conditioning requires equal-size blocks"); + const int64_t block_drafts = n_tokens / n_blocks; + + const int64_t n_freq = 128; + const int64_t half = n_freq / 2; + const float min_snr = hparams.dspark_min_log_snr; + const float max_snr = hparams.dspark_max_log_snr; + + // host-side: exact port of LogSnrEmbed.forward's featurization fused + // with the anchor/mask pattern above. A pure function of + // n_tokens/block_drafts/min_log_snr/max_log_snr, all known here, so + // precomputing on the host (rather than chaining + // ggml_arange/ggml_sin/ggml_cos in-graph) keeps it directly auditable + // line-for-line against the python reference. + std::vector feat((size_t) (n_freq * n_tokens)); + for (int64_t pos = 0; pos < n_tokens; ++pos) { + const float log_snr = (pos % block_drafts == 0) ? max_snr : min_snr; + const float t = (log_snr - min_snr) / (max_snr - min_snr) * 1000.0f; + for (int64_t i = 0; i < half; ++i) { + const float freq = expf(-logf(10000.0f) * (float) i / (float) half); + const float angle = t * freq; + feat[(size_t) (pos * n_freq + i)] = sinf(angle); + feat[(size_t) (pos * n_freq + half + i)] = cosf(angle); + } + } + + auto logsnr_input = std::make_unique(std::move(feat)); + logsnr_input->feat = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_freq, n_tokens); + ggml_set_input(logsnr_input->feat); + ggml_set_name(logsnr_input->feat, "dspark_log_snr_feat"); + ggml_tensor * snr_feat = logsnr_input->feat; + res->add_input(std::move(logsnr_input)); + + ggml_tensor * snr_hidden = build_lora_mm(model.dspark_log_snr_fc1_w, snr_feat); + snr_hidden = ggml_add(ctx0, snr_hidden, model.dspark_log_snr_fc1_b); + snr_hidden = ggml_silu(ctx0, snr_hidden); + cb(snr_hidden, "dspark_log_snr_fc1", -1); + + ggml_tensor * snr_embed = build_lora_mm(model.dspark_log_snr_fc2_w, snr_hidden); + snr_embed = ggml_add(ctx0, snr_embed, model.dspark_log_snr_fc2_b); + cb(snr_embed, "dspark_log_snr_fc2", -1); + + inpL = ggml_add(ctx0, inpL, snr_embed); + cb(inpL, "dspark_draft_embd_snr", -1); + } + for (int il = 0; il < n_layer; ++il) { const auto & layer = model.layers[il];