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:
@@ -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) {
|
for (int il = 0; il < n_layer; ++il) {
|
||||||
const auto & layer = model.layers[il];
|
const auto & layer = model.layers[il];
|
||||||
|
|
||||||
ggml_tensor * Kcur = build_lora_mm(layer.wk, 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);
|
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);
|
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);
|
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);
|
cb(noise_norm, "attn_conv_in", il);
|
||||||
}
|
}
|
||||||
|
|
||||||
ggml_tensor * Qcur = build_lora_mm(layer.wq, 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);
|
ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm, layer.wk_s);
|
||||||
ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm);
|
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);
|
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);
|
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
|
// cache-aware, non-causal attention
|
||||||
ggml_tensor * cur = use_iswa
|
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_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, NULL, 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) {
|
if (attn_dynamic) {
|
||||||
cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);
|
cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);
|
||||||
|
|||||||
Reference in New Issue
Block a user