furyhawk/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-07-21 02:05:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

openpangu: use fused indexer top_k with -fidx (CPU-only op) (#2111)

Browse Source 
Route DSA indexer selection through ggml_indexer_topk when -fidx is set: one op computes the weighted-relu head sum plus causal mask and returns top-k rows without materializing the [n_kv, n_ihead, T] score tensor. Off by default; the unfused chunked/full paths are unchanged. Composes with gathered DSA, deferred attention-chunk masks, and the set_rows mask path.

Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com>
This commit is contained in:
Joel Farthing
2026-07-12 07:09:12 +03:00
committed by GitHub
co-authored by Joel Farthing
parent 3a9d373411
commit bdb23e8763
2 changed files with 31 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/graphs/build_openpangu.cpp
+25 -1
View File
@@ -408,7 +408,31 @@ ggml_tensor * llm_build_context::build_openpangu_attention(
!dsa_gather_allowed &&
openpangu_att_score_should_chunk(n_kv, hparams.param_sink_number, n_head, n_tokens,
OPENPANGU_ATT_SCORE_CHUNK, OPENPANGU_ATT_FULL_KQ_MAX_MIB);
if (chunk_scores) {
if (lctx.cparams.fused_idx_topk) {
// Fused indexer top-k (CPU-only op): one op computes sum_g w * relu(q.k) + causal
// mask -> top-k without materializing the [n_kv, n_ihead, T] score tensor, so no
// score chunking is needed. CUDA backends do not implement GGML_OP_INDEXER_TOPK;
// the scheduler runs it on CPU and copies the operands across the backend boundary.
// The op reads the mask row-strided, so the raw view suffices.
ggml_tensor * idx_mask = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens,
KQ_mask->nb[1], 0);
sel_idx = ggml_indexer_topk(ctx0, k_all_idx, q_idx, w_idx, idx_mask,
GGML_UNARY_OP_RELU, (int) topk); // [topk, T] i32
if (il == 0) ggml_set_name(sel_idx, "opg0_idx_sel");
dsa_gather_engaged = dsa_gather_allowed;
if (dsa_gather_engaged) {
GGML_ASSERT(n_kv >= topk + (int64_t) pad + n_tokens);
} else if (defer_sel_mask_to_att_chunks) {
sel_mask = nullptr;
} else {
ggml_tensor * base = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv, n_tokens);
base = ggml_fill(ctx0, base, -1e30f);
ggml_tensor * zeros = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, topk, n_tokens);
zeros = ggml_fill(ctx0, zeros, 0.0f);
sel_mask = ggml_set_rows(ctx0, base, zeros, sel_idx);
sel_mask = ggml_reshape_2d(ctx0, sel_mask, n_kv, n_tokens);
}
} else if (chunk_scores) {
ggml_tensor * sel_mask_parts = nullptr;
for (int64_t c0 = 0; c0 < n_tokens; c0 += OPENPANGU_IDX_SCORE_CHUNK) {
const int64_t tc = std::min<int64_t>(OPENPANGU_IDX_SCORE_CHUNK, n_tokens - c0);
src/llama.cpp
+6
View File
@@ -7772,6 +7772,12 @@ struct llama_context * llama_init_from_model(
LLAMA_LOG_INFO("%s: dsa_top_k = %d\n", __func__, cparams.dsa_top_k);
}
}
if (model->arch == LLM_ARCH_OPENPANGU) {
LLAMA_LOG_INFO("%s: dsa_idx_topk = %d\n", __func__, cparams.fused_idx_topk);
if (cparams.dsa_top_k > 0) {
LLAMA_LOG_INFO("%s: dsa_top_k = %d\n", __func__, cparams.dsa_top_k);
}
}
LLAMA_LOG_INFO("%s: k_cache_hadam = %d\n", __func__, cparams.k_cache_hadamard);
LLAMA_LOG_INFO("%s: v_cache_hadam = %d\n", __func__, cparams.v_cache_hadamard);
LLAMA_LOG_INFO("%s: split_mode_graph_scheduling = %d\n", __func__, cparams.split_mode_graph_scheduling);