Skip to content

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

Merged
ikawrakow merged 1 commit into
ikawrakow:mainfrom
joelfarthing:filament/openpangu-fidx-cpu
Jul 12, 2026
Merged

openpangu: use fused indexer top_k with -fidx (CPU-only op)#2111
ikawrakow merged 1 commit into
ikawrakow:mainfrom
joelfarthing:filament/openpangu-fidx-cpu

Conversation

@joelfarthing

Copy link
Copy Markdown
Contributor

openPangu's indexer scoring is exactly the shape #2098's GGML_OP_INDEXER_TOPK computes, so this PR routes openPangu's DSA selection through the op when -fidx is set. One branch in build_openpangu.cpp, off by default, unfused paths unchanged. It composes with the gathered DSA decode path, the deferred attention-chunk masks, and the set_rows mask path, and the loader now echoes dsa_idx_topk for openPangu the way it does for GLM-DSA. The op is CPU-only, so on a CUDA rig the scheduler places it on CPU and its operands cross the backend boundary each eval.

The point on a hybrid rig is the compute-buffer ceiling. Measured on a 12 GB RTX 4070 (the #2065 Q4_K_M, -fa off -ngl 999 -ot exps=CPU -ub 2048, q8_0 caches):

fidx-cpu-pr-graph
ctx -fidx load CUDA0 compute buffer CUDA_Host compute buffer peak VRAM
32768 off FIT 2808.27 MiB 350.16 MiB 6687 MiB
32768 on FIT 2808.27 MiB 320.72 MiB 6685 MiB
65536 off FIT 3945.52 MiB 663.80 MiB 8771 MiB
65536 on FIT 2808.27 MiB 580.97 MiB 7633 MiB
131072 off OOM allocation fails - -
131072 on FIT 2808.27 MiB 1101.47 MiB 9527 MiB

The paths diverge at 64K, and at 128K the unfused path fails to allocate while -fidx holds the CUDA0 buffer at its 32K floor, because the score tensors never exist on the GPU. The cost is throughput: with the op on CPU and attention on CUDA, the warm A/B sweep (figure, right panels) runs S_PP -8% at N_KV 2048 growing to -44% at 32K, and S_TG -4% to -26%. A -ngl 0 probe on this 8-core desktop showed S_TG at parity and S_PP 3 to 12 percent behind, so on this class of CPU the flag is not a throughput feature, and can be left off when the unfused path fits.

Correctness is output-level: a 4.7K-token needle at temp 0 (past top_k 2048, so selection prunes) returns identically with -fidx off and on across the base graph, MTP speculative decode, the deferred-chunk regime, and all four accepted indexer cache types (12/12 arms; top-k tie-break identity is not asserted).

The branch leaves op placement to the scheduler, so if a CUDA implementation of the op lands (the direction #2103 is exploring), openPangu picks it up here with no further arch changes, and the round-trip cost above disappears. And the op carries #2098's platform contract (AVX2 and ARM dotprod builds); other builds should leave -fidx off, same as GLM-DSA today.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants