From bdb23e8763004467a2ea1bb00fa35de96580f329 Mon Sep 17 00:00:00 2001 From: Joel Farthing Date: Sat, 11 Jul 2026 23:09:12 -0500 Subject: [PATCH] openpangu: use fused indexer top_k with -fidx (CPU-only op) (#2111) 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> --- src/graphs/build_openpangu.cpp | 26 +++++++++++++++++++++++++- src/llama.cpp | 6 ++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/graphs/build_openpangu.cpp b/src/graphs/build_openpangu.cpp index b0c6faac4..4dfb1bdd2 100644 --- a/src/graphs/build_openpangu.cpp +++ b/src/graphs/build_openpangu.cpp @@ -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(OPENPANGU_IDX_SCORE_CHUNK, n_tokens - c0); diff --git a/src/llama.cpp b/src/llama.cpp index c06b7194a..56c6b2bad 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -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);