From 873e5d8e39feb34a376e0efd01bf3f665dfffeb5 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Fri, 21 Aug 2026 18:54:29 +0200 Subject: [PATCH] model: use ggml_rope_set_offset() (#27382) * model: use ggml_rope_set_offset() * partially apply to deepseek2 --- docs/development/HOWTO-add-model.md | 13 +++++ src/models/deepseek2.cpp | 40 ++++++++------ src/models/deepseek4.cpp | 83 +++++------------------------ src/models/dflash.cpp | 12 +---- src/models/minicpm3.cpp | 28 ++++------ src/models/plm.cpp | 28 ++++------ 6 files changed, 71 insertions(+), 133 deletions(-) diff --git a/docs/development/HOWTO-add-model.md b/docs/development/HOWTO-add-model.md index fcc87f165..31b3f2686 100644 --- a/docs/development/HOWTO-add-model.md +++ b/docs/development/HOWTO-add-model.md @@ -166,6 +166,19 @@ Examples: - Some models require scaling the input position. For example, `[0, 1, 2, ...]` becomes `[0, 0.5, 1, ...]`. In this case, you can provide the scaling via `freq_scale = 0.5f`. - Some models use learned RoPE frequencies instead of relying on `powf(freq_base, -2.0 * i / n_dims)`. In this case, you can provide the learned frequencies via the `rope_freqs` tensor (corresponding to the `c` argument in `ggml_rope_ext`), then set `freq_base = 1.0f`. An important note is that `rope_freqs` in GGML is the **inverse** (`theta = pos[i] / rope_freqs`), so you may need to invert `rope_freqs` during conversion. +### Rotating only a part of the head + +Many models rotate only a part of each head and leave the rest untouched (often called the "nope" part). Do not build this with views plus `ggml_concat`, it's not efficient. Both layouts can be done with a single RoPE op: + +- `[rope|nope]`, rotated dims first: pass `n_dims` smaller than the head size to `ggml_rope_ext`. Dims from `n_dims` to the end are copied as-is. +- `[nope|rope]`, rotated dims last: call `ggml_rope_set_offset(cur, n_offs)` on the result of the RoPE, where `n_offs` is the size of the leading untouched part. Dims outside `[n_offs, n_offs + n_dims)` are copied as-is. + +`n_offs` must be even, `n_offs + n_dims` must fit in the row, and vision RoPE is not supported. Note that the frequencies are computed relative to the rotated window. + +Example: DeepSeek-V4 uses `[nope|rope]` for its query, key and compressed KV tensors, so `src/models/deepseek4.cpp` ropes the whole tensor and then calls `ggml_rope_set_offset(cur, n_embd_head_nope)`. + +Exception: some models apply an extra op to the `nope` part, for example `deepseek32.cpp`, and may not use this optimization. While RoPE can be applied selectively to a part of the head, the extra op may not, so these models still need views plus `ggml_concat`. + ## GGUF specification https://github.com/ggml-org/ggml/blob/master/docs/gguf.md diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp index ba90c0d07..e0e537e00 100644 --- a/src/models/deepseek2.cpp +++ b/src/models/deepseek2.cpp @@ -524,17 +524,9 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p q = ggml_mul_mat(ctx0, model.layers[il].wq, cur); cb(q, "q", il); } - // split into {n_embd_head_qk_nope, n_head, n_tokens} - ggml_tensor * q_nope = - ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k), - ggml_row_size(q->type, n_embd_head_k) * n_head, 0); - cb(q_nope, "q_nope", il); - - // and {n_embd_head_qk_rope, n_head, n_tokens} - ggml_tensor * q_pe = ggml_view_3d( - ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k), - ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope)); - cb(q_pe, "q_pe", il); + // {n_embd_head_k, n_head, n_tokens} + q = ggml_reshape_3d(ctx0, q, n_embd_head_k, n_head, n_tokens); + cb(q, "q", il); ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur); cb(kv_cmpr_pe, "kv_cmpr_pe", il); @@ -552,10 +544,6 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); cb(k_pe, "k_pe", il); - q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); - cb(q_pe, "q_pe", il); - k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); cb(k_pe, "k_pe", il); @@ -564,6 +552,20 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p cb(kv_cmpr, "kv_cmpr", il); if (is_mla) { + // split into {n_embd_head_qk_nope, n_head, n_tokens} + ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, + q->nb[1], q->nb[2], 0); + cb(q_nope, "q_nope", il); + + // and {n_embd_head_qk_rope, n_head, n_tokens} + ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, + q->nb[1], q->nb[2], ggml_row_size(q->type, n_embd_head_qk_nope)); + cb(q_pe, "q_pe", il); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(q_pe, "q_pe", il); + // {n_embd_head_qk_nope, n_tokens, n_head} q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); cb(q_nope, "q_nope_perm", il); @@ -623,10 +625,14 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p Vcur = ggml_cont(ctx0, Vcur); cb(Vcur, "Vcur_cont", il); - ggml_tensor * Qcur = ggml_concat(ctx0, q_nope, q_pe, 0); + // RoPE is applied to the trailing dims only + ggml_tensor * Qcur = ggml_rope_ext(ctx0, q, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + Qcur = ggml_rope_set_offset(Qcur, n_embd_head_qk_nope); cb(Qcur, "Qcur", il); - ggml_tensor * Kcur = ggml_concat(ctx0, k_nope, ggml_repeat(ctx0, k_pe, q_pe), 0); + ggml_tensor * Kcur = ggml_concat(ctx0, k_nope, + ggml_repeat_4d(ctx0, k_pe, n_embd_head_qk_rope, n_head, n_tokens, 1), 0); cb(Kcur, "Kcur", il); if (inp_attn_scale) { diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 366ca2e54..1c278e435 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -501,21 +501,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_hca_compressed_kv_from_state( comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il); cb(comp, name, il); - ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp, n_embd_head_nope, 1, n_blocks, - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head), - 0); - ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp, n_embd_head_rope, 1, n_blocks, - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head_nope)); - - comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig, + comp = ggml_rope_ext(ctx0, comp, comp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig, hparams.dsv4_compress_rope_base, freq_scale, ext_factor, dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow); - cb(comp_pe, name, il); - - comp = ggml_concat(ctx0, comp_nope, comp_pe, 0); + comp = ggml_rope_set_offset(comp, n_embd_head_nope); cb(comp, name, il); return comp; @@ -585,21 +574,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_overlap_compressed_kv_from_sta comp = build_norm(comp, norm, nullptr, LLM_NORM_RMS, il); cb(comp, name, il); - ggml_tensor * comp_nope = ggml_view_3d(ctx0, comp, n_embd_head_nope, 1, n_blocks, - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head), - 0); - ggml_tensor * comp_pe = ggml_view_3d(ctx0, comp, n_embd_head_rope, 1, n_blocks, - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head), - ggml_row_size(comp->type, n_embd_head_nope)); - - comp_pe = ggml_rope_ext(ctx0, comp_pe, comp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig, + comp = ggml_rope_ext(ctx0, comp, comp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig, hparams.dsv4_compress_rope_base, freq_scale, ext_factor, dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow); - cb(comp_pe, name, il); - - comp = ggml_concat(ctx0, comp_nope, comp_pe, 0); + comp = ggml_rope_set_offset(comp, n_embd_head_nope); cb(comp, name, il); return comp; @@ -628,21 +606,12 @@ ggml_tensor * llama_model_deepseek4::graph::build_lid_top_k( indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, nt); cb(indexer_q, "lid_q", il); - ggml_tensor * indexer_q_nope = ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, nt, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head)*n_indexer_head, - 0); - ggml_tensor * indexer_q_pe = ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, nt, - ggml_row_size(indexer_q->type, n_embd_indexer_head), - ggml_row_size(indexer_q->type, n_embd_indexer_head)*n_indexer_head, - ggml_row_size(indexer_q->type, n_embd_indexer_head_nope)); - - indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_embd_indexer_head_rope, + indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, nullptr, n_embd_indexer_head_rope, rope_type, n_ctx_orig, hparams.dsv4_compress_rope_base, freq_scale, ext_factor, dsv4_rope_attn_factor(freq_scale, ext_factor), beta_fast, beta_slow); - cb(indexer_q_pe, "lid_q_pe", il); + indexer_q = ggml_rope_set_offset(indexer_q, n_embd_indexer_head_nope); + cb(indexer_q, "lid_q_rope", il); - indexer_q = ggml_concat(ctx0, indexer_q_nope, indexer_q_pe, 0); indexer_q = llama_mul_mat_hadamard(ctx0, indexer_q, inp_lid.k_rot); cb(indexer_q, "lid_q_rot", il); @@ -945,18 +914,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl( q = ggml_rms_norm(ctx0, q, norm_rms_eps); cb(q, "q_norm", il); - ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_nope, n_head, nt, - ggml_row_size(q->type, n_embd_head), - ggml_row_size(q->type, n_embd_head)*n_head, - 0); - ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_rope, n_head, nt, - ggml_row_size(q->type, n_embd_head), - ggml_row_size(q->type, n_embd_head)*n_head, - ggml_row_size(q->type, n_embd_head_nope)); - q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, + q = ggml_rope_ext(ctx0, q, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); - cb(q_pe, "q_pe", il); - q = ggml_concat(ctx0, q_nope, q_pe, 0); + q = ggml_rope_set_offset(q, n_embd_head_nope); cb(q, "q", il); ggml_tensor * kv = build_lora_mm(layer.wkv, cur); @@ -964,18 +924,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl( kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, nt); cb(kv, "kv_norm", il); - ggml_tensor * kv_nope = ggml_view_3d(ctx0, kv, n_embd_head_nope, 1, nt, - ggml_row_size(kv->type, n_embd_head), - ggml_row_size(kv->type, n_embd_head), - 0); - ggml_tensor * kv_pe = ggml_view_3d(ctx0, kv, n_embd_head_rope, 1, nt, - ggml_row_size(kv->type, n_embd_head), - ggml_row_size(kv->type, n_embd_head), - ggml_row_size(kv->type, n_embd_head_nope)); - kv_pe = ggml_rope_ext(ctx0, kv_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, + kv = ggml_rope_ext(ctx0, kv, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); - cb(kv_pe, "kv_pe", il); - kv = ggml_concat(ctx0, kv_nope, kv_pe, 0); + kv = ggml_rope_set_offset(kv, n_embd_head_nope); cb(kv, "kv", il); const int64_t ratio = hparams.dsv4_compress_ratios[il]; @@ -1245,17 +1196,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl( } out = ggml_reshape_3d(ctx0, out, n_embd_head, n_head, nt); - ggml_tensor * out_nope = ggml_view_3d(ctx0, out, n_embd_head_nope, n_head, nt, - ggml_row_size(out->type, n_embd_head), - ggml_row_size(out->type, n_embd_head)*n_head, - 0); - ggml_tensor * out_pe = ggml_view_3d(ctx0, out, n_embd_head_rope, n_head, nt, - ggml_row_size(out->type, n_embd_head), - ggml_row_size(out->type, n_embd_head)*n_head, - ggml_row_size(out->type, n_embd_head_nope)); - out_pe = ggml_rope_ext_back(ctx0, out_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, + out = ggml_rope_ext_back(ctx0, out, inp_pos, nullptr, n_embd_head_rope, rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l, ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l); - out = ggml_concat(ctx0, out_nope, out_pe, 0); + out = ggml_rope_set_offset(out, n_embd_head_nope); cb(out, "attn_derope", il); out = ggml_reshape_3d(ctx0, out, o_group_dim, n_groups, nt); diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 5b70a5179..d3c919b35 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -591,17 +591,9 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_ kv = build_norm(kv, layer.attn_kv_norm, nullptr, LLM_NORM_RMS, il); kv = ggml_reshape_3d(ctx0, kv, n_embd_head, 1, n_tokens); - ggml_tensor * kv_nope = ggml_view_3d(ctx0, kv, n_embd_head_nope, 1, n_tokens, - ggml_row_size(kv->type, n_embd_head), - ggml_row_size(kv->type, n_embd_head), - 0); - ggml_tensor * kv_pe = ggml_view_3d(ctx0, kv, n_embd_head_rope, 1, n_tokens, - ggml_row_size(kv->type, n_embd_head), - ggml_row_size(kv->type, n_embd_head), - ggml_row_size(kv->type, n_embd_head_nope)); - kv_pe = ggml_rope_ext(ctx0, kv_pe, inp_pos, nullptr, n_embd_head_rope, rope_type, 0, + kv = ggml_rope_ext(ctx0, kv, inp_pos, nullptr, n_embd_head_rope, rope_type, 0, freq_base, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); - kv = ggml_concat(ctx0, kv_nope, kv_pe, 0); + kv = ggml_rope_set_offset(kv, n_embd_head_nope); cb(kv, "kv_injected", il); if (inp_attn->self_k_rot_swa) { diff --git a/src/models/minicpm3.cpp b/src/models/minicpm3.cpp index e011b1ff0..7820d5224 100644 --- a/src/models/minicpm3.cpp +++ b/src/models/minicpm3.cpp @@ -115,19 +115,9 @@ llama_model_minicpm3::graph::graph(const llama_model & model, const llm_graph_pa q = ggml_mul_mat(ctx0, model.layers[il].wq_b, q); cb(q, "q", il); - // split into {n_head * n_embd_head_qk_nope, n_tokens} - ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, - ggml_row_size(q->type, hparams.n_embd_head_k()), - ggml_row_size(q->type, hparams.n_embd_head_k() * n_head), - 0); - cb(q_nope, "q_nope", il); - - // and {n_head * n_embd_head_qk_rope, n_tokens} - ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, - ggml_row_size(q->type, hparams.n_embd_head_k()), - ggml_row_size(q->type, hparams.n_embd_head_k() * n_head), - ggml_row_size(q->type, n_embd_head_qk_nope)); - cb(q_pe, "q_pe", il); + // {n_embd_head_k, n_head, n_tokens}, RoPE is applied to the trailing dims only + q = ggml_reshape_3d(ctx0, q, hparams.n_embd_head_k(), n_head, n_tokens); + cb(q, "q", il); // {n_embd, kv_lora_rank + n_embd_head_qk_rope} * {n_embd, n_tokens} -> {kv_lora_rank + n_embd_head_qk_rope, n_tokens} ggml_tensor * kv_pe_compresseed = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur); @@ -172,12 +162,13 @@ llama_model_minicpm3::graph::graph(const llama_model & model, const llm_graph_pa v_states = ggml_cont(ctx0, v_states); cb(v_states, "v_states", il); - q_pe = ggml_rope_ext( - ctx0, q_pe, inp_pos, rope_factors, + q = ggml_rope_ext( + ctx0, q, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow ); - cb(q_pe, "q_pe", il); + q = ggml_rope_set_offset(q, n_embd_head_qk_nope); + cb(q, "q_rope", il); // shared RoPE key k_pe = ggml_rope_ext( @@ -187,10 +178,11 @@ llama_model_minicpm3::graph::graph(const llama_model & model, const llm_graph_pa ); cb(k_pe, "k_pe", il); - ggml_tensor * q_states = ggml_concat(ctx0, q_nope, q_pe, 0); + ggml_tensor * q_states = q; cb(q_states, "q_states", il); - ggml_tensor * k_states = ggml_concat(ctx0, k_nope, ggml_repeat(ctx0, k_pe, q_pe), 0); + ggml_tensor * k_states = ggml_concat(ctx0, k_nope, + ggml_repeat_4d(ctx0, k_pe, n_embd_head_qk_rope, n_head, n_tokens, 1), 0); cb(k_states, "k_states", il); cur = build_attn(inp_attn, diff --git a/src/models/plm.cpp b/src/models/plm.cpp index 8ca325f5e..5abefd53b 100644 --- a/src/models/plm.cpp +++ b/src/models/plm.cpp @@ -81,19 +81,9 @@ llama_model_plm::graph::graph(const llama_model & model, const llm_graph_params q = ggml_mul_mat(ctx0, model.layers[il].wq, cur); cb(q, "q", il); - // split into {n_head * n_embd_head_qk_nope, n_tokens} - ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, - ggml_row_size(q->type, hparams.n_embd_head_k()), - ggml_row_size(q->type, hparams.n_embd_head_k() * n_head), - 0); - cb(q_nope, "q_nope", il); - - // and {n_head * n_embd_head_qk_rope, n_tokens} - ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, - ggml_row_size(q->type, hparams.n_embd_head_k()), - ggml_row_size(q->type, hparams.n_embd_head_k() * n_head), - ggml_row_size(q->type, n_embd_head_qk_nope)); - cb(q_pe, "q_pe", il); + // {n_embd_head_k, n_head, n_tokens}, RoPE is applied to the trailing dims only + q = ggml_reshape_3d(ctx0, q, hparams.n_embd_head_k(), n_head, n_tokens); + cb(q, "q", il); // {n_embd, kv_lora_rank + n_embd_head_qk_rope} * {n_embd, n_tokens} -> {kv_lora_rank + n_embd_head_qk_rope, n_tokens} ggml_tensor * kv_pe_compresseed = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur); @@ -143,12 +133,13 @@ llama_model_plm::graph::graph(const llama_model & model, const llm_graph_params 0); cb(v_states, "v_states", il); - q_pe = ggml_rope_ext( - ctx0, q_pe, inp_pos, nullptr, + q = ggml_rope_ext( + ctx0, q, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow ); - cb(q_pe, "q_pe", il); + q = ggml_rope_set_offset(q, n_embd_head_qk_nope); + cb(q, "q_rope", il); // shared RoPE key k_pe = ggml_rope_ext( @@ -158,10 +149,11 @@ llama_model_plm::graph::graph(const llama_model & model, const llm_graph_params ); cb(k_pe, "k_pe", il); - ggml_tensor * q_states = ggml_concat(ctx0, q_nope, q_pe, 0); + ggml_tensor * q_states = q; cb(q_states, "q_states", il); - ggml_tensor * k_states = ggml_concat(ctx0, k_nope, ggml_repeat(ctx0, k_pe, q_pe), 0); + ggml_tensor * k_states = ggml_concat(ctx0, k_nope, + ggml_repeat_4d(ctx0, k_pe, n_embd_head_qk_rope, n_head, n_tokens, 1), 0); cb(k_states, "k_states", il); cur = build_attn(inp_attn,