From c478bc5bab9e3097bf48612e1b6619295fea3dbe Mon Sep 17 00:00:00 2001 From: bri-prism <288398250+bri-prism@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:51:24 -0700 Subject: [PATCH] dflash: let the log-SNR input opt into graph reuse The input inherited the default can_reuse() == false, so turning on log-SNR conditioning rebuilt the whole decoder graph for every ubatch. v_feat is a function of n_tokens and n_seqs_unq, which llm_graph_params::allow_reuse already compares, and of min/max_log_snr, which are fixed per model. So the input can reuse whenever those checks pass. Acceptance is unchanged at 65.891%, mean length 3.58, on the same drafter and prompts. --- src/llama-graph.cpp | 6 ++++++ src/llama-graph.h | 2 ++ 2 files changed, 8 insertions(+) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 211d37273e1..49c865e7fcf 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -75,6 +75,12 @@ void llm_graph_input_dspark_logsnr::set_input(const llama_ubatch * ubatch) { } } +bool llm_graph_input_dspark_logsnr::can_reuse(const llm_graph_params & params) { + // v_feat is a function of n_tokens and n_seqs_unq, both already checked by + // llm_graph_params::allow_reuse, and of min/max_log_snr, fixed per model + return feat && feat->ne[1] == params.ubatch.n_tokens; +} + 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 b8efa7f0e3b..492574f69a1 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -152,6 +152,8 @@ class llm_graph_input_dspark_logsnr : public llm_graph_input_i { void set_input(const llama_ubatch * ubatch) override; + bool can_reuse(const llm_graph_params & params) override; + ggml_tensor * feat = nullptr; // F32 [128, n_tokens] std::vector v_feat;