nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

DFlash: support gpt-oss drafts with attention bias (#2061) 

Adds support for the z-lab gpt-oss DFlash drafts (e.g. z-lab/gpt-oss-20b-DFlash),
which use a Qwen3 backbone trained with attention_bias=true.

- dflash-draft loader: create optional attention bias tensors bq/bk/bv/bo
  (TENSOR_NOT_REQUIRED); the Qwen3.5/MiMo drafts have none and are unaffected.
- build_dflash: add those biases at the q/k/v/o projections (cross-context K/V
  and the noise block), each guarded so bias-free drafts are unchanged.
- Narrow the DFlash graph contract validator to reject only fused qkv biases
  (bqkv/bqk/bkv), which the graph still does not implement, and remove a dead
  duplicate of the validator in llama.cpp (the live copy is in llama-dflash.cpp).

The gpt-oss tokenizer recognition needed to convert the draft is a separate,
general gpt-oss conversion fix submitted independently.

Verified on gpt-oss-20b: coherent output and 42% draft acceptance
(~2.95 accepted tokens/cycle) at cross_ctx=128.

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
This commit is contained in:
Joel Farthing 2026-06-30 02:02:18 -05:00 committed by GitHub
parent f74a6fb87b
commit 615c3e11b8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 21 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/graphs/build_dflash.cpp
View File

@ -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);

14
src/llama-dflash.cpp
View File

@ -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;
}

8
src/llama-load-tensors.cpp
View File

@ -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);

10
src/llama.cpp
View File

@ -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