From cc231cb0da565440cf6a3e5b55dfeba477972cb6 Mon Sep 17 00:00:00 2001 From: JamePeng Date: Sun, 30 Aug 2026 16:34:39 +0800 Subject: [PATCH] dflash: pass missing NVFP4 scales to attention operations (#28000) - DFlash2 NVFP4 draft models produced almost no accepted speculative tokens because the Q, K, V, and output projection scales were not passed to the corresponding graph operations. --- src/models/dflash.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index f2c7d1d24..037723a90 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -615,8 +615,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra for (int il = 0; il < n_layer; ++il) { const auto & layer = model.layers[il]; - ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g); - ggml_tensor * Vcur = build_lora_mm(layer.wv, inp_g); + ggml_tensor * Kcur = build_lora_mm(layer.wk, inp_g, layer.wk_s); + ggml_tensor * Vcur = build_lora_mm(layer.wv, inp_g, layer.wv_s); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); @@ -698,9 +698,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra 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); + ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm, layer.wq_s); + ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm, layer.wk_s); + ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm, layer.wv_s); Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); @@ -717,8 +717,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra // cache-aware, non-causal attention ggml_tensor * cur = use_iswa - ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il) - : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il); + ? build_attn(inp_attn_iswa, layer.wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il) + : build_attn(inp_attn, layer.wo, NULL, layer.wo_s, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il); if (attn_dynamic) { cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);