diff --git a/src/graphs/build_dflash.cpp b/src/graphs/build_dflash.cpp index 4f6c7cd0..38df86a7 100644 --- a/src/graphs/build_dflash.cpp +++ b/src/graphs/build_dflash.cpp @@ -39,6 +39,7 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { GGML_ASSERT(il < (int32_t) lctx.dflash.kv.v_ctx_cache.size()); ggml_tensor * Kcur_ctx_proj = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, fused_target); + if (model.layers[il].bk) { Kcur_ctx_proj = ggml_add(ctx0, Kcur_ctx_proj, model.layers[il].bk); } cb(Kcur_ctx_proj, "dflash_kv_k_proj", il); ggml_tensor * Kcur_ctx = ggml_reshape_3d(ctx0, Kcur_ctx_proj, n_embd_head_k, n_head_kv, update_rows); @@ -54,6 +55,7 @@ ggml_cgraph * llm_build_context::build_dflash_kv_cache() { cb(Kcur_ctx, "dflash_kv_k_physical", il); ggml_tensor * Vcur_ctx = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, fused_target); + if (model.layers[il].bv) { Vcur_ctx = ggml_add(ctx0, Vcur_ctx, model.layers[il].bv); } cb(Vcur_ctx, "dflash_kv_v_proj", il); if (std::abs(hparams.f_attn_v_scale - 1.0f) > 1e-4f) { Vcur_ctx = ggml_scale(ctx0, Vcur_ctx, hparams.f_attn_v_scale); @@ -234,6 +236,9 @@ ggml_cgraph * llm_build_context::build_dflash() { ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur); ggml_tensor * Kcur_noise = llm_build_lora_mm(lctx, ctx0, model.layers[il].wk, cur); ggml_tensor * Vcur_noise = llm_build_lora_mm(lctx, ctx0, model.layers[il].wv, cur); + if (model.layers[il].bq) { Qcur = ggml_add(ctx0, Qcur, model.layers[il].bq); } + if (model.layers[il].bk) { Kcur_noise = ggml_add(ctx0, Kcur_noise, model.layers[il].bk); } + if (model.layers[il].bv) { Vcur_noise = ggml_add(ctx0, Vcur_noise, model.layers[il].bv); } cb(Qcur, "Qcur", il); cb(Kcur_noise, "Kcur_noise", il); cb(Vcur_noise, "Vcur_noise", il); @@ -251,7 +256,7 @@ ggml_cgraph * llm_build_context::build_dflash() { Kcur_noise = ggml_reshape_3d(ctx0, Kcur_noise, n_embd_head_k, n_head_kv, n_tokens); Kcur_noise = llm_build_norm(ctx0, Kcur_noise, hparams, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, cb, il); - cb(Qcur, "Kcur_normed", il); + cb(Kcur_noise, "Kcur_normed", il); Kcur_noise = ggml_rope_ext(ctx0, Kcur_noise, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); @@ -338,6 +343,7 @@ ggml_cgraph * llm_build_context::build_dflash() { cb(cur, "flash_attn_reshaped", il); cur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wo, cur); + if (model.layers[il].bo) { cur = ggml_add(ctx0, cur, model.layers[il].bo); } cb(cur, "kqv_out", il); cur = ggml_add(ctx0, cur, inpSA); diff --git a/src/llama-dflash.cpp b/src/llama-dflash.cpp index f419aab0..ae630b31 100644 --- a/src/llama-dflash.cpp +++ b/src/llama-dflash.cpp @@ -215,13 +215,11 @@ static void llama_graph_compute_sched( } static bool dflash_layer_has_attention_bias(const llama_layer & layer) { - return layer.bq != nullptr || - layer.bk != nullptr || - layer.bv != nullptr || - layer.bo != nullptr || - layer.bqkv != nullptr || - layer.bqk != nullptr || - layer.bkv != nullptr; + // build_dflash() now applies separate q/k/v/o attention biases (bq/bk/bv/bo). + // Only fused qkv-style biases remain unimplemented in the DFlash graph. + return layer.bqkv != nullptr || + layer.bqk != nullptr || + layer.bkv != nullptr; } static bool validate_dflash_graph_contract(const llama_context & lctx) { @@ -309,7 +307,7 @@ static bool validate_dflash_graph_contract(const llama_context & lctx) { } if (dflash_layer_has_attention_bias(model.layers[il])) { - LLAMA_LOG_ERROR("%s: DFlash graph does not implement attention bias tensors, but layer %d requires them\n", + LLAMA_LOG_ERROR("%s: DFlash graph implements only separate q/k/v/o attention bias; layer %d uses an unsupported fused qkv bias\n", __func__, il); return false; } diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 6b85bdd1..77b70d16 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -2315,6 +2315,14 @@ bool create_tensors_helper::create_dflash_tensors(const LLM_TN & tn) { layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + + // optional attention biases: present on gpt-oss-style DFlash drafts (attention_bias=true), + // absent on the Qwen3.5/MiMo drafts -> mark not-required so both load. + layer.bq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "bias", i), {n_embd_head_k * n_head}, llama_model_loader::TENSOR_NOT_REQUIRED); + layer.bk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "bias", i), {n_embd_gqa}, llama_model_loader::TENSOR_NOT_REQUIRED); + layer.bv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "bias", i), {n_embd_gqa}, llama_model_loader::TENSOR_NOT_REQUIRED); + layer.bo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, llama_model_loader::TENSOR_NOT_REQUIRED); + layer.attn_sinks = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, llama_model_loader::TENSOR_NOT_REQUIRED); layer.ffn_gate = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); diff --git a/src/llama.cpp b/src/llama.cpp index 5b665c74..f8c13eb0 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -5140,16 +5140,6 @@ static bool prepare_mtp_graph_inputs( return true; } -static bool dflash_layer_has_attention_bias(const llama_layer & layer) { - return layer.bq != nullptr || - layer.bk != nullptr || - layer.bv != nullptr || - layer.bo != nullptr || - layer.bqkv != nullptr || - layer.bqk != nullptr || - layer.bkv != nullptr; -} - // decode a batch of tokens by evaluating the transformer // // - lctx: llama context