prism-v6: dspark Path A, converge on upstream dflash with log-SNR conditioning and drafter-own embeddings

src/llama-arch.cpp

@@ -317,6 +317,9 @@ static const std::map<llm_kv, const char *> 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, const char *> 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, llm_tensor_info> 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) {}

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,
 };
 
 

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;

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<float> 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<float> 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:

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 <arch>.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;
 

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<int32_t> target_layer_ids;
 

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<false>::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<float> 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<llm_graph_input_dspark_logsnr>(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];