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.
This commit is contained in:
JamePeng 2026-08-30 16:34:39 +08:00 committed by GitHub
parent bebc9350ec
commit cc231cb0da
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -615,8 +615,8 @@ llama_model_dflash::graph<false>::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<false>::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<false>::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);