openPangu: fused/specialized ops for prefill and decode throughput - #2159
openPangu: fused/specialized ops for prefill and decode throughput#2159joelfarthing wants to merge 4 commits into
Conversation
The quantized-K path under-sized its q8_1 scratch buffer: it allocated q->ne[1]*max_rows blocks using the unpadded head dim, but quantize_mmq_q8_1_cuda writes q_padded/QK8_1 blocks per row and must process all q->ne[1]*nrows rows. Size the buffer by (q_padded/QK8_1) blocks x (q->ne[1]*max_rows) rows and pass the full q->ne[1]*nrows row count so the scratch cannot be overrun and every query row is quantized. CUDA graph identity: INDEXER_TOPK dispatches a source-type-specialized kernel (dense F16 vs quantized cache, F32 vs F16 mask). Source addresses alone do not identify the captured kernel, so snapshot each source type in ggml_graph_node_properties and force re-capture when an INDEXER_TOPK source type changes, preventing a reused graph from replaying the wrong kernel variant when sources are reallocated at the same address. CPU backend: report INDEXER_TOPK support via iqk_indexer_topk_supported (guarded by GGML_USE_IQK_MULMAT) so the scheduler places the node on a backend that can run it. Harden the scheduler's pass-5 node-assignment check from assert to GGML_ASSERT so a node no backend supports fails as a defined abort under NDEBUG instead of indexing sched->backends[-1].
…ated) Adds the OpenPangu-2.0-Flash fused/specialized op path: latent attention over a packed K/V cache (prefix + indexed variants), fused RMSNorm+add, offset RoPE, batched mix, and pack-cache-rows, wired into the OpenPangu graph builder alongside the shared INDEXER_TOPK op. Each op ships a CUDA kernel plus the exact legacy ggml chain it replaces. Selection is capability-gated: an op is used only when the scheduled backend reports support for the exact node (ggml_backend_supports_op, anchored on per-layer weights), and falls back to the legacy chain otherwise, so a partially-supporting backend never strands a node. GGML_OPENPANGU_LEGACY_OPS is a single compile-time bring-up switch (build with -DGGML_OPENPANGU_LEGACY_OPS=1) that forces every legacy chain at once for A/B baselines; it dominates an explicit -fidx request, and the effective policy is reported in the model load log (openpangu_ops). The fallback also covers operand contract, not just backend capability: each fused candidate is prechecked with a non-aborting validity predicate (ggml_latent_attn_ext_valid, ggml_rope_ext_offset_valid, ggml_pack_cache_rows_valid, ggml_batched_mix_ext_valid) before construction, so a layout the op cannot represent takes the legacy chain instead of aborting in the constructor. The CPU backend reports no support for LATENT_ATTN, so its scalar reference implementation stays a correctness/testing path and a CPU-only or -ngl 0 run keeps the vectorized legacy attention rather than silently adopting the reference op. Builds on the INDEXER_TOPK quantized-CUDA scratch-sizing fix in the parent commit; the OpenPangu graph selects that op on backends that support it.
| const int ne00, const int ne01, const int ne02, const int nb00, const int nb01, const int nb02, | ||
| const int nb03, const int ne10, const int ne11, const int ne12, const int nb10, const int nb11, | ||
| const int nb12, const int nb13, char ** cdst_indirect, int graph_cpynode_index) { | ||
| const int nb12, const int nb13, char ** cdst_indirect, int graph_write_index) { |
There was a problem hiding this comment.
Even if the variable name does not correspond to your taste, I see zero reason to change 119 LOC for no real reason.
| bool use_cpy_indirection = false; | ||
| std::vector<char *> cpy_dest_ptrs; | ||
| char ** dest_ptrs_d; | ||
| // CPY and destination-rooted mutation ops share one graph-order pointer table. |
There was a problem hiding this comment.
Please no renaming. What is it that we gained by changing the name other than to increase the review burden?
Revert the graph_cpynode_index / use_cpy_indirection / cpy_dest_ptrs field names and the ggml_cuda_cpy_dest_ptrs_copy helper to their original spellings. PACK_CACHE_ROWS reuses this graph-order pointer table, so the machinery is no longer CPY-only, but renaming it is not required for the sharing and only inflates the diff. Keep the original names and note the shared use in a comment. No functional change: the destructor cudaFree(dest_ptrs_d) leak fix, the dest_ptrs_d nullptr init, the ggml_graph_node_properties::src_type field, the PACK_CACHE_ROWS branch, and the op-conditional destination push_back are all retained.
|
Thanks for review! Reverted to the original names. PACK_CACHE_ROWS reuses the same table, so |
| local_max = fmaxf(local_max, v); | ||
| } | ||
| shbuf[tid] = local_max; | ||
| __syncthreads(); |
There was a problem hiding this comment.
We have things such as warp_reduce_sum, warp_reduce_max, etc.
| return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && | ||
| sink_s >= 1 && sink_s <= 8 && op->src[0]->ne[0] == (int64_t) sink_s*sink_s; | ||
| } | ||
| case GGML_OP_LATENT_ATTN: { |
There was a problem hiding this comment.
If you are writing such extensive checks then please add a function that does that. This applies to all the long checks that you have added.
| #include "../llama-build-context.h" | ||
| #include "../llama-model.h" | ||
| #include "../llama-context.h" | ||
| #include "openpangu-op-policy.h" |
There was a problem hiding this comment.
I was disappointed to find out that after adding multiple new ops and thousands of lines of implementations, the build graph function actually became longer rather than shorter, as one would expect given the many fused ops.
|
Tested the This PR
Main branch
|
|
As a general comment: I feel too much new stuff is being added. I need to think some more. |
| ggml_tensor * q_idx_offset = nullptr; | ||
| if (openpangu_rope_offset_enabled() && | ||
| ggml_rope_ext_offset_valid(q_idx, inp_pos, nullptr, n_rot, 0, rope_type)) { | ||
| q_idx_offset = ggml_rope_ext_offset(ctx0, q_idx, inp_pos, nullptr, n_rot, 0, rope_type, |
There was a problem hiding this comment.
This is funny. We add a whole new op so we can skip the so called "legacy route". We then use a rope offset of zero. Which is 100% equivalent to just using the exiting ggml_rope_ext directly.
If in doubt, just change ggml_rope_ext_offset to ggml_rope_ext, remove the extra 0 argument between n_rot and rope_type, and see what happens. When you find that it works like that, remove all the checking if we can use that, remove the new GGML_OP_ROPE_OFFSET op and associated ggml_rope_ext_offset API, etc.
The reason you find this pattern of making two views of a tensor, doing RoPE one one of the views, and then concatenating them back together is I guess laziness to check what gets copied unmodified for the different RoPE types. But in this particular case, it just worksTM.
On a second thought: don't remove the op. I see that the DS4 implementation is full of the "legacy" pattern.
There was a problem hiding this comment.
On a second thought: don't remove the op. I see that the DS4 implementation is full of the "legacy" pattern.
Right, the zero-offset validates here but was designed for nonzero-offset to replace DS4's repeated split/RoPE/concat chains. I tried to make all of these ops portable in that way.
As a general comment: I feel too much new stuff is being added. I need to think some more.
I totally understand that, and would also understand keeping this as an unmerged reference branch if that's better.
Also: I made an IQ4_NL of openPangu that reduces overhead a bit and seems to outperform the Q4_K_M on my usual agent tests. It's at https://huggingface.co/ji-farthing/openPangu-2.0-Flash-ik-llama-GGUF in case helpful.
There was a problem hiding this comment.
So, I did a quick experiment with DS4 where avoiding the concatenation between RoPE'd and not RoPE'd parts could matter. I see zero difference on the CPU. On the GPU, I see ~1% performance gain for PP and no performance gain for TG. So, basically, I don't think it is worth the extra code and extra op.
But if you think it is worth it, please extract the RoPE change into a separate PR. Please no new op, and please don't add I don't know how many checks. Strictly speaking the addition is a boolean flag, and not an offset. With an offset there is the possibility to have a [no-RoPE, RoPE, noRoPE] configuration, which I don't think is useful in practice. We only care about [RoPE, noRoPE] (handled by the current implementation) or [no RoPE, RoPE] (triggered by an added boolean flag that can go into op_params[15]. So, then, one can simply set the flag in the graph building code, but if we wanted to be more professional, there would be a function in the API such as ggml_rope_set_flipped that sets op_params[15] to 1(after checking that n_dims is actually less than src[0]->ne[0].
|
OK, below is a Given this, lets do the following: lets split out the Looking at the arguments of Once that happens, and hopefully I have also merged the DS4 branch by then, we can go one by one with the remaining ops, where it should be demonstrated that they are useful for openPangu and DS4 (and possibly other arches if applicable).
|
|
Oh, one more thing. As it stands, the PR does basically nothing for CPU-only inference. Here is what I get on my Ryzen-3995WX CPU with the PR (with all changes included):
And this is on the main branch
|
|
This PR superseded by #2168 et seq. |
2 participants
[superseded by #2168 et seq.]
This is the second of two stacked PRs; it depends on #2158, the shared
indexer_topkfix (q8_1 scratch sizing, a CPU support predicate, a CUDA graph-node identity fix), split out so the one change reaching DeepSeek2 and GLM-DSA is reviewable on its own diff. This one is the openPangu operator suite that builds on it. I'm aware it is large, and I don't necessarily expect a quick review or a clean merge. I wrote the operators to be usable beyond openPangu, with GLM-DSA and DeepSeek v4 in mind, so I hope some are useful even if openPangu adoption isn't.Currently, openPangu-2.0-Flash runs its DSA attention, indexer, and MTP heads on a bespoke manual graph. It works, but leaves throughput and compute-buffer on the table at five sites: the latent-attention downstream reconstruction, the MTP post-norm RMSNorm+residual, the offset RoPE on the DSA indexer K/Q, the DSA cache-row writes, and (without
-fidx) the top-k selection. This PR adds a fused or specialized op for each and routes openPangu through them, capability-gated: a layer adopts the fused path only when its scheduled backend can execute the candidate node, and each candidate is prechecked with a non-aborting validity predicate, so an unsupported backend or a layout the op cannot represent falls back to the byte-identical legacy chain instead of a scheduler-inserted CPU island or a constructor abort. The CPU backend reports no support for the latent-attention op, so a CPU-only or-ngl 0run keeps the vectorized legacy attention.Selection is not driven by per-op tunables: the default is fixed and capability-gated, and the only override is a compile-time
-DGGML_OPENPANGU_LEGACY_OPS=1, which forces every legacy chain at once (all six ops revert together) and dominates an explicit-fidx. It's not read from the environment, so a compiled binary carries one fixed policy. The effective policy prints in the load log (openpangu_ops = fused (capability-gated)orlegacy-forced, plus adsa_idx_topk_policyline).The operators, grouped by what they do:
ggml_latent_attn_prefix_extandggml_latent_attn_indexed_extreplace the gathered/scored latent-attention reconstruction on the openPangu attention path (dense, SWA, and the gathered DSA route). This is the largest single piece and the biggest measured win. The op takes an F32 query and reads the raw latent K cache directly in F32, F16, or Q8_0, dequantizing internally; the optional mask is F32. When a layer's backend cannot run the produced node, the layer keeps the full legacy chain, including the get_rows subchunk loops.ggml_fused_rms_norm_addfolds the post-norm RMSNorm and its residual add into one op at the two openPangu MTP post-norm sites. The fusedGGML_OP_FUSED_RMS_NORM_ADDnode is produced only for the canonical case (contiguous F32 weight matching the normalized dimension); an F16 or otherwise non-canonical weight is decomposed by the constructor into rms_norm + mul + add, and a backend that cannot run the fused node likewise takes the decomposed path. This op is only in the MTP graph, so it is dormant in ordinary decode; see the validation caveats.ggml_rope_ext_offsetapplies NEOX RoPE over a channel interval with a same-pass passthrough of the untouched channels, used at the two DSA indexer K/Q sites. Inputs are F32. The main-Q pack is deliberately not converted, because its no-rope branch is a projection rather than a split/reconcat. Both openPangu call sites pass a zero interval offset, which is byte-for-byte the existing NEOX kernel; the nonzero-offset generality is carried for the rope-first packing a future GLM/DeepSeek adopter needs, and the shared CUDA and CPU NEOX paths are verified equivalent at offset zero so no existing NEOX model changes.ggml_pack_cache_rowspacks two F32 sources into one destination cache row per token, with F32/F16/Q8_0 destinations and a block-aligned split for the quantized path, replacing the ggml_cpy-based DSA cache writes. It carries a small openPangu cache-write layer that validates the base and MTP write sets and retargets the write nodes for graph reuse.GGML_OP_INDEXER_TOPKautomatically when the scheduled backend can execute it, and keeps the legacy chain when the indexer projection placement is split or unresolved. One honest exception to the capability gate: an explicit-fidxforces this op on regardless of backend support, preserving the existing DSA opt-in semantics, so-fidxon a backend that cannot run the op is the one way to get a scheduler-inserted CPU island here. Every other op in this PR is gate-only, with no force override. A-DGGML_OPENPANGU_LEGACY_OPS=1build forces the legacy chain and takes precedence even over an explicit-fidx. The correctness work on the op's CUDA path is not in this PR: it is the prerequisiteindexer_topkfix this one depends on (see Scope and blast radius).ggml_batched_mix_extfuses the mHC post-block combine, the per-token contractionR_new[h,s,t] = sum_j m[s,j,t] * R[h,j,t]over the Sinkhorn doubly-stochastic mix, into one strided token-batched matrix op (r [D,J,T]bymix [J,O,T]to[D,O,T]). It produces all S output streams in a single pass. The existingggml_mul_multi_addfuses only the single-output reduce (sum_j r[j]*w[j]), so an S-output mix would take S passes through it, which is exactly the per-slot legacy loop this op replaces. F32, capability-gated like the rest, one openPangu call site.Scope and blast radius. The new operators are constructed only from
build_openpangu.cpp(and, for pack-cache, the openPangu cache-write helper). No other model builder calls any of them; I searched every builder to confirm it. The new enum entries, name/symbol tables, CPU compute functions, and CUDA dispatch/support cases in the shared files (ggml.h,ggml.c,ggml-cuda.cu) are dormant unless the op is actually built, so they are inert for every other architecture.Two touches genuinely reach past openPangu:
src_typegraph identity) is the prerequisite indexer_topk: fix quantized q8_1 scratch sizing on CUDA #2158, carried on its own diff. Here openPangu just adopts the op under the capability gate;build_deepseek2.cppand the op's numerical contract are untouched. This PR does extend indexer_topk: fix quantized q8_1 scratch sizing on CUDA #2158'ssrc_typegraph-identity snapshot to the new type-specialized ops (LATENT_ATTN, PACK_CACHE_ROWS, BATCHED_MIX), and the op-params comparison to LATENT_ATTN and BATCHED_MIX, so their captured kernels cannot replay stale; no existing op's graph reuse is loosened.cdst_indirect/graph_write_index). It is gated:use_write_indirectiondefaults false, and a null indirect pointer falls through to the original direct destination, so generic CPY is behavior-identical for every non-openPangu model — a signature change, not a behavior change. This PR also frees the shareddest_ptrs_dtable in theggml_cuda_graphdestructor, which base leaks on every graph that uses indirection (the realloc path already frees it, so freeing at destruction is safe); the fix closes the leak for every CPY-indirection user.No other architecture's graph, cache layout, or state handling changes. I traced the shared files to confirm the additions are registration and dispatch, not edits to existing op behavior.
For GLM-DSA and DeepSeek v4 use: these are generic ggml ops, not openPangu-hardcoded.
ggml_batched_mix_ext,ggml_rope_ext_offset, andggml_pack_cache_rowsare F32 with no openPangu assumptions;ggml_fused_rms_norm_addtakes F32/F16 weights; the latent-attention ops carry the MLA split (dv,dv_off) and the indexed route as runtime arguments, so any latent-KV-plus-indexer architecture supplies its own dimensions behind the same gate. The ggml core carries no openPangu constants; the openPangu specifics stay inbuild_openpangu.cpp.I've watched #2147 (the DeepSeek v4 work) and checked this suite against it directly. It does not reimplement any of these operators: its DeepSeek v4 attention is assembled from
ggml_flash_attn_ext, aggml_set_rowstop-k mask, and the sharedggml_indexer_topk, and none ofggml_latent_attn_*,ggml_rope_ext_offset,ggml_pack_cache_rows,ggml_batched_mix_ext, orggml_fused_rms_norm_addappear in its diff. Where the two touch the same ground it is confirming rather than conflicting: DeepSeek v4 calls the sharedGGML_OP_INDEXER_TOPKfrom its own builder, joiningbuild_deepseek2.cppand openPangu as a further in-tree caller of the op this stack fixes. The two diffs intersect in eight files, but the overlap is almost entirely additive registration and dispatch: op enums and names inggml.h/ggml.c, the CUDA dispatch inggml-cuda.cu, the legacy-ops CMake option, and config plumbing inllama-context.h/llama-cparams.h/llama.cpp. The one substantive shared edit isbuild_openpangu.cpp, where #2147 swaps amulplussum_rows_extfor the existingggml_mul_multi_addat the mHC pre-combine, a few lines from this PR's post-combine work; it resolves on rebase. Myggml_latent_attn_indexed_extcomputes in one op the same indexed latent attention that DeepSeek v4 currently assembles from three, so if that consolidation is useful there, the adoption pattern is the openPangu one, gated the same way.Cross-architecture usefulness is a design-intent and shared-op claim, not validated. The exception is INDEXER_TOPK, which is already shared and now has a further caller in #2147.
Contracts. The inference-only constructors reject gradient-bearing inputs; the latent-attention op requires a float-aligned mask row stride (CUDA matched byte-exact to the CPU reference); quantized latent caches and the pack-cache quantized split must be block-aligned. The constructors assert these.
Validation. Re-established on the converged branch (head
2879e4466). On an RTX 4070, CUDA 13.3, GCC 13.3, openPangu-2.0-Flash-Q4_K_M with-ngl 999 -ot exps=CPU -c 32768 -ictk q8_0 -fidx -dsatk 64greedy: the fused build and a-DGGML_OPENPANGU_LEGACY_OPS=1legacy build both generate coherent output, and the load log confirms the effective policy on each:openpangu_ops = fused (capability-gated)withdsa_idx_topk_policy = explicit-force, versuslegacy-forced (-DGGML_OPENPANGU_LEGACY_OPS=1)withdsa_idx_topk_policy = legacy-forcedeven with-fidxpassed, so the compile-time switch dominates an explicit request.Quality on the fused build, four fixtures generated and then run rather than eyeballed:
bulgaria(1300 years, coherent, no repetition, no CJK bleed),sayap(correct alignment, no preamble),quicksort(sorts correctly over 200 random cases),double-linked-listwith its pytest (14 of 14 pass). All four finished onstop, no reasoning-token leak, decode stable at 17.07 to 17.17 tok/s.Lossless fused-versus-legacy: at the converged head, in the deployed regime (
-c 34816 -ctk q8_0 -ictk q8_0 -fidx), the fused build and a compile-time-DGGML_OPENPANGU_LEGACY_OPS=1build both generate coherent output; on the earlier pre-split build the fused ops' greedy streams were byte-identical to the legacy chains. The restoration of 32K retrieval and near-lossless q8_0 perplexity relative to current main comes from the prerequisite indexer PR #2158; this PR's job is to add the fused ops without regressing that.The per-op focused CPU/CUDA correctness tests pass on the operators; a few assertions in the held-out suite pinned pre-split behavior and are realigned to this PR's contract. The tests are held out of this diff to keep it to the production change, and I can supply them on request.
Performance. Warm A/B/A/B, medians of two passes each, measured on this branch's head
2879e4466against current main9d07d8681(bothllama-sweep-benchbinaries' sha256 recorded with the logs). openPangu-2.0-Flash-Q4_K_M on an RTX 4070, experts on CPU:-ngl 999 -ot exps=CPU -fa 0 -ctk q8_0 -ictk q8_0 -fidx -dsatk 64, sweep-c 34816 -b 2048 -ub 2048. This is the deployed long-context regime (q8_0 K-cache and q8_0 indexer cache).Exact medians (
llama-sweep-bench):Prefill is a steady ~1.11× across the curve; decode holds while main sags (1.14× at 0 to 1.25× at 32K; 93.4% vs 84.6% retention), with no collapse at the ceiling. The throughput is the fused ops; the prerequisite indexer fix is a correctness change, expected to be throughput-neutral (not separately swept). The 32K CUDA compute buffer is 2.29× smaller (2873.27 to 1256.02 MiB, 14361 to 8473 graph nodes); KV (1009.71 MiB) and host (312.28 MiB) buffers are identical, so the reduction is the fused graph, not a cache change.
MTP is still not a win.
ggml_fused_rms_norm_addlives only at the MTP post-norm sites, so it's not exercised in ordinary decode, and MTP self-speculation loses in this CPU-expert-offload regime (heads=1 runs 10.7 to 12.1 tok/s versus 17.1 non-MTP, and it degrades further with draft depth and with heads=3), because the batched verify pays the CPU expert cost every step. Measured in MTP with the op on versus off, its throughput contribution is within noise and the greedy token streams match. It earns its place on correctness and on removing a graph node, not on a decode-speed number in this regime. I suspect that MTP needs the whole model on GPU for a win.