From b90939934add9ba4fbb37e8c6470809a70b78f0a Mon Sep 17 00:00:00 2001 From: Joel Farthing Date: Sat, 11 Jul 2026 04:29:20 -0500 Subject: [PATCH] model: add openPangu-2.0-Flash (92B-A6B) with MLA-latent cache, DSA/SWA, mHC, and multi-head MTP (#2065) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * openpangu: Stage-1 converter probe for openPangu-2.0-Flash Add OpenPanguV2ForCausalLM conversion support (converter-only; runtime graph is Stage-2). Registers a new LLM_ARCH_OPENPANGU on the Python/gguf-py side: - gguf-py/constants.py: MODEL_ARCH.OPENPANGU + name, indexer KV keys, 22 new tensor enums (DSA indexer x4, MoME convs x3, param-sink x2, mHC/Hyper- Connections x12, block-post-norm), and the full MODEL_TENSORS list reusing the deepseek MLA + MoE + NextN bricks. - tensor_mapping.py: arch-specific block mappings that disambiguate the sandwich norms (post_attention/pre_mlp/post_mlp) and pin every Pangu-only tensor; non-block global mHC merge module. - convert_hf_to_gguf.py: OpenPanguV2Model (subclasses DeepseekV2Model) with set_gguf_parameters (MLA/MoE/indexer/mHC/param-sink/DSA+SWA metadata), modify_tensors (expert merge, kv_b split, no MTP skip), and the OpenPanguV2Tokenizer pre-tokenizer hash. Validated offline against the real 50-shard safetensors index: all 37,587 tensors map to a GGUF target (0 unmapped), and set_gguf_parameters reads only hparams present in config.json. No weights downloaded; no GPU. Pinned on the ik/dsa_loop_hadamard_blend DSA substrate. * openpangu: Stage-2 arch scaffold (LLM_ARCH_OPENPANGU) — loadable, compiles New arch on main (DSA-decoupled). Declares openPangu-2.0-Flash to the runtime so the model loads into memory; the compute graph is the next step. - llama-arch.{h,cpp}: LLM_ARCH_OPENPANGU + name; 3 KV keys (mhc_num_stream, mhc_recur_norm, param_sink_number); 18 tensor enums (mHC x12, MoME conv x3, param-sink x2, block-post-norm). - llama-model.cpp: OPENPANGU tensor-name block, strings matched to the converter. - llama-model.h: layer + model struct fields (mHC / conv / sink / block-post / merge). - llama-hparams.{h,cpp}: reader (MLA + MoE + sigmoid gate + indexer + mHC + param-sink + NextN); n_layer_kv_from_start = n_layer - nextn (MTP skipped). - llama-load-tensors.cpp: create_openpangu_tensors (GLM-DSA MLA/MoE base + Pangu tensors; indexer loaded-but-unused for dense fallback); dispatch + is_mla_attn. Builds clean (CPU-only libllama). Dense-fallback design: no DSA indexer / SWA windowing / MTP for first generation (exact <=512 tokens). Graph is Stage-2b. * openpangu: fix compresskv_conv dim (kv_lora_rank, not +rope); pin attention order in spec * openpangu: end-to-end runtime — build_openpangu graph runs, generates (garbled) First full forward pass of openPangu-2.0-Flash on ik_llama. Pipeline works end to end: new LLM_ARCH_OPENPANGU loads the Q4 GGUF, the graph executes, and llama-cli generates 40 tokens (EXIT=0). Output is currently garbled (tensor-layout bug to debug), but the structure is proven. graphs/build_openpangu.cpp: dense decompressed-MHA attention + 4-stream mHC (Hyper-Connections) with 20-iter Sinkhorn + MoE(sigmoid+shared) + sandwich norms + entry stream-repeat/tail-merge + inp_out_ids selection. Bring-up fixes to load+run: - llama-vocab.cpp: register 'openpangu' pre-tokenizer (QWEN2 family) - llama.cpp: OPENPANGU -> LLAMA_ROPE_TYPE_NORM (was defaulting to NONE=-1) - llama-load-tensors.cpp: full wkv_b load; k_b/v_b as flattened 2D; block_post_norm dim = S*H (10240); conv weights 2D {3,C}; mHC alpha/beta/gamma + param_sink + merge params use bare (no-.weight) tensor names - llama-model.cpp: OPENPANGU is NOT is_mla_attn (decompressed MHA, standard KV cache) - graph loop bounded to base layers (skip NextN/MTP) v0 deferrals (need conv-state cache / manual attention path, all documented): MoME convs (passthrough), o_conv, param_sink. Next: fix the layout bug to coherence. * openpangu: COHERENT generation — NEOX rope, Sinkhorn orientation, MoME convs, param_sink Four correctness fixes on top of the end-to-end scaffold, verified checkpoint-by- checkpoint against a Python golden reference running on the GGUF's own dequantized weights (block-0 activations now match to rounding at full fidelity): - rope: NORM -> NEOX. Pangu config rope_interleave=false; the Infer source maps it as is_neox_style = not rope_interleave (rotary_mode='half'). - mHC Sinkhorn: the flat h_res block is torch-[r,c] row-major, so a bare ggml reshape lands column-fastest; the doubly-stochastic iteration ran transposed (Sinkhorn is not transpose-symmetric). One transpose at input fixes the whole chain including the mhc_post application. - MoME convs (qa/compresskv/o): were passthrough stubs. Implemented as out = x + causal_conv1d(x) (every Infer call site uses residual_connection=1; tap stats confirm the perturbation form). Taps cast f16->f32 for ggml_mul. Batch-local v0: exact for fresh-sequence prefill; decode steps miss the t-1/t-2 taps until a conv-state cache exists. - param_sink: 128 learned latent-KV entries prepended per layer via a manual attention path (kv_store + explicit soft_max over [sinks ++ cache]); huge effect at short context. o_conv now applied pre-o_proj on the same path. flash_attn forced off for OPENPANGU (FA kernel cannot see the sinks). - converter: add_bos_token=true (HF prepends <|pangu_text_start|> via the post-processor; the key was absent so ik dropped BOS). Greedy Q4_K_M smoke, chat template + : coherent CoT reasoning and a correct answer. Layer-0 instrumentation (opg0_* names) kept for now. * openpangu: MoME conv-state cache — decode steps get real t-1/t-2 taps Allocate a per-layer cache_s_l tensor for OPENPANGU base layers holding the last two pre-conv latents of the three MoME sites, packed [qa 2*1024 | compresskv 2*512 | o 2*6144] f32 (~60KB/layer). The conv helper reads the [C,2] history window (zeros at sequence start, kv_head==0), builds xx = [hist ++ x], and writes the last two columns back each ubatch — the concat naturally handles both prefill chaining and the T==1 shift. Read precedes write in graph order; the fixed-offset copy is graph-reuse safe. Verified: prefill anchors unchanged (bit-path identical, zero-history branch); -ub 1 token-by-token run matches the golden reference at t4 (qlora_conv 0.084, R_block 0.008 rel; attn_out 0.15 on one channel = f16 KV-cache rounding, washes out by post-norm); final logits differ from full-batch only by a common-mode shift that softmax cancels. Chat-template greedy smoke: think-block repetition is gone — clean structured CoT and correct answer. v0 limits documented in the helper: one state slot (single sequence); cache rewinds leave the state stale. * openpangu: NextN/MTP speculative decoding — 1.7-1.8x TG on CPU Wire the three NextN layers (46-48) into ik's MTP speculative framework (--spec-type mtp). v0 drafts with head 1 (layer 46), self-chained by the framework. - llama.cpp: add OPENPANGU to the cparams.mtp arch allowlist (it was silently zeroed, which left the target context without a logits buffer once the server enabled embeddings -> GGML_ASSERT(lctx.logits) in speculative_is_compat). - load-tensors: MTP layers carry no mHC tensors (tail_use_mhc=false in the reference) — create them only for base layers. nextn.* tensors were already wired by the Stage-1 probe. - build_openpangu: extract the attention sublayer into build_openpangu_attention (shared base/MTP); add build_openpangu_mtp: eh_proj(cat(enorm(embed), hnorm(prev_hidden))) -> one plain-residual Pangu block (sandwich norms, convs+param_sink, MoE+shexp, no mHC/block_post_norm) -> shared_head norm+head. MTP branch returns the draft graph when mtp_op_type != NONE; main graph keeps all-token outputs under cparams.mtp. MTP convs run batch-local (no conv-state slot) — affects acceptance only. A/B (Q4_K_M, CPU, greedy, 192-token chat CoT completion, warm back-to-back, medians of 3, bracketed B/A/B): no-spec: 2.44 t/s (2.34-2.86) --spec-type mtp:n_max=3: 4.23 / 4.49 t/s (brackets) => ~1.7-1.8x Draft acceptance 34% on CoT prose (46% on repetitive text); spec and no-spec greedy outputs are byte-identical. Headroom: conv-state for MTP drafts, n_max tuning, true 3-head chaining (spec_step_idx). * server: include draft_n/draft_n_accepted in /completion timings get_formated_timings() (the /completion path) omitted the speculative counters that get_timings() (the OAI path) already reports; add them, guarded by n_draft_total > 0 like the OAI path. * openpangu: position-indexed MoME conv-state ring — rollback-safe spec decoding + MTP draft chaining The v0 single-slot conv state held the last-2 pre-conv latents of the most recent batch, so any speculative draft rejection left latents of REJECTED positions in the state and every later decode ran with wrong t-1/t-2 taps (3 conv sites x 46 layers). At 192-token greedy runs every spec config diverged from no-spec, each differently (rejection-pattern dependent). Replace it with a per-layer ring cache_s_l [n_lora_q+n_lora_kv+n_head*v_dim, 16]: column pos%16 holds position pos's pre-conv latents ([qa|ckv|o] packed). Invariant: reads target only positions before the first batch token, which are committed, and committed latents depend only on the committed prefix - rollback-safe by construction, no checkpointing. Writes cover the last min(T,16) batch positions in <=2 contiguous cpy segments; the copy sources are views of the [hist ++ x] concat so the history read is an ancestor of every write (read-before-write by graph dependency). The ring is also allocated for the NextN/MTP layers, so the draft head chains real conv taps across WARMUP -> sequential DRAFT_GEN steps (was batch-local zero-history per draft token). graph_reuse is forced off for the arch: ring view offsets are position- baked and the reuse patcher only updates the standard K/V-store copies. Measured cost on the CPU server path: none visible. ggml_set_rows driven by an input index tensor is the future reuse-safe shape. Verified (Q4_K_M, CPU, greedy 192-tok chat-CoT, warm single process): - no-spec output byte-identical to pre-ring build - spec output byte-identical to no-spec below the n_predict cap, for all of n_max in {1,2,3,4,6} x p_min in {0,0.3,0.6} (old build: all diverged) - acceptance n3-p0: 33.9% -> 60.9%; n3-p0.3: 58.1% -> 68.9% - TG medians: no-spec 3.19-3.32 t/s; mtp:n_max=3,p_min=0.3 6.97 t/s (~2.1x) * openpangu: DSA lightning indexer + SWA schedule — long-context correctness past the dense fallback The dense fallback was exact only <=512 tokens (SWA window). This wires the real DSA/SWA hybrid schedule, self-contained from GGUF keys the converter already writes (openpangu.swa_layers + sliding_window_list; absent keys keep the old dense fallback): - SWA layers (30 base @512): the generic inp_KQ_mask_swa path, per-layer mask choice in the builder. The NextN/MTP layers are SWA @2048 in the checkpoint schedule; MTP graphs run in their own context, so the mask fill picks hparams.n_swa_mtp when built with an MTP op type. - DSA layers (16, every 3rd): lightning indexer implemented in-graph from the Infer reference semantics (jointfix _pangu_torch_calib): q_idx = wq_b on the post-conv post-norm q-lora latent (24x128), k_idx = rms-normed wk(x) shared across heads, both NEOX-roped on the FIRST n_rot channels; score = sum_g w_g * relu(q_g . k) in f32, causal-masked, exact top-k via argsort + ggml_set_rows scatter into a -1e30 base -> additive selection mask on the existing manual soft_max seam. Selection engages only when the causal window exceeds index_top_k (2048); below that the layer is exactly dense. - Indexer keys cached per position (cache_idx_l, f32 [128, kv_size], DSA layers only) with the same committed-position invariant as the conv-state ring, so speculative rollbacks stay safe. - param sinks remain outside both the window and the selection budget, matching the reference. Verified (Q4_K_M, CPU): - <=512 tokens: byte-identical to the dense build (96/160-token greedy) - indexer scores vs a GGUF-dequant golden reference at 2101 tokens: 1e-3 rel (f16 weight rounding); top-3 selection indices exact on all compared queries - >512 coherence clean; 3.4K-token needle retrieval through active selection (needle outside every SWA window, ~1300 positions pruned) answers exactly * openpangu: MLA-latent KV cache — attention absorbed into the 512-latent, 14x smaller cache, ~2.2x TG Store per position only [ckv_norm 512 | roped k_pe 64] (f32, k_l) plus the transposed 512-latent (f32, v_l, v_trans layout); per-head K/V are never materialized. q_nope is absorbed through attn_k_b (loaded 2D from the converter split for base layers; derived at load via llm_prepare_mla for the NextN layers - now guarded for layers without attention weights, e.g. the idle NextN heads 2/3). The value side is the latent itself, up-projected through attn_v_b after the weighted sum, matching the Infer _forward_dsa reference. param sinks are native latent-space entries, which removes the per-step full-cache concat+cast that dominated long-context decode. llama_state row sizes now come from llama_kv_k_row_embd/llama_kv_v_row_embd (arch-aware), fixing an out-of-bounds crash in the server prompt-cache save path (hparams-derived 9216-wide rows vs actual 576-wide latent rows). Verified (Q4_K_M, CPU): layer-0 attention output matches an f32 golden reference computed from the same GGUF weight encodings (~1e-2 on O(1) values); MTP spec output byte-identical to no-spec; 3.4K needle retrieval through active DSA selection exact under greedy. Output differs from the materialized build at the token level because attn_k_b/attn_v_b are independently quantized tensors - both are legitimate Q4-fidelity encodings. Perf (CPU, warm): no-spec TG 3.2-3.3 -> 6.9-7.1 t/s; mtp:n_max=3,p_min=0.3 -> 11.1 t/s (byte-exact, 67% acceptance); prefill 30.5 t/s at 3.4K; KV self size at 4K ctx: 5.5 GiB -> 391 MiB. Not yet supported on the latent cache: K-shift/defrag (context shifting) - unreached in current usage. * openpangu: fence unsupported serving modes, truth-pass comments, drop dead weight/keys Post-audit hardening. The cache's position-indexed side state (MoME conv ring, DSA indexer keys) made several generic serving paths silently unsound; they are now fenced loudly instead of documented as unsupported: - s_l_position_ring flag on llama_kv_cache: the qnext-state predicate no longer claims the conv ring, so per-seq state save, seq_cp and the s_copy graph skip it - state save/restore refused for the arch at every llama_state_* entry (the ring and idx_l are not in the state format; restoring without them diverges silently) - K-shift/self-extend assert, defrag skips with a warning, server ctx_shift off via new llama_model_supports_ctx_shift() - single sequence enforced at context creation (n_seq_max > 1 refused) - server prompt-cache reuse limited to pure extension via new llama_model_supports_partial_kv_reuse(): mid-cache divergence reprocesses from scratch (the 16-column ring cannot rewind); multi-turn continuation stays fast - MTP draft length clamped to 13 via new llama_model_max_draft_tokens() so a rejected draft can never overwrite the ring columns the next decode reads - K/V cache types forced to f32 for the arch so the KV size log reports the truth - cache_size(): real latent-cache branch (was falling through to the ~14x larger materialized estimate used for offload planning) - unused fused wkv_b no longer loaded (TENSOR_SKIP; the graph runs entirely on the pre-split k_b/v_b), llm_prepare_mla openPangu special-case removed (it was a no-op) - stale v0 comments rewritten to describe the shipped graph; converter stops writing dead keys (dsa_layers, block_post_layernorm_idx) and the tokenizer pre-hash is registered in convert_hf_to_gguf_update.py Gates on this build: greedy spec output byte-identical to no-spec (EOS-terminated, sha-equal); 3.4K needle retrieved exactly; -np 2 / state save / n_max=20 / stale prefix reuse all refused or clamped with clear messages. * openpangu: assert kv_head == first batch position at graph build The ring, indexer and latent stores are addressed by absolute position through kv_head; the fences make append-only decode the only reachable mode, but the invariant was unchecked. Assert it at both graph entries (base and MTP) so any future cache plumbing that breaks it fails at build instead of corrupting output. Worst-case measurement builds pass pos = null and are exempt. * openpangu: cont h_pre before the mHC broadcast mul (CUDA binbcast misreads strided views) h_pre is a row-slice view of the fused mixes tensor. The CPU mul handles the strides; the CUDA broadcast path reads the view as if contiguous, so token 0 mixes correctly and every later token gets h_post/h_res rows instead. First divergent node in the whole graph (oracle rel 0.36 at opg0_attn_mhcpre_x, fixed to 7.5e-5). Sibling views h_post/h_res were already cont-wrapped, which is why only h_pre was exposed. * openpangu: keep DSA zero-trick sources finite (CUDA clamp propagates the 0*(-inf) NaN) The selection-mask base and zeros were built by scaling the MASKED scores by zero, but post-mask sc contains -inf and 0 * -inf = NaN. The CPU clamp launders NaN back to -1e30 (fminf/fmaxf ignore NaN); the CUDA clamp propagates it, so every DSA layer emitted NaN masks at n_kv > top_k and logits collapsed (observed: eval-callback CLAMP sum -1.3e36 on CPU vs nan on CUDA, 11748 NaNs downstream). Scale the pre-mask finite scores instead, which is correct on any backend regardless of clamp NaN semantics. Also defensively cont the strided KQ_mask slice feeding the score add (same strided-view kernel class as the mHC h_pre fix; unproven here but cheap). Gates after fix: 2600-token probe coherent, 3.4K needle exact ('7391') with and without MTP speculation, PP ~120 t/s. * openpangu: f16 latent KV cache option (explicit -ctk/-ctv f16 halves cache memory, f32 stays default) Track explicit cache-type requests through CLI/env; openPangu resolves no-request to f32 (unchanged), accepts explicit f32/f16, warns and falls back to f32 for BF16/quantized. Sink and cached-token KQ paths stay separate until after KQ so the latent cache is read directly without the f32-only concat; value is the sum of the sink and cache matmuls. Ring and DSA indexer caches stay f32; cache_size() follows the resolved types. * openpangu: enable graph reuse * openpangu: wire multi-head MTP drafting * openpangu: add MTP heads override * openpangu: keep MTP update logits last * openpangu: scope MTP warmup heads * speculative: apply per-request MTP heads before warmup * openpangu: fix multi-head MTP warmup computing on unwritten inputs Each chained head called the build_inp_* helpers itself, so the warmup and update graphs held one inp_tokens/inp_pos/inp_out_ids/KQ_mask tensor per head while llama_set_inputs only fills the tensors the lctx pointers reference, i.e. the last head's copies. Every head but the last read unwritten compute-buffer memory: with heads=3 active even head 1's ring, latent cache, and cached one-token draft were computed from garbage, which is why depth-1 acceptance measured 4% against 98% for the heads=1 control. Create the batch inputs once in build_openpangu and pass them to every build_openpangu_mtp call, and fix the two chaining errors that were hiding behind the garbage inputs: - Shift the chained hidden: head k+1's row at position p consumes head k's output row at p-1, the same convention head 1 uses for the target's conditioned hidden rows. The predecessor of a batch's first row lives in the previous warmup/update, carried across decodes through a new inp_mtp_carry input backed by lctx.mtp_carry (written back per ubatch, zeroed when a prompt warmup restarts from position 0). - Fill head 3's cache row at draft step 2: each draft step runs one head, so head 3's own decode at step 3 attended over a never-written row at the step-2 position. Pre-write it from the committed carry. Also include the active head count in the graph-reuse key next to the existing step index (reuse stays forced off for this arch). * speculative: default MTP drafting to a single head A stage without an explicit heads= override previously resolved to 0, meaning all model heads, so multi-head drafting was silently on by default for models that carry more than one NextN layer. Keep it opt-in (heads=N or heads=0 for all) until multi-head measures a win over the single-head config; single-head models are unaffected either way. * speculative: fence MTP head upshift over a warmed prefix Deeper NextN heads only hold valid cache rows for spans that were warmed with them. A request drafting with more MTP heads than the cached prefix was warmed with (e.g. a heads=1 conversation continued with heads=3, a pure extension the divergence fence deliberately allows) would read never-written deeper-head rows: verification keeps the output correct, but acceptance quietly collapses and any measurement taken there is misleading. Track the minimum head count the committed context has been warmed with since position 0 and have the server reprocess from scratch when a request asks for more. Also announce the model's NextN head count and the single-head default once at MTP context setup. * openpangu: skip dead MTP chain compute and stall-free carry readback The update chain's last head and the draft-time row fill only matter for their latent-cache and conv-ring writes; their FFN, norms, and shared head fed nothing. Add a cache-writes-only mode to the MTP block builder that returns after the attention block, and use it at both sites. The multi-head carry readback previously synchronized the scheduler after every warmup/update decode, a hard stall on CUDA. Issue the device-to-host copy async on the backend stream instead (stream order protects the source buffer from later graphs) and synchronize lazily when the host buffer is next consumed or resized. * openpangu: stop emitting fused kv_b tensor * openpangu: default latent cache to f16 * openpangu: refuse unsupported latent cache types * Window OpenPangu SWA cache reads * Gather OpenPangu DSA decode reads Gather DSA decode attention over the selected latent rows for OpenPangu base-model decode and verify graphs. The gathered branch now uses ggml_top_k order directly, runs maskless softmax over sinks plus selected rows for T <= 14, and derives values from the gathered k_l rows instead of the transposed latent cache. * Chunk OpenPangu indexer prefill scoring * Chunk OpenPangu prefill attention * Gather OpenPangu sparse prefill attention * Drop OpenPangu value cache * Add OpenPangu indexer cache type flag * Add OpenPangu q8_0 latent cache type Store the OpenPangu MLA latent K cache as q8_0 via -ctk q8_0 (about 0.53x of f16); the default stays f16 so behavior is unchanged without the flag. Latent V stays f16/f32. The q8 latent cache is a storage format only: it is dequanted to F32 before all compute. K reads go through openpangu_build_k_latent_for_read, V derivation through openpangu_build_v_latent_from_k (full 576-wide row to F32, then slice), and the DSA gather paths already dequant via get_rows. Feeding a q8 latent view directly into the KQ mul_mat corrupts large-context prefill, so that path is removed for quantized caches. The cache write stages ckv and kpe through F32 and writes one full 576-wide q8 row per token. Verified on a small discriminator model: the default f16 path is byte-identical to the prior code; the first-DSA-layer attention envelope is within 0.6% of the f16 cache (linf_rel 0.0057); top-k selection is bit-identical between cache types; the q8 latent cache is 0.531x the f16 size at 8K and 32K context; and generation stays coherent on both the dense and DSA-gather paths at all tested context lengths. * Remove OpenPangu debug trace env knobs and redundant DSA_TOPK override Drop the five LLAMA_OPENPANGU_*_TRACE debug-logging knobs (DSA_GATHER_TRACE, IDX_CHUNK_TRACE, ATT_CHUNK_TRACE, PREFILL_GATHER_TRACE, SWA_WINDOW_TRACE) and the LLAMA_OPENPANGU_DSA_TOPK override, which duplicated the -dsatk / --dsa-top-k CLI flag; top-k now comes solely from cparams.dsa_top_k. The five perf-tuning knobs (DSA_GATHER, IDX_CHUNK, ATT_CHUNK, ATT_KQ_MAX_MIB, PREFILL_GATHER) are retained pending the perf battery. No change to default behavior. * Subchunk OpenPangu DSA prefill gather to fit CUDA grid limit The prefill gathered-attention ggml_get_rows produced dst rows = topk * token_chunk (2048 * 256 = 524288) mapped to the CUDA grid.y dimension, which caps at 65535, crashing with GET_ROWS invalid argument at long context (N_KV around 10.5K with the natural topk of 2048). Split the prefill gather into token subchunks so topk * subchunk_tokens stays within the grid limit, and guard the decode gather with the same fit check (falling back to the dense masked path if a pathological topk would not fit). The subchunking is over the token dimension only, so per-token attention is unchanged and the result is numerically identical. Verified: the GPU sweep runs past the old crash boundary to 22K+ with zero CUDA errors; CPU and -ctk q8_0 paths unaffected. * openpangu: fix scheduler node budget for chunked DSA prefill; drop unused attn_kv_b; remove env tunables - Size the scheduler graph node budget for the chunked DSA prefill so 32K/ub2048 no longer trips the hash-set reservation assert; derive the extra budget from the builder's chunk/top-k/window structure with a fixed safety margin. - Remove LLAMA_OPENPANGU_* environment tunables from both the node-budget estimator and build_openpangu.cpp; use fixed constants in both so they stay in sync. - Converter: emit only the split attn_k_b/attn_v_b projections and drop the unused fused attn_kv_b tensor. * openpangu: restore DeepSeek converter kv_b; drop trace env + dead code; fix dense-fallback node budget - convert_hf_to_gguf.py: restore fused attn_kv_b in DeepseekV2Model (shared parent); openPangu subclass keeps split-only k_b/v_b. Stops newly-converted DeepSeek GGUFs from failing to load. - src/llama.cpp: remove LLAMA_GRAPH_REUSE_TRACE getenv, hit/miss counters, and the unconditional destructor log (no getenv or behavior change for any arch); node-budget estimator now covers the dense-fallback (n_swa==0) attention-chunk loop while skipping absent idx/top-k terms, preserving a strict overcount; remove unreachable openPangu split-cache block. - src/llama-context.h: drop now-dead graph_reuse_hits/misses members. - include/llama.h: move type_k/type_v/idx_type_k *_explicit bools to struct end to avoid a mid-struct ABI shift for out-of-tree consumers. - src/graphs/build_openpangu.cpp: replace vestigial env-struct singletons with the OPENPANGU_* constants; drop a redundant Sinkhorn permute round-trip (one transpose; greedy output verified byte-identical). Decode output unchanged (byte-identical greedy generation verified); shared-file changes are openPangu-gated or restore the pre-PR baseline. * openpangu: chat-parser support (reasoning split + thinking toggle) Two openPangu-only fixes, both gated on the arch-unique token <|pangu_text_start|> so no other model's parsing changes. - chat-diff-analyzer: add a workarounds entry that force-sets TAG_BASED reasoning with an empty start and a end. openPangu prefills in the generation prompt, so the output is delimited only by ; the differential detector otherwise learns start="" from the assistant-history form and fails to split, leaking reasoning into content. Same shape as the existing Laguna prefill patch. - chat.cpp: bridge enable_thinking to the template's `thinking` variable. openPangu's template gates reasoning on `thinking` rather than the ecosystem-standard `enable_thinking`, so the standard toggle was inert. An explicit `thinking` chat_template_kwarg still overrides via the extra_context merge. Blast radius: test-chat-auto-parser 437/437 unchanged; the sole test-chat-template diff is a pre-existing GLM trailing-newline. * openpangu: use ggml_cast for latent dequant reads Replace ggml_cpy(view, ggml_new_tensor_2d(F32, ...)) with ggml_cast in the MLA latent V-from-K and K-read helpers. ggml_cast emits the identical GGML_OP_CPY node into a fresh f32 tensor, so behavior is unchanged; it is the idiomatic form. Per review. * openpangu: narrow SWA reuse-key fields to 32-bit The openpangu_swa_window_view reuse key stored n_kv/n_tokens/window/pad as int64_t, but these are bounded well under 2^31 (window/pad are uint32_t at source; n_kv/n_tokens <= context length). Narrow to int32_t/uint32_t and drop the widening casts. w_view/win_off stay int64_t: they feed ggml view dims/offsets. Per review. * openpangu: precompute param_sink derived tensors at load The per-layer attention-sink block (sink_blk [576,NS]) and its transposed latent (s_lat_t [NS,512]) are pure functions of the layer weights, yet were rebuilt every eval across all 49 layers (RMS-norm + cast + concat + transpose). Compute them once at load, mirroring the wk_b derived-weight precompute, and read the stored tensors in build_openpangu_attention. Numerically identical; removes per-token work at decode. * openpangu: replace conv position-ring with ggml_ssm_conv + spec-rollback checkpoint Migrate the MoME depthwise causal conv (three sites per attention sublayer: qa-lora, compressed-kv, attn-out) from the bespoke 16-column position-indexed ring onto the core ggml_ssm_conv op with a recurrent conv-state slot. Cache: s_l becomes [2*conv_col_ne, qnext_state_slots], holding the (d_conv-1)=2 history taps per channel for the three sites (float offsets 0 / 2*n_lora_q / 2*(n_lora_q+n_lora_kv)). Drops the conv_hist_idx / conv_write_idx graph inputs and their fill in llama_set_inputs; adds one single-sequence sq input for ggml_ssm_conv shared across the three sites and the MTP head. Speculative rollback: the position ring self-healed rejected draft columns by absolute position; a recurrent slot does not, since seq_rm is a no-op for recurrent state. openPangu is admitted at the three spec-checkpoint save/init gates so the whole-slot shadow checkpoint (gpu-fallback) snapshots the conv slot before drafting and restores it before the accepted-token replay. The restore path is already keyed on ckpt.valid, so no gate change is needed there. Per-step checkpoint mode is declined for openPangu, which has no SSM recurrent term, so auto mode resolves to the whole-slot shadow. Gated: non-spec needle unchanged; MTP-spec needle correct with healthy draft acceptance (rollback verified via the acceptance canary). * openpangu: single ggml_concat copy for the latent cache store The non-quantized latent store split the [ckv | roped k_pe] row into two views and two cache copies, with a base_offset field on the CacheCopy struct to place the second one. Match the quantized path: concat the two parts and do one copy into the cache row. This drops the second cache-copy slot (OPENPANGU_COPY_K_KPE) and removes base_offset from CacheCopy entirely. Cache contents are unchanged: the concat writes the same [ckv 512 | k_pe 64] bytes to the same row. Gated on the needle for both the f16 latent path (the one that changed) and the q8 latent path, plus coherence. * openpangu: reuse the shared kr_l indexer cache instead of a separate idx_l The DSA lightning indexer stored its per-position keys in an openPangu-only idx_l cache, parallel to the kr_l indexer cache GLM-DSA already uses. Both have the same storage contract: [indexer_head_size, kv_size], idx_type_k dtype, one row per KV cell, written at kv_head and read [dim, n_kv] from zero. openPangu now allocates its indexer keys into kr_l and shares the dsa_cache_copies graph-reuse fixup. The fixup patch is factored into a helper that both the generic path and the openPangu update_cache_copies branch call, so the openPangu indexer copy is repointed to the current kv_head on graph reuse like every other cache write. This drops the idx_l vector, its allocation and memory accounting, and the openPangu third cache-copy slot (now one latent copy per layer). Per-arch allocation predicates stay separate (GLM uses indexer_is_full, openPangu uses the window==0 DSA schedule); only the kr_l storage and the copy fixup are shared. openPangu keeps its no-shift/no-defrag/no-state-I/O behavior, and the GLM Hadamard/k-shift logic stays GLM-gated. Gated: needle correct on f16 and q8 latent caches and under MTP speculation (acceptance unchanged at 0.67), plus coherence. * openpangu: discard pos-0 graphs from reuse; retire stale conv-state comments The ggml_ssm_conv refactor bakes the pos-0 conv-state reset into graph topology (a scale-by-zero node on the state view). A graph built at pos 0 could be reused at pos > 0 when the batch shape and padded n_kv match (a 1-token prompt followed by TG is the concrete case), zeroing the conv history on every reused decode. Admit openPangu at the existing reset_previous gate so pos-0 graphs are discarded from reuse, the same guard the qnext recurrent state relies on. Also retire the internal phase-plan comments the conv refactor left behind: they claimed the spec-checkpoint wiring had not landed in the commit that landed it, and misdescribed the s_l slot as awaiting rollback support. Gated: needle 8457 on f16 and q8 latent, MTP-spec needle (drafts fully accepted), coherence. * openpangu: drop the _explicit cache-type plumbing; validate unconditionally Review follow-up (item 1 of the second review). The explicit/default distinction carried less than claimed: the latent K/V fallback was f16, which is already the -ctk/-ctv and API default, so distinguishing unset from set-to-the-default bought nothing, and the two bools were behaviorally redundant. The only load-bearing use was the indexer cache, where openPangu defaulted to f32 while -ictk defaults to f16. Gating the f16 indexer directly (needle on f16 and q8 latent paths, MTP speculation, coherence) shows no quality difference, so openPangu now takes the standard f16 indexer default and the f32 special case is gone. Default indexer cache memory halves (64 -> 32 MiB at c 8192). Removes type_k_explicit/type_v_explicit/idx_type_k_explicit from llama.h, the cparams/mparams plumbing, and common; the resolve helpers become plain unconditional validators, so -ctk q8_0 is honored and an unsupported type errors out at load instead of silently coercing. Gated: needle 8457 on the new f16-indexer default, on q8 latent with MTP speculation, and with -ictk f32 explicitly honored (64 MiB f32 buffer in the load log); -ctk q4_0 and -ictk q4_1 refused with a clear error. * openpangu: keep MTP draft decodes position-contiguous under speculation The MTP framework's one-token draft shortcut caches a prediction one row past the accepted prefix during the accepted-token update, then skips re-decoding the last sampled token at the next draft round. A mask-addressed cache tolerates the resulting position gap; openPangu's position-addressed append-only cache (cell == position) does not: after a rollback the next draft decode lands one cell behind its position, and after a full acceptance the cache head sits one row ahead of the next draft base, either way tripping the kv_head == pos[0] invariant and aborting the server. The checkpoint admission in the conv refactor made this the standard openPangu speculative flow; the needle-first gates never generated enough draft rounds against a short prompt to reach it. Decline the shortcut re-seed for openPangu in mtp_accept_batch (restoring the drafting behavior all measured acceptance numbers were taken on) and trim rows at or beyond the draft base in mtp_speculative_gen_draft, so every draft decode stays position-contiguous with the cache head. Gated: the crashing flow (short prompt, 512-token spec generation, then a second request) completes with acceptance 0.60 prose / 0.87 code, matching the pre-checkpoint baseline profile; needle 8457 plus coherence on f16+spec and q8+spec. * openpangu: remove stale ring limits and fix MTP graph reuse * cli: preserve speculative carry on fallback Decode an already-emitted pending token when a draft cannot be used instead of sampling unchanged logits and duplicating output. Document single-head MTP as the default and multi-head modes as experimental. --------- Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com> --- common/chat-diff-analyzer.cpp | 21 + common/chat.cpp | 7 + common/common.cpp | 18 +- common/common.h | 3 + common/speculative.cpp | 118 ++- common/speculative.h | 7 + convert_hf_to_gguf.py | 145 ++++ convert_hf_to_gguf_update.py | 1 + docs/parameters.md | 4 +- docs/speculative.md | 1 + examples/main/main.cpp | 62 +- examples/server/server-context.cpp | 88 +- examples/server/server-context.h | 3 + examples/server/server-task.cpp | 18 + examples/server/server-task.h | 4 + gguf-py/gguf/constants.py | 119 +++ gguf-py/gguf/tensor_mapping.py | 88 ++ include/llama.h | 12 + src/CMakeLists.txt | 1 + src/graphs/build_openpangu.cpp | 1198 ++++++++++++++++++++++++++++ src/llama-arch.cpp | 5 + src/llama-arch.h | 26 + src/llama-build-context.cpp | 26 + src/llama-build-context.h | 35 + src/llama-context.h | 61 +- src/llama-hparams.cpp | 72 ++ src/llama-hparams.h | 10 + src/llama-load-tensors.cpp | 148 ++++ src/llama-model-loader.cpp | 1 + src/llama-model.cpp | 95 ++- src/llama-model.h | 28 +- src/llama-vocab.cpp | 1 + src/llama.cpp | 754 +++++++++++++++-- 33 files changed, 3073 insertions(+), 107 deletions(-) create mode 100644 src/graphs/build_openpangu.cpp diff --git a/common/chat-diff-analyzer.cpp b/common/chat-diff-analyzer.cpp index 8eba8bff4..95f1e8ed1 100644 --- a/common/chat-diff-analyzer.cpp +++ b/common/chat-diff-analyzer.cpp @@ -130,6 +130,27 @@ static std::vector in the generation prompt (like the Laguna case above), + // so the generated reasoning starts immediately and is delimited only by . The + // is concatenated into a larger literal ('...assistant\n'), so the + // standalone-literal reasoning detector does not pick it up; set the markers explicitly. + // Tool calls (<|tool_call_start|>[{...}]<|tool_call_end|>) are already handled by the auto-parser. + [](const common_chat_template & tmpl, autoparser & analysis) -> void { + if (tmpl.src.find("<|pangu_text_start|>") != std::string::npos) { + // Force-set (do not gate on mode==NONE): the differential detector sees the + // assistant-history form reasoning and sets start="", but at + // generation time is prompt-prefilled, so the output is delimited only by + // (start=""). Same shape as the Laguna patch above. + analysis.reasoning.mode = reasoning_mode::TAG_BASED; + analysis.reasoning.start = ""; + analysis.reasoning.end = ""; + if (std::find(analysis.preserved_tokens.begin(), analysis.preserved_tokens.end(), "") == + analysis.preserved_tokens.end()) { + analysis.preserved_tokens.push_back(""); + } + LOG_DBG(ANSI_ORANGE "[Patch: openPangu-2.0 thinking template]\n" ANSI_RESET); + } } }); diff --git a/common/chat.cpp b/common/chat.cpp index eda1b1e49..60ea5c403 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -781,6 +781,13 @@ static std::string common_chat_template_direct_apply_impl( {"eos_token", tmpl.eos_token()}, {"enable_thinking", inputs.enable_thinking}, }; + // openPangu's chat template gates reasoning on a `thinking` variable rather than the + // ecosystem-standard `enable_thinking`, so the normal toggle never reaches it. Bridge the + // standard control to it here so reasoning works on and off through `enable_thinking`. An + // explicit `thinking` chat_template_kwarg still wins (merged via extra_context below). + if (tmpl.source().find("<|pangu_text_start|>") != std::string::npos) { + inp["thinking"] = inputs.enable_thinking; + } if (tools_override.has_value() || !inputs.tools.empty()) { inp["tools"] = tools_override.has_value() ? *tools_override : inputs.tools; } diff --git a/common/common.cpp b/common/common.cpp index c05eec9df..6f0f1afa8 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -157,6 +157,9 @@ common_params_speculative common_params_speculative::with_stage_overrides(const if (stage.has_p_min_override()) { result.p_min = stage.p_min; } + if (stage.has_mtp_heads_override()) { + result.mtp_heads = stage.mtp_heads; + } if (stage.has_dflash_cross_ctx_override()) { result.dflash_cross_ctx = stage.dflash_cross_ctx; } @@ -182,6 +185,7 @@ common_params_speculative common_params_speculative::with_stage_overrides(const result.n_max = std::max(result.n_max, 0); result.n_min = std::max(0, std::min(result.n_min, result.n_max)); + result.mtp_heads = std::max(result.mtp_heads, 0); result.stages.clear(); return result; @@ -792,8 +796,8 @@ void gpt_params_parse_from_env(gpt_params & params) { get_env("LLAMA_ARG_CONT_BATCHING", params.cont_batching); get_env("LLAMA_ARG_HOST", params.hostname); get_env("LLAMA_ARG_PORT", params.port); - get_env("LLAMA_ARG_CACHE_TYPE_K", params.cache_type_k); - get_env("LLAMA_ARG_CACHE_TYPE_V", params.cache_type_v); + get_env("LLAMA_ARG_CACHE_TYPE_K", params.cache_type_k); + get_env("LLAMA_ARG_CACHE_TYPE_V", params.cache_type_v); get_env("LLAMA_ARG_MLOCK", params.use_mlock); get_env("LLAMA_ARG_K_CACHE_HADAMARD", params.k_cache_hadamard); get_env("LLAMA_ARG_V_CACHE_HADAMARD", params.v_cache_hadamard); @@ -926,6 +930,13 @@ static void common_speculative_stage_apply_kv( } return; } + if (key == "heads" || key == "mtp_heads") { + stage.mtp_heads = std::stoi(value_raw); + if (stage.mtp_heads < 0) { + throw std::invalid_argument("speculative stage mtp_heads must be >= 0"); + } + return; + } if (key == "cross_ctx" || key == "dflash_cross_ctx") { stage.dflash_cross_ctx = std::stoi(value_raw); if (stage.dflash_cross_ctx < 1) { @@ -3320,7 +3331,8 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param " cpu serialise state via llama_state_seq; re-decode on rejection" }); options.push_back({ "*", "--spec-type SPEC[:k=v,...]", "canonical speculative stage entry; repeat for a supported two-stage chain.\n" "types: none, draft, dflash, mtp, ngram-cache, ngram-simple, ngram-map-k, ngram-map-k4v, ngram-mod, suffix\n" - "canonical keys: n_max,n_min,p_min,cross_ctx,ngram_size_n,ngram_size_m,ngram_min_hits,suffix_min_match_len,suffix_max_depth,suffix_corpus\n" + "canonical keys: n_max,n_min,p_min,heads,cross_ctx,ngram_size_n,ngram_size_m,ngram_min_hits,suffix_min_match_len,suffix_max_depth,suffix_corpus\n" + "MTP heads: heads=1 is the default; heads>1 and heads=0 (all model heads) are experimental\n" "for comma-bearing string values, quote the value inside the stage payload for normal shell use\n" "if argv is passed directly without shell unescaping, the parser also accepts escaped commas as \\,\n" "examples: --spec-type mtp:n_max=1,p_min=0.0\n" diff --git a/common/common.h b/common/common.h index c590dcdd7..1eaf66444 100644 --- a/common/common.h +++ b/common/common.h @@ -172,6 +172,7 @@ struct common_speculative_stage_params { int32_t n_max = -1; int32_t n_min = -1; float p_min = -1.0f; + int32_t mtp_heads = -1; int32_t dflash_cross_ctx = -1; uint16_t ngram_size_n = 0; @@ -185,6 +186,7 @@ struct common_speculative_stage_params { bool has_n_max_override() const { return n_max >= 0; } bool has_n_min_override() const { return n_min >= 0; } bool has_p_min_override() const { return p_min >= 0.0f; } + bool has_mtp_heads_override() const { return mtp_heads >= 0; } bool has_dflash_cross_ctx_override() const { return dflash_cross_ctx >= 0; } bool has_ngram_size_n_override() const { return ngram_size_n > 0; } bool has_ngram_size_m_override() const { return ngram_size_m > 0; } @@ -218,6 +220,7 @@ struct common_params_speculative { int32_t n_max = 16; // number of tokens to draft during speculative decoding int32_t n_min = 0; // minimum number of tokens to draft during speculative decoding std::vector stages; // explicit stage chain for single-spec or self-spec + model fallback + int32_t mtp_heads = 1; // MTP heads to use; 1 is the default, while >1 and 0 (all model heads) are experimental int32_t dflash_cross_ctx = 512; // target-feature context window for DFlash float p_split = 0.1f; // speculative decoding split probability diff --git a/common/speculative.cpp b/common/speculative.cpp index d91e0296e..9091d26d6 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -24,6 +24,8 @@ #define SPEC_VOCAB_CHECK_START_TOKEN_ID 5 void llama_set_mtp_target_context(struct llama_context * ctx, struct llama_context * target_ctx); +void llama_set_mtp_step_idx(struct llama_context * ctx, int32_t mtp_step_idx); +void llama_set_mtp_n_heads(struct llama_context * ctx, int32_t mtp_n_heads); const std::vector common_speculative_types = { COMMON_SPECULATIVE_TYPE_NONE, @@ -212,12 +214,13 @@ static std::vector mtp_speculative_gen_draft( struct llama_context * ctx, int n_draft, float p_min, + int32_t mtp_heads, llama_token id_last, llama_pos n_past, llama_seq_id seq_id, bool constant_draft_positions = false); -static int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch & batch, bool is_prompt_warmup); +static int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch & batch, bool is_prompt_warmup, int32_t mtp_heads); struct mtp_last_embd { std::vector embd; @@ -228,7 +231,19 @@ struct mtp_last_embd { struct common_speculative_state_mtp : public common_speculative_state { llama_context * ctx_tgt; llama_context * ctx_mtp = nullptr; + int32_t mtp_heads_active = 0; + // number of NextN heads the model carries, and the minimum head count the committed + // context has been warmed with since position 0 (deeper heads' caches only hold valid + // rows for spans warmed with them; a request drafting with MORE heads than the cached + // prefix was warmed with must reprocess from scratch). Single-sequence by design, like + // the rest of the openPangu MTP state. + int32_t n_heads_model = 1; + int32_t mtp_warmed_heads = 0; common_sampler * smpl; + + int32_t resolved_heads() const { + return mtp_heads_active > 0 ? std::min(mtp_heads_active, n_heads_model) : n_heads_model; + } // For Gemma 4 external MTP assistant: draft positions are held constant bool constant_draft_positions = false; int n_embd = 0; @@ -252,9 +267,15 @@ struct common_speculative_state_mtp : public common_speculative_state { smpl = common_sampler_init(llama_get_model(ctx_mtp), sparams); llama_set_mtp_target_context(ctx_mtp, ctx_tgt); n_embd = llama_mtp_state_n_embd(ctx_mtp); + n_heads_model = std::max(1, llama_model_n_nextn_layer(llama_get_model(ctx_mtp))); LOG_INF("%s: MTP context ready (n_ctx=%d, constant_draft_positions=%s)\n", __func__, llama_n_ctx(ctx_mtp), constant_draft_positions ? "true" : "false"); + if (n_heads_model > 1) { + LOG_INF("%s: model carries %d NextN/MTP heads; heads=1 is the default, while " + "heads>1 and heads=0 (all model heads) are experimental\n", + __func__, n_heads_model); + } } ~common_speculative_state_mtp() override { @@ -304,6 +325,7 @@ struct common_speculative_state_mtp : public common_speculative_state { } llama_context * ctx = ctx_mtp; + mtp_heads_active = std::max(0, params.mtp_heads); const auto hidden_it = target_hidden_by_seq.find(seq_id); if (hidden_it == target_hidden_by_seq.end() || (int) hidden_it->second.size() != n_embd) { @@ -323,6 +345,7 @@ struct common_speculative_state_mtp : public common_speculative_state { ctx, params.n_max, params.p_min, + params.mtp_heads, id_last, n_past, seq_id, @@ -1072,6 +1095,7 @@ static common_params_speculative common_speculative_get_runtime_params( result.n_max = stage.has_n_max_override() ? stage.n_max : params.n_max; result.n_min = stage.has_n_min_override() ? stage.n_min : params.n_min; result.p_min = stage.has_p_min_override() ? stage.p_min : params.p_min; + result.mtp_heads = stage.has_mtp_heads_override() ? stage.mtp_heads : params.mtp_heads; if (config.type == COMMON_SPECULATIVE_TYPE_SUFFIX) { result.suffix_min_match_len = stage.has_suffix_min_match_len_override() @@ -1081,11 +1105,43 @@ static common_params_speculative common_speculative_get_runtime_params( result.n_max = std::max(result.n_max, 0); result.n_min = std::max(0, std::min(result.n_min, result.n_max)); + result.mtp_heads = std::max(result.mtp_heads, 0); result.stages.clear(); return result; } +bool common_speculative_mtp_requires_fresh_warmup(const common_speculative * spec) { + const auto * mtp_state = common_speculative_get_mtp_state(spec); + if (mtp_state == nullptr || mtp_state->n_heads_model <= 1) { + return false; + } + + // drafting with more heads than the cached prefix was warmed with would read + // never-written deeper-head cache rows; the caller must reprocess from position 0 + return mtp_state->resolved_heads() > mtp_state->mtp_warmed_heads && mtp_state->mtp_warmed_heads > 0; +} + +void common_speculative_prepare_request(common_speculative * spec, common_params_speculative & params) { + if (spec == nullptr) { + return; + } + + const auto runtime_stages = params.get_resolved_stages(); + const bool use_runtime_stage_overrides = common_speculative_stage_chain_matches(runtime_stages, spec->configs); + + for (size_t i = 0; i < spec->impls.size(); ++i) { + auto * mtp_state = dynamic_cast(spec->impls[i].get()); + if (mtp_state == nullptr) { + continue; + } + + const auto & runtime_stage = use_runtime_stage_overrides ? runtime_stages[i] : spec->configs[i].stage; + common_params_speculative impl_params = common_speculative_get_runtime_params(spec->configs[i], params, runtime_stage); + mtp_state->mtp_heads_active = std::max(0, impl_params.mtp_heads); + } +} + static common_ngram_map get_common_ngram_map(const common_speculative_config & config) { uint16_t size_key = config.params.ngram_size_n; uint16_t size_value = config.params.ngram_size_m; @@ -1281,7 +1337,8 @@ common_speculative * common_speculative_init( configs.push_back(common_speculative_config(stage, stage_params)); } - if (!configs.empty() && llama_model_has_recurrent(llama_get_model(ctx_tgt))) { + if (!configs.empty() && (llama_model_has_recurrent(llama_get_model(ctx_tgt)) || + llama_model_is_openpangu(llama_get_model(ctx_tgt)))) { const int ckpt_tokens = std::max(1, params.get_max_stage_n_max() + 1); const int actual_mode = llama_spec_ckpt_init(ctx_tgt, params.recurrent_ckpt_mode, ckpt_tokens); if (actual_mode == LLAMA_SPEC_CKPT_NONE) { @@ -1407,6 +1464,7 @@ common_speculative * common_speculative_init( /* .configs = */ std::move(configs), /* .impls = */ std::move(impls) }; + common_speculative_prepare_request(result, params); // initialize autotune if requested if (params.autotune && params.has_composite_stage_chain()) { @@ -2683,10 +2741,20 @@ static int32_t mtp_accept_batch( if (!llama_set_draft_input_hidden_state_copy(state.ctx_mtp, hidden_rows, hidden_rows_floats)) { return -1; } - if (mtp_update_kv_cache(state.ctx_mtp, accepted_batch, false) != 0) { + if (mtp_update_kv_cache(state.ctx_mtp, accepted_batch, false, state.mtp_heads_active) != 0) { return -1; } + if (llama_model_is_openpangu(llama_get_model(state.ctx_mtp))) { + // The one-token draft shortcut re-seeded below would skip re-decoding the last + // sampled token next round, leaving a hole at its position. openPangu's KV cache + // is position-addressed append-only (cell == position), so draft decodes must be + // position-contiguous; decline the shortcut and let the next round decode the + // sampled token normally. + mtp_invalidate_cached_draft(state, seq_id); + return 0; + } + auto & last = mtp_get_last_embd(state, seq_id); const float * embd = llama_get_embeddings_ith(state.ctx_mtp, accepted_batch.n_tokens - 1); if (embd != nullptr) { @@ -2767,6 +2835,18 @@ int32_t common_speculative_on_target_batch( const float * last_hidden = hidden_rows_storage.data() + (size_t) (batch.n_tokens - 1) * features.width; mtp_store_target_hidden(*mtp_state, seq_id, last_hidden, features.width); + // track the minimum head count the committed context has been warmed with: a fresh + // position-0 warmup resets it, everything after can only narrow it + { + const int32_t resolved = mtp_state->resolved_heads(); + if (is_prompt_warmup && batch.pos != nullptr && batch.n_tokens > 0 && batch.pos[0] == 0) { + mtp_state->mtp_warmed_heads = resolved; + } else { + mtp_state->mtp_warmed_heads = mtp_state->mtp_warmed_heads > 0 + ? std::min(mtp_state->mtp_warmed_heads, resolved) : resolved; + } + } + if (mtp_state->constant_draft_positions) { mtp_invalidate_cached_draft(*mtp_state, seq_id); return 0; @@ -2804,7 +2884,7 @@ int32_t common_speculative_on_target_batch( if (!llama_set_draft_input_hidden_state_copy(mtp_state->ctx_mtp, conditioned_hidden_rows, hidden_rows_storage.size())) { return -1; } - const int32_t ret = mtp_update_kv_cache(mtp_state->ctx_mtp, batch, true); + const int32_t ret = mtp_update_kv_cache(mtp_state->ctx_mtp, batch, true, mtp_state->mtp_heads_active); mtp_invalidate_cached_draft(*mtp_state, seq_id); return ret; } @@ -2839,6 +2919,7 @@ std::vector mtp_speculative_gen_draft( struct llama_context * ctx, int n_draft, float p_min, + int32_t mtp_heads, llama_token id_last, llama_pos n_past, llama_seq_id seq_id, @@ -2856,7 +2937,22 @@ std::vector mtp_speculative_gen_draft( common_sampler_reset(smpl); + if (llama_model_is_openpangu(llama_get_model(ctx)) && + llama_kv_cache_seq_pos_max(ctx, seq_id) >= n_past) { + // Position-addressed cache: drafting restarts at n_past, so any rows at or beyond + // it (the accepted-update writes one row past the accepted prefix) must be dropped + // first to keep the draft decode position-contiguous with the cache head. + llama_kv_cache_seq_rm(ctx, seq_id, n_past, -1); + } + + const int n_embd = llama_mtp_state_n_embd(ctx); + const int n_mtp_heads_model = std::max(1, llama_model_n_nextn_layer(llama_get_model(ctx))); + const int n_mtp_heads = mtp_heads > 0 + ? std::max(1, std::min((int) mtp_heads, n_mtp_heads_model)) + : n_mtp_heads_model; + llama_batch mtp_batch = llama_batch_init(1, 0, 1); + llama_set_mtp_n_heads(ctx, n_mtp_heads); llama_set_mtp_op_type(ctx, MTP_OP_DRAFT_GEN); float prob; @@ -2864,7 +2960,6 @@ std::vector mtp_speculative_gen_draft( llama_token current_input_id = id_last; llama_pos current_n_past = n_past; - const int n_embd = llama_mtp_state_n_embd(ctx); auto & last = mtp_get_last_embd(state, seq_id); int i0 = 0; @@ -2878,6 +2973,8 @@ std::vector mtp_speculative_gen_draft( current_n_past++; if (!llama_set_draft_input_hidden_state_copy(ctx, last.embd.data(), last.embd.size())) { llama_batch_free(mtp_batch); + llama_set_mtp_step_idx(ctx, 0); + llama_set_mtp_n_heads(ctx, 0); llama_set_mtp_op_type(ctx, MTP_OP_NONE); return drafts; } @@ -2889,6 +2986,7 @@ std::vector mtp_speculative_gen_draft( mtp_batch.n_tokens = 0; const llama_pos draft_pos = constant_draft_positions ? n_past : current_n_past; common_batch_add(mtp_batch, current_input_id, draft_pos, {seq_id}, true); + llama_set_mtp_step_idx(ctx, std::min(i, n_mtp_heads - 1)); ++n_decode; if (llama_decode(ctx, mtp_batch) != 0) { @@ -2922,6 +3020,8 @@ std::vector mtp_speculative_gen_draft( } } llama_batch_free(mtp_batch); + llama_set_mtp_step_idx(ctx, 0); + llama_set_mtp_n_heads(ctx, 0); llama_set_mtp_op_type(ctx, MTP_OP_NONE); // Purge the metadata for the draft tokens. @@ -2938,7 +3038,7 @@ std::vector mtp_speculative_gen_draft( } -int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch& batch, bool is_prompt_warmup) { +int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch& batch, bool is_prompt_warmup, int32_t mtp_heads) { if (batch.n_tokens == 0) { return 0; } @@ -2962,12 +3062,18 @@ int32_t mtp_update_kv_cache(struct llama_context * ctx, const llama_batch& batch } mtp_batch.logits[mtp_batch.n_tokens-1] = true; if (is_prompt_warmup) { + llama_set_mtp_n_heads(ctx, mtp_heads); + llama_set_mtp_step_idx(ctx, 0); llama_set_mtp_op_type(ctx, MTP_OP_WARMUP); } else { + llama_set_mtp_n_heads(ctx, mtp_heads); + llama_set_mtp_step_idx(ctx, 0); llama_set_mtp_op_type(ctx, MTP_OP_UPDATE_ACCEPTED); } const int32_t ret = llama_decode(ctx, mtp_batch); + llama_set_mtp_step_idx(ctx, 0); + llama_set_mtp_n_heads(ctx, 0); llama_set_mtp_op_type(ctx, MTP_OP_NONE); return ret; } diff --git a/common/speculative.h b/common/speculative.h index da740c0b6..cf0fad6e1 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -82,6 +82,13 @@ void common_speculative_free(common_speculative * spec); // optionally call once at the beginning of a new generation void common_speculative_begin(common_speculative * spec, const llama_tokens & prompt); +// apply per-request runtime parameters before prompt warmup can touch companion state +void common_speculative_prepare_request(common_speculative * spec, common_params_speculative & params); + +// true when the active request drafts with more MTP heads than the cached prefix was +// warmed with; the caller must then reprocess the prompt from position 0 +bool common_speculative_mtp_requires_fresh_warmup(const common_speculative * spec); + // sample up to n_draft tokens and add them to the batch using the draft model // draft_base_pos/draft_seq_id override the MTP position for id_last llama_tokens common_speculative_draft( diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index bfc42f6e6..f3540dea9 100644 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -592,6 +592,9 @@ class Model: if chkhsh == "66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273": # ref: https://huggingface.co/alvarobartt/grok-2-tokenizer res = "grok-2" + if chkhsh == "65df2fe396b537a53433301848c0a739f56d56f67ad3d35eba27961ac33c12bb": + # ref: https://huggingface.co/openpangu/openPangu-2.0-Flash + res = "openpangu" if chkhsh == "972da7b59cec44d1f0a490a86c96df53859e486e481563e5dddac155013d87ac": # ref: https://huggingface.co/poolside/Laguna-XS.2 res = "laguna" @@ -4588,6 +4591,148 @@ class DeepseekV2Model(Model): raise ValueError(f"Unprocessed experts: {experts}") +@Model.register("OpenPanguV2ForCausalLM") +class OpenPanguV2Model(DeepseekV2Model): + # openPangu-2.0-Flash: MLA + DSA/SWA hybrid + MoE + mHC(Hyper-Connections) + MoME convs. + # Emits a complete, self-contained GGUF: weights (incl. pre-split attn_k_b/attn_v_b for + # the latent-attention graph) plus the DSA/SWA schedule and mHC/MoME/sink metadata. + model_arch = gguf.MODEL_ARCH.OPENPANGU + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # MTP/NextN layers (num_hidden_layers .. +num_nextn_predict_layers-1) are real blocks + # in ik_llama's layout (n_layer includes NextN; n_layer_kv_from_start excludes them). + self._nextn = int(self.hparams.get("num_nextn_predict_layers", 0) or 0) + self.block_count = int(self.hparams["num_hidden_layers"]) + self._nextn + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + def set_vocab(self): + # OpenPanguV2Tokenizer is a GPT2/BPE-style tokenizer. The pre-tokenizer hash is + # registered in Model.get_vocab_base_pre (see the openpangu entry). + self._set_vocab_gpt2() + # HF prepends <|pangu_text_start|> via the tokenizer post-processor; tokenizer_config + # has no add_bos_token key, so state it explicitly for the GGUF. + self.gguf_writer.add_add_bos_token(True) + + def set_gguf_parameters(self): + # Base transformer params (block_count now includes NextN layers). + Model.set_gguf_parameters(self) + hparams = self.hparams + arch = gguf.MODEL_ARCH_NAMES[self.model_arch] + + self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"]) + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"]) + self.gguf_writer.add_kv_lora_rank(hparams["kv_lora_rank"]) + self.gguf_writer.add_key_length(hparams["qk_nope_head_dim"] + hparams["qk_rope_head_dim"]) + self.gguf_writer.add_value_length(hparams["v_head_dim"]) + + # MoE + self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_count(hparams["n_routed_experts"]) + self.gguf_writer.add_expert_used_count(hparams["num_experts_per_tok"]) + self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"]) + self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"]) + # router_enable_expert_bias => sigmoid gating with e_score_correction bias + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + + # RoPE + self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"]) + self.gguf_writer.add_rope_freq_base(hparams["rope_theta"]) + + # NextN / MTP + self.gguf_writer.add_uint32( + gguf.Keys.LLM.NEXTN_PREDICT_LAYERS.format(arch=arch), self._nextn + ) + + # DSA lightning indexer + self.gguf_writer.add_uint32( + gguf.Keys.Attention.INDEXER_HEAD_COUNT.format(arch=arch), hparams["index_n_heads"] + ) + self.gguf_writer.add_uint32( + gguf.Keys.Attention.INDEXER_KEY_LENGTH.format(arch=arch), hparams["index_head_dim"] + ) + self.gguf_writer.add_uint32( + gguf.Keys.Attention.INDEXER_TOP_K.format(arch=arch), hparams["index_topk"] + ) + + # SWA (window; last few layers widen to 2048 per sliding_window_list) + if hparams.get("sliding_window") is not None: + self.gguf_writer.add_sliding_window(hparams["sliding_window"]) + + # Pangu-specific structural metadata (consumed by the OPENPANGU graph). The runtime + # derives everything else from these plus tensor presence: DSA layers are the + # windowless base layers (dsa_layers is redundant with swa_layers, so it is not + # written), and block_post_norm placement follows the tensors themselves. + self.gguf_writer.add_uint32(f"{arch}.mhc_num_stream", hparams["mhc_num_stream"]) + self.gguf_writer.add_uint32(f"{arch}.mhc_recur_norm", hparams["mhc_recur_norm"]) + self.gguf_writer.add_uint32(f"{arch}.param_sink_number", hparams["param_sink_number"]) + self.gguf_writer.add_array(f"{arch}.swa_layers", hparams["swa_layers"]) + if hparams.get("sliding_window_list") is not None: + self.gguf_writer.add_array(f"{arch}.sliding_window_list", hparams["sliding_window_list"]) + + _experts: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # Sigmoid router bias: rename to the deepseek-style ".bias" the tensor map expects. + if name.endswith("e_score_correction_bias"): + name = name.replace("e_score_correction_bias", "e_score_correction.bias") + + # NOTE: unlike DeepseekV2, we do NOT skip MTP layers — they are real blocks here. + + # Merge routed experts into stacked 3D tensors (same layout as deepseek/glm4moe). + if name.find("mlp.experts") != -1: + n_experts = self.hparams["n_routed_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + for w_name in ["down_proj", "gate_proj", "up_proj"]: + datas: list[Tensor] = [] + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + data_torch = torch.stack(datas, dim=0) + merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + tensors.append((self.map_tensor_name(merged_name), data_torch)) + return tensors + else: + return [] + + # Split the fused MLA kv_b into k_b / v_b (deepseek MLA layout). OpenPangu's + # graph consumes the pre-split tensors directly, so do not emit the fused copy. + if name.endswith("kv_b_proj.weight"): + name_kb = name.replace("kv_b_proj", "k_b_proj") + name_vb = name.replace("kv_b_proj", "v_b_proj") + + n_head_kv = self.hparams["num_attention_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + + assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope_head_dim) + + kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope_head_dim, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) + k_b = k_b.transpose(1, 2) + k_b = k_b.reshape(n_head_kv * data_torch.shape[-1], qk_nope_head_dim) + v_b = v_b.reshape(n_head_kv * v_head_dim, data_torch.shape[-1]) + + return [ + (self.map_tensor_name(name_kb), k_b), + (self.map_tensor_name(name_vb), v_b), + ] + + # Everything else (attn/norms/mHC/MoME conv/param-sink/indexer/nextn) maps by name. + return [(self.map_tensor_name(name), data_torch)] + + @Model.register("T5WithLMHeadModel") @Model.register("T5ForConditionalGeneration") @Model.register("MT5ForConditionalGeneration") diff --git a/convert_hf_to_gguf_update.py b/convert_hf_to_gguf_update.py index 0c9b6dfd3..709e3d6e2 100755 --- a/convert_hf_to_gguf_update.py +++ b/convert_hf_to_gguf_update.py @@ -107,6 +107,7 @@ models = [ {"name": "minimax-m2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/MiniMaxAI/MiniMax-M2", }, {"name": "mellum2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/JetBrains/Mellum2-12B-A2.5B-Base", }, {"name": "gpt-4o", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/openai/gpt-oss-20b", "chkhsh": "ccc2ef013c104be7bae2965776d611e1d7a8a2a9c547dd93a682c9a9fc80352e", }, # o200k_harmony shares the GPT-4o pre-tokenizer + {"name": "openpangu", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/openpangu/openPangu-2.0-Flash", "chkhsh": "65df2fe396b537a53433301848c0a739f56d56f67ad3d35eba27961ac33c12bb", }, # primary source: https://gitcode.com/ascend-tribe/openPangu-2.0-Flash ] diff --git a/docs/parameters.md b/docs/parameters.md index f627e3b28..38baf22a2 100644 --- a/docs/parameters.md +++ b/docs/parameters.md @@ -132,7 +132,7 @@ Check the details [here](./speculative.md). | `-ctkd, --cache-type-k-draft TYPE` | KV cache data type for K for the draft model | - | For draft model, see: `-ctk` | | `-ctvd, --cache-type-v-draft TYPE` | KV cache data type for V for the draft model | - | For draft model, see: `-ctk` | | `-draft, --draft-params` | Comma-separated list of draft model parameters | - | | -| `--spec-type SPEC[:k=v,...]` | Canonical speculative stage entry; repeat to configure the supported two-stage chain | - | Types: `none`, `draft`, `dflash`, `mtp`, `ngram-cache`, `ngram-simple`, `ngram-map-k`, `ngram-map-k4v`, `ngram-mod`, `suffix`. Canonical keys: `n_max`, `n_min`, `p_min`, `cross_ctx`, `ngram_size_n`, `ngram_size_m`, `ngram_min_hits`, `suffix_min_match_len`, `suffix_max_depth`, `suffix_corpus`. String values may escape commas as `\,` or quote the value inside the stage payload. Examples: `--spec-type ngram-mod:n_max=64,n_min=2,ngram_size_n=8 --spec-type mtp:n_max=1,p_min=0.0`, `--model-draft draft.gguf --spec-type dflash:n_max=4,cross_ctx=512` | +| `--spec-type SPEC[:k=v,...]` | Canonical speculative stage entry; repeat to configure the supported two-stage chain | - | Types: `none`, `draft`, `dflash`, `mtp`, `ngram-cache`, `ngram-simple`, `ngram-map-k`, `ngram-map-k4v`, `ngram-mod`, `suffix`. Canonical keys include `n_max`, `n_min`, `p_min`, `heads`, `cross_ctx`, `ngram_size_n`, `ngram_size_m`, `ngram_min_hits`, `suffix_min_match_len`, `suffix_max_depth`, `suffix_corpus`. For MTP, `heads=1` is the default; values above `1` and `heads=0` (all model heads) are experimental. String values may escape commas as `\,` or quote the value inside the stage payload. Examples: `--spec-type ngram-mod:n_max=64,n_min=2,ngram_size_n=8 --spec-type mtp:n_max=1,p_min=0.0`, `--model-draft draft.gguf --spec-type dflash:n_max=4,cross_ctx=512` | | `--spec-autotune` | Automatically tune speculative params to maximize tokens/sec | - | Automatically determines the near-optimal arguments for the type of speculation being performed [PR 1595](https://github.com/ikawrakow/ik_llama.cpp/pull/1595) | | `--recurrent-ckpt-mode MODE` | Checkpoint strategy for recurrent/hybrid speculative decoding | auto | One of: - `auto` auto-select: per-step if CUDA full-GPU, gpu-fallback otherwise - `per-step` save SSM state per draft step in VRAM; no re-decode on rejection - `gpu-fallback` copy state to GPU buffer; re-decode on rejection - `cpu` serialise state via llama_state_seq; re-decode on rejection [PR 1669](https://github.com/ikawrakow/ik_llama.cpp/pull/1669) [PR 1774](https://github.com/ikawrakow/ik_llama.cpp/pull/1774) | @@ -383,7 +383,7 @@ WIP | `--override-kv KEY=TYPE:VALUE` | Override model metadata by key | - | Advanced option to override model metadata by key. May be specified multiple times. types: int, float, bool, str. Example: `--override-kv tokenizer.ggml.add_bos_token=bool:false` | | `-m, --model FNAME` | Model path | models/$filename | Mandatory, the GGUF model file to be served. | | `-md, --model-draft FNAME` | Draft model for speculative decoding | unused | Required when an explicit `draft` stage is used. | -| `--spec-type SPEC[:k=v,...]` | Canonical speculative stage entry; repeat for the supported two-stage chain | none | Use stage-local keys like `n_max`, `n_min`, `p_min`, `ngram_size_n`, `ngram_size_m`, `ngram_min_hits`, `suffix_min_match_len`, `suffix_max_depth`, and `suffix_corpus`. | +| `--spec-type SPEC[:k=v,...]` | Canonical speculative stage entry; repeat for the supported two-stage chain | none | Use stage-local keys like `n_max`, `n_min`, `p_min`, `heads`, `ngram_size_n`, `ngram_size_m`, `ngram_min_hits`, `suffix_min_match_len`, `suffix_max_depth`, and `suffix_corpus`. For MTP, `heads=1` is the default; values above `1` and `heads=0` (all model heads) are experimental. | ### Request-Level Speculative Overrides diff --git a/docs/speculative.md b/docs/speculative.md index 6665dc3a5..af99c9867 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -128,6 +128,7 @@ Canonical stage keys: | `n_max` | Maximum drafted tokens for that stage | | `n_min` | Minimum usable drafted tokens for that stage | | `p_min` | Minimum speculative probability threshold | +| `heads` | MTP heads to use; `1` is the default, while values above `1` and `0` (all model heads) are experimental | | `ngram_size_n` | Lookup n-gram size | | `ngram_size_m` | Draft m-gram size | | `ngram_min_hits` | Minimum matching hits for n-gram map stages | diff --git a/examples/main/main.cpp b/examples/main/main.cpp index 2a45c3a67..1760ad0bb 100644 --- a/examples/main/main.cpp +++ b/examples/main/main.cpp @@ -941,17 +941,20 @@ int main(int argc, char ** argv) { const int n_predict_budget = n_remain < 0 ? std::numeric_limits::max() : n_remain; bool used_speculative = false; + const bool sampled_before_from_carry = have_speculative_sampled; + llama_token sampled_before = LLAMA_TOKEN_NULL; + if (sampled_before_from_carry) { + sampled_before = speculative_sampled; + have_speculative_sampled = false; + speculative_sampled = LLAMA_TOKEN_NULL; + } + bool sampled_before_ready = sampled_before_from_carry; if (spec != nullptr && n_predict_budget != 1) { - const bool sampled_before_from_carry = have_speculative_sampled; - llama_token sampled_before = LLAMA_TOKEN_NULL; - if (sampled_before_from_carry) { - sampled_before = speculative_sampled; - have_speculative_sampled = false; - speculative_sampled = LLAMA_TOKEN_NULL; - } else { + if (!sampled_before_ready) { sampled_before = common_sampler_sample_legacy(ctx_sampling, ctx, ctx_guidance); common_sampler_accept(ctx_sampling, ctx, sampled_before, /* apply_grammar= */ true); + sampled_before_ready = true; } static const llama_tokens empty_speculative_tokens; const llama_tokens & draft_history = @@ -981,7 +984,7 @@ int main(int argc, char ** argv) { const int min_usable_draft = params.speculative.get_min_usable_stage_n_min(); if ((int) draft.size() >= min_usable_draft && (!draft.empty() || n_predict_budget > 1)) { - if (llama_model_has_recurrent(model)) { + if (llama_model_has_recurrent(model) || llama_model_is_openpangu(model)) { if (!common_speculative_before_draft( spec, model, @@ -1066,23 +1069,42 @@ int main(int argc, char ** argv) { } if (!used_speculative) { - const llama_token id = common_sampler_sample_legacy(ctx_sampling, ctx, ctx_guidance); - common_sampler_accept(ctx_sampling, ctx, id, /* apply_grammar= */ true); + if (sampled_before_ready) { + // A carried token was already displayed by the previous speculative step but + // still needs to be decoded. If a new draft is unusable, stage that token instead + // of sampling the unchanged logits again and emitting it twice. + embd.push_back(sampled_before); + embd_is_prompt = false; + input_echo = true; - LOG("last: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx, ctx_sampling->prev).c_str()); + if (sampled_before_from_carry) { + emitted_generated = false; + } else { + emitted = embd; + emitted_generated = true; + --n_remain; + } - embd.push_back(id); - emitted = embd; - embd_is_prompt = false; - emitted_generated = true; + LOG("n_remain: %d\n", n_remain); + } else { + const llama_token id = common_sampler_sample_legacy(ctx_sampling, ctx, ctx_guidance); + common_sampler_accept(ctx_sampling, ctx, id, /* apply_grammar= */ true); - // echo this to console - input_echo = true; + LOG("last: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx, ctx_sampling->prev).c_str()); - // decrement remaining sampling budget - --n_remain; + embd.push_back(id); + emitted = embd; + embd_is_prompt = false; + emitted_generated = true; - LOG("n_remain: %d\n", n_remain); + // echo this to console + input_echo = true; + + // decrement remaining sampling budget + --n_remain; + + LOG("n_remain: %d\n", n_remain); + } } else { input_echo = true; } diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index fcdae78be..61c0b7e9f 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -164,7 +164,15 @@ static common_speculative_stage_params server_parse_speculative_stage_json(const continue; } - throw std::runtime_error("Error: per-request speculative.stages only support type, n_max, n_min, and p_min; structural stage overrides are startup-only"); + if (item.key() == "heads" || item.key() == "mtp_heads") { + stage.mtp_heads = item.value().get(); + if (stage.mtp_heads < 0) { + throw std::runtime_error("Error: speculative.stages[].heads must be >= 0"); + } + continue; + } + + throw std::runtime_error("Error: per-request speculative.stages only support type, n_max, n_min, p_min, and heads; structural stage overrides are startup-only"); } return stage; @@ -509,6 +517,8 @@ void server_slot::reset() { // Reset speculative decoding stats n_draft_total = 0; n_draft_accepted = 0; + n_draft_by_depth.clear(); + n_draft_accepted_by_depth.clear(); chat_msg = {}; json_schema = json(); generated_tool_call_ids.clear(); @@ -608,7 +618,7 @@ void server_slot::release() { json server_slot::get_formated_timings() const { - return json{ + json timings = json{ {"prompt_n", n_prompt_tokens_processed}, {"prompt_ms", t_prompt_processing}, {"prompt_per_token_ms", t_prompt_processing / n_prompt_tokens_processed}, @@ -622,6 +632,27 @@ json server_slot::get_formated_timings() const { {"n_ctx", n_ctx}, {"n_past", n_past}, }; + if (n_draft_total > 0) { + timings["draft_n"] = n_draft_total; + timings["draft_n_accepted"] = n_draft_accepted; + json by_depth = json::array(); + for (size_t i = 0; i < n_draft_by_depth.size(); ++i) { + if (n_draft_by_depth[i] <= 0) { + continue; + } + const int32_t accepted = i < n_draft_accepted_by_depth.size() + ? n_draft_accepted_by_depth[i] : 0; + by_depth.push_back({ + {"depth", (int32_t) i + 1}, + {"draft_n", n_draft_by_depth[i]}, + {"draft_n_accepted", accepted}, + }); + } + if (!by_depth.empty()) { + timings["draft_by_depth"] = by_depth; + } + } + return timings; } result_timings server_slot::get_timings() const { @@ -644,6 +675,8 @@ result_timings server_slot::get_timings() const { if (n_draft_total > 0) { timings.draft_n = n_draft_total; timings.draft_n_accepted = n_draft_accepted; + timings.draft_n_by_depth = n_draft_by_depth; + timings.draft_n_accepted_by_depth = n_draft_accepted_by_depth; } return timings; @@ -1158,6 +1191,9 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task) if (stage_override.has_p_min_override()) { slot.params.speculative.stages[i].p_min = stage_override.p_min; } + if (stage_override.has_mtp_heads_override()) { + slot.params.speculative.stages[i].mtp_heads = stage_override.mtp_heads; + } } const auto resolved = slot.params.speculative.get_resolved_stages(); @@ -1195,6 +1231,7 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task) if (!common_speculative_validate_chain(slot.params.speculative, &spec_error)) { throw std::runtime_error("Error: invalid speculative request configuration: " + spec_error); } + common_speculative_prepare_request(slot.spec, slot.params.speculative); } catch (const std::exception & e) { send_error(task, e.what(), ERROR_TYPE_INVALID_REQUEST); return false; @@ -1762,6 +1799,12 @@ bool server_context::launch_slot_with_task(server_slot& slot, server_task& task) LOG_WARNING("%s\n", "ctx_shift is not implemented for split mode graph, it will be disabled"); } } + if (!llama_model_supports_ctx_shift(llama_get_model(slot.ctx))) { + if (params_base.ctx_shift) { + params_base.ctx_shift = false; + LOG_WARNING("%s\n", "ctx_shift is not supported by this model's KV cache, it will be disabled"); + } + } { const auto& stop = data.find("stop"); if (stop != data.end() && stop->is_array()) { @@ -3488,6 +3531,12 @@ void server_context::add_sampled_tokens() { } else { // keep track of total number of drafted tokens tested slot.n_draft_total += draft.size(); + if (slot.n_draft_by_depth.size() < draft.size()) { + slot.n_draft_by_depth.resize(draft.size(), 0); + } + for (size_t i = 0; i < draft.size(); ++i) { + slot.n_draft_by_depth[i]++; + } // add all drafted tokens to the batch for (size_t i = 0; i < draft.size(); i++) { @@ -3822,6 +3871,30 @@ void server_context::batch_pending_prompt(const int32_t n_ubatch, const int32_t slot.n_past = prefix.first; slot.n_past_prompt = prefix.second; slot.n_past_offset = slot.n_past_prompt - slot.n_past; + if (!llama_model_supports_partial_kv_reuse(model) && + slot.n_past < (int32_t) slot.cache_tokens.size()) { + // the cache diverges from the new prompt mid-sequence; this + // model can only extend or reset a cached sequence (per-position + // side state past the divergence point is already lost) + LLAMA_LOG_INFO("%s: cached sequence diverges at %d/%d and this model does not support partial KV reuse - reprocessing from scratch\n", + __func__, (int) slot.n_past, (int) slot.cache_tokens.size()); + slot.n_past = 0; + slot.n_past_prompt = 0; + slot.n_past_offset = 0; + } + + if (slot.n_past > 0 && slot.spec != nullptr && + common_speculative_mtp_requires_fresh_warmup(slot.spec)) { + // the request drafts with more MTP heads than the cached + // prefix was warmed with; deeper-head cache rows for the + // reused span were never written + LLAMA_LOG_INFO("%s: request drafts with more MTP heads than the cached prefix was warmed with - reprocessing from scratch\n", + __func__); + slot.n_past = 0; + slot.n_past_prompt = 0; + slot.n_past_offset = 0; + } + if ((slot.n_past + size_threshold < slot.cache_tokens.size())) { int32_t back = 4; @@ -4118,7 +4191,14 @@ void server_context::speculative_decoding_accept() { slot.t_token_generation = std::max(1, t_current - slot.t_start_generation) / 1e3; // update how many tokens out of those tested were accepted - slot.n_draft_accepted += ids.size() - 1; + const size_t n_draft_accepted = ids.size() - 1; + slot.n_draft_accepted += n_draft_accepted; + if (slot.n_draft_accepted_by_depth.size() < n_draft_accepted) { + slot.n_draft_accepted_by_depth.resize(n_draft_accepted, 0); + } + for (size_t i = 0; i < n_draft_accepted; ++i) { + slot.n_draft_accepted_by_depth[i]++; + } // rollback to the state before sampling the draft tokens slot.cache_tokens.keep_first(slot.cache_tokens.n_tokens() - n_draft); @@ -4722,7 +4802,7 @@ void server_context::update_slots() { // make sure we're in the right embedding mode llama_set_embeddings(ctx, batch_type == 1); - if (llama_model_has_recurrent(model)) { + if (llama_model_has_recurrent(model) || llama_model_is_openpangu(model)) { const int ckpt_mode = params_base.speculative.recurrent_ckpt_mode; for (auto & slot : slots) { diff --git a/examples/server/server-context.h b/examples/server/server-context.h index 015311a5a..03f00046f 100644 --- a/examples/server/server-context.h +++ b/examples/server/server-context.h @@ -6,6 +6,7 @@ #include #include +#include @@ -163,6 +164,8 @@ struct server_slot { // speculative decoding stats int32_t n_draft_total = 0; // Total draft tokens generated int32_t n_draft_accepted = 0; // Draft tokens actually accepted + std::vector n_draft_by_depth; + std::vector n_draft_accepted_by_depth; int32_t n_past_se = 0; // self-extend diff --git a/examples/server/server-task.cpp b/examples/server/server-task.cpp index 82be8b454..c35ce7107 100644 --- a/examples/server/server-task.cpp +++ b/examples/server/server-task.cpp @@ -20,6 +20,24 @@ json result_timings::to_json() const { if (draft_n > 0) { base["draft_n"] = draft_n; base["draft_n_accepted"] = draft_n_accepted; + if (!draft_n_by_depth.empty()) { + json by_depth = json::array(); + for (size_t i = 0; i < draft_n_by_depth.size(); ++i) { + if (draft_n_by_depth[i] <= 0) { + continue; + } + const int32_t accepted = i < draft_n_accepted_by_depth.size() + ? draft_n_accepted_by_depth[i] : 0; + by_depth.push_back({ + {"depth", (int32_t) i + 1}, + {"draft_n", draft_n_by_depth[i]}, + {"draft_n_accepted", accepted}, + }); + } + if (!by_depth.empty()) { + base["draft_by_depth"] = by_depth; + } + } } return base; diff --git a/examples/server/server-task.h b/examples/server/server-task.h index 417d0a3f7..9685cd28c 100644 --- a/examples/server/server-task.h +++ b/examples/server/server-task.h @@ -8,6 +8,8 @@ // TODO: prevent including the whole server-common.h as we only use server_tokens #include "server-common.h" +#include + using json = nlohmann::ordered_json; enum stop_type { @@ -130,6 +132,8 @@ struct result_timings { // Optional speculative metrics - only included when > 0 int32_t draft_n = 0; int32_t draft_n_accepted = 0; + std::vector draft_n_by_depth; + std::vector draft_n_accepted_by_depth; json to_json() const; }; diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index e3b05643b..0433d5cea 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -113,6 +113,9 @@ class Keys: CAUSAL = "{arch}.attention.causal" Q_LORA_RANK = "{arch}.attention.q_lora_rank" KV_LORA_RANK = "{arch}.attention.kv_lora_rank" + INDEXER_HEAD_COUNT = "{arch}.attention.indexer.head_count" + INDEXER_KEY_LENGTH = "{arch}.attention.indexer.key_length" + INDEXER_TOP_K = "{arch}.attention.indexer.top_k" REL_BUCKETS_COUNT = "{arch}.attention.relative_buckets_count" SLIDING_WINDOW = "{arch}.attention.sliding_window" SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" @@ -260,6 +263,7 @@ class MODEL_ARCH(IntEnum): ARCTIC = auto() DEEPSEEK2 = auto() GLM4_MOE = auto() + OPENPANGU = auto() CHATGLM = auto() BITNET = auto() BITNET_25 = auto() @@ -386,6 +390,33 @@ class MODEL_TENSOR(IntEnum): MTP_CENTROIDS = auto() DFLASH_FC = auto() DFLASH_HIDDEN_NORM = auto() + # openPangu-2.0 (DSA lightning indexer) + INDEXER_K_NORM = auto() + INDEXER_PROJ = auto() # weights_proj + INDEXER_ATTN_K = auto() # wk + INDEXER_ATTN_Q_B = auto() # wq_b + # openPangu-2.0 (MoME causal-conv on MLA latents) + ATTN_QA_CONV = auto() + ATTN_KV_CONV = auto() # compresskv_conv + ATTN_O_CONV = auto() + # openPangu-2.0 (learned static param sink) + ATTN_PARAM_SINK_KV = auto() # param_sink_compressed_kv + ATTN_PARAM_SINK_K_PE = auto() # param_sink_k_pe + # openPangu-2.0 (mHC / Hyper-Connections: per-attn, per-mlp, global merge) + MHC_ATTN_PHI = auto() + MHC_ATTN_ALPHA = auto() + MHC_ATTN_BETA = auto() + MHC_ATTN_GAMMA = auto() + MHC_MLP_PHI = auto() + MHC_MLP_ALPHA = auto() + MHC_MLP_BETA = auto() + MHC_MLP_GAMMA = auto() + MHC_MERGE_PHI = auto() + MHC_MERGE_ALPHA = auto() # branch_alpha_pre + MHC_MERGE_BETA = auto() # branch_beta_pre + MHC_MERGE_GAMMA = auto() + # openPangu-2.0 (sandwich norm: extra whole-block post-norm on a layer subset) + BLOCK_POST_NORM = auto() MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { @@ -437,6 +468,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.DEEPSEEK2: "deepseek2", MODEL_ARCH.CHATGLM: "chatglm", MODEL_ARCH.GLM4_MOE: "glm4moe", + MODEL_ARCH.OPENPANGU: "openpangu", MODEL_ARCH.BITNET: "bitnet", MODEL_ARCH.BITNET_25: "bitnet-25", MODEL_ARCH.T5: "t5", @@ -563,6 +595,29 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.MTP_CENTROIDS: "mtp_centroids", MODEL_TENSOR.DFLASH_FC: "dflash_fc", MODEL_TENSOR.DFLASH_HIDDEN_NORM: "dflash_hidden_norm", + # openPangu-2.0 + MODEL_TENSOR.INDEXER_K_NORM: "blk.{bid}.attn_indexer_k_norm", + MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.attn_indexer_weights_proj", + MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.attn_indexer_k", + MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.attn_indexer_q_b", + MODEL_TENSOR.ATTN_QA_CONV: "blk.{bid}.attn_qa_conv", + MODEL_TENSOR.ATTN_KV_CONV: "blk.{bid}.attn_compresskv_conv", + MODEL_TENSOR.ATTN_O_CONV: "blk.{bid}.attn_o_conv", + MODEL_TENSOR.ATTN_PARAM_SINK_KV: "blk.{bid}.attn_param_sink_kv", + MODEL_TENSOR.ATTN_PARAM_SINK_K_PE: "blk.{bid}.attn_param_sink_k_pe", + MODEL_TENSOR.MHC_ATTN_PHI: "blk.{bid}.attn_mhc_phi", + MODEL_TENSOR.MHC_ATTN_ALPHA: "blk.{bid}.attn_mhc_alpha", + MODEL_TENSOR.MHC_ATTN_BETA: "blk.{bid}.attn_mhc_beta", + MODEL_TENSOR.MHC_ATTN_GAMMA: "blk.{bid}.attn_mhc_gamma", + MODEL_TENSOR.MHC_MLP_PHI: "blk.{bid}.mlp_mhc_phi", + MODEL_TENSOR.MHC_MLP_ALPHA: "blk.{bid}.mlp_mhc_alpha", + MODEL_TENSOR.MHC_MLP_BETA: "blk.{bid}.mlp_mhc_beta", + MODEL_TENSOR.MHC_MLP_GAMMA: "blk.{bid}.mlp_mhc_gamma", + MODEL_TENSOR.MHC_MERGE_PHI: "merge_mhc_phi", + MODEL_TENSOR.MHC_MERGE_ALPHA: "merge_mhc_alpha", + MODEL_TENSOR.MHC_MERGE_BETA: "merge_mhc_beta", + MODEL_TENSOR.MHC_MERGE_GAMMA: "merge_mhc_gamma", + MODEL_TENSOR.BLOCK_POST_NORM: "blk.{bid}.block_post_norm", } MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { @@ -1298,6 +1353,70 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.OPENPANGU: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + # MLA attention (deepseek-style) + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, # post_attention_layernorm (sandwich) + # DSA lightning indexer + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + # MoME causal convs + param sink + MODEL_TENSOR.ATTN_QA_CONV, + MODEL_TENSOR.ATTN_KV_CONV, + MODEL_TENSOR.ATTN_O_CONV, + MODEL_TENSOR.ATTN_PARAM_SINK_KV, + MODEL_TENSOR.ATTN_PARAM_SINK_K_PE, + # MoE (routed + shared + sigmoid bias) and dense-lead FFN + MODEL_TENSOR.FFN_NORM, # pre_mlp_layernorm + MODEL_TENSOR.FFN_POST_NORM, # post_mlp_layernorm (sandwich) + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + # mHC / Hyper-Connections + MODEL_TENSOR.MHC_ATTN_PHI, + MODEL_TENSOR.MHC_ATTN_ALPHA, + MODEL_TENSOR.MHC_ATTN_BETA, + MODEL_TENSOR.MHC_ATTN_GAMMA, + MODEL_TENSOR.MHC_MLP_PHI, + MODEL_TENSOR.MHC_MLP_ALPHA, + MODEL_TENSOR.MHC_MLP_BETA, + MODEL_TENSOR.MHC_MLP_GAMMA, + MODEL_TENSOR.MHC_MERGE_PHI, + MODEL_TENSOR.MHC_MERGE_ALPHA, + MODEL_TENSOR.MHC_MERGE_BETA, + MODEL_TENSOR.MHC_MERGE_GAMMA, + # sandwich extra block post-norm (layer subset) + MODEL_TENSOR.BLOCK_POST_NORM, + # NextN / MTP tail (3 layers) + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], MODEL_ARCH.DFLASH: [ MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 2a5f38413..a4cba86f8 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -87,6 +87,20 @@ class TensorNameMap: MODEL_TENSOR.ROPE_FACTORS_LONG: (), MODEL_TENSOR.ROPE_FACTORS_SHORT: (), + + # openPangu-2.0 global mHC stream-merge module (non-block) + MODEL_TENSOR.MHC_MERGE_PHI: ( + "model.merge_mhc_module.phi", + ), + MODEL_TENSOR.MHC_MERGE_ALPHA: ( + "model.merge_mhc_module.branch_alpha_pre", + ), + MODEL_TENSOR.MHC_MERGE_BETA: ( + "model.merge_mhc_module.branch_beta_pre", + ), + MODEL_TENSOR.MHC_MERGE_GAMMA: ( + "model.merge_mhc_module.norm_gamma", + ), } block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = { @@ -697,6 +711,80 @@ class TensorNameMap: "model.layers.{bid}.post_attention_layernorm", ), }, + # openPangu-2.0: disambiguate the sandwich norms (post_attention/pre_mlp/post_mlp) + # from the generic collisions, and pin all Pangu-only block tensors. + MODEL_ARCH.OPENPANGU: { + MODEL_TENSOR.ATTN_POST_NORM: ( + "model.layers.{bid}.post_attention_layernorm", + ), + MODEL_TENSOR.FFN_NORM: ( + "model.layers.{bid}.pre_mlp_layernorm", + ), + MODEL_TENSOR.FFN_POST_NORM: ( + "model.layers.{bid}.post_mlp_layernorm", + ), + MODEL_TENSOR.BLOCK_POST_NORM: ( + "model.layers.{bid}.block_post_layernorm", + ), + MODEL_TENSOR.FFN_EXP_PROBS_B: ( + "model.layers.{bid}.mlp.e_score_correction", + ), + # DSA lightning indexer + MODEL_TENSOR.INDEXER_K_NORM: ( + "model.layers.{bid}.self_attn.indexer.k_norm", + ), + MODEL_TENSOR.INDEXER_PROJ: ( + "model.layers.{bid}.self_attn.indexer.weights_proj", + ), + MODEL_TENSOR.INDEXER_ATTN_K: ( + "model.layers.{bid}.self_attn.indexer.wk", + ), + MODEL_TENSOR.INDEXER_ATTN_Q_B: ( + "model.layers.{bid}.self_attn.indexer.wq_b", + ), + # MoME causal convs + MODEL_TENSOR.ATTN_QA_CONV: ( + "model.layers.{bid}.self_attn.qa_conv", + ), + MODEL_TENSOR.ATTN_KV_CONV: ( + "model.layers.{bid}.self_attn.compresskv_conv", + ), + MODEL_TENSOR.ATTN_O_CONV: ( + "model.layers.{bid}.self_attn.o_conv", + ), + # learned static param sink + MODEL_TENSOR.ATTN_PARAM_SINK_KV: ( + "model.layers.{bid}.self_attn.param_sink_compressed_kv", + ), + MODEL_TENSOR.ATTN_PARAM_SINK_K_PE: ( + "model.layers.{bid}.self_attn.param_sink_k_pe", + ), + # mHC / Hyper-Connections (per attn + per mlp) + MODEL_TENSOR.MHC_ATTN_PHI: ( + "model.layers.{bid}.attn_mhc_module.phi", + ), + MODEL_TENSOR.MHC_ATTN_ALPHA: ( + "model.layers.{bid}.attn_mhc_module.branch_alpha", + ), + MODEL_TENSOR.MHC_ATTN_BETA: ( + "model.layers.{bid}.attn_mhc_module.branch_beta", + ), + MODEL_TENSOR.MHC_ATTN_GAMMA: ( + "model.layers.{bid}.attn_mhc_module.norm_gamma", + ), + MODEL_TENSOR.MHC_MLP_PHI: ( + "model.layers.{bid}.mlp_mhc_module.phi", + ), + MODEL_TENSOR.MHC_MLP_ALPHA: ( + "model.layers.{bid}.mlp_mhc_module.branch_alpha", + ), + MODEL_TENSOR.MHC_MLP_BETA: ( + "model.layers.{bid}.mlp_mhc_module.branch_beta", + ), + MODEL_TENSOR.MHC_MLP_GAMMA: ( + "model.layers.{bid}.mlp_mhc_module.norm_gamma", + ), + }, } mapping: dict[str, tuple[MODEL_TENSOR, str]] diff --git a/include/llama.h b/include/llama.h index c2c183c15..ce5458335 100644 --- a/include/llama.h +++ b/include/llama.h @@ -698,6 +698,9 @@ extern "C" { LLAMA_API bool llama_model_has_recurrent(const struct llama_model * model); + // Returns true if the model is openPangu (conv-only recurrent state that rides the spec-rollback checkpoint) + LLAMA_API bool llama_model_is_openpangu(const struct llama_model * model); + // Returns true if the model is a Gemma 4 MTP assistant (external frozen-KV speculative drafter) LLAMA_API bool llama_model_is_gemma4_mtp_assistant(const struct llama_model * model); @@ -705,6 +708,15 @@ extern "C" { LLAMA_API bool llama_model_is_split_mode_graph(const struct llama_model * model); + // Returns false for models whose KV cache cannot be re-positioned after the fact + // (K-shift / context shift / self-extend), e.g. openPangu's latent cache. + LLAMA_API bool llama_model_supports_ctx_shift(const struct llama_model * model); + + // Returns false for models that can only reuse a cached sequence as a pure extension: + // rewinding into the middle of a decoded sequence loses per-position side state + // (e.g. openPangu keeps only the current recurrent conv state). + LLAMA_API bool llama_model_supports_partial_kv_reuse(const struct llama_model * model); + LLAMA_API const char * llama_model_arch_string(const struct llama_model * model); // Returns 0 on success diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 4a9ed6657..0ee65e188 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -111,6 +111,7 @@ add_library(llama graphs/build_gptneox.cpp graphs/build_arctic.cpp graphs/build_deepseek2.cpp + graphs/build_openpangu.cpp graphs/build_glm4.cpp graphs/build_bitnet.cpp graphs/build_cohere2.cpp diff --git a/src/graphs/build_openpangu.cpp b/src/graphs/build_openpangu.cpp new file mode 100644 index 000000000..b0c6faac4 --- /dev/null +++ b/src/graphs/build_openpangu.cpp @@ -0,0 +1,1198 @@ +#include "../llama-build-context.h" +#include "../llama-model.h" +#include "../llama-context.h" + +#include +#include +#include + +static constexpr int OPENPANGU_CACHE_COPIES_PER_LAYER = 1; +static constexpr int OPENPANGU_COPY_K_CKV = 0; +static constexpr int64_t OPENPANGU_DSA_GATHER_MIN_RATIO = 2; +// Keep these fixed constants in sync with llama_openpangu_chunked_graph_nodes() +// in llama.cpp; the scheduler budget mirrors the chunk loops below. +static constexpr int64_t OPENPANGU_IDX_SCORE_CHUNK = 256; +static constexpr int64_t OPENPANGU_ATT_SCORE_CHUNK = 256; +static constexpr int64_t OPENPANGU_ATT_FULL_KQ_MAX_MIB = 1024; +static constexpr int64_t OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX = 65535; + +static std::vector & openpangu_cache_copies(llama_context & lctx) { + return lctx.cparams.mtp_op_type == MTP_OP_NONE ? lctx.openpangu_cache_copies : lctx.openpangu_cache_copies_mtp; +} + +static void openpangu_clear_cache_copies(llama_context & lctx) { + auto & copies = openpangu_cache_copies(lctx); + std::fill(copies.begin(), copies.end(), llama_context::CacheCopy{}); + std::fill(lctx.dsa_cache_copies.begin(), lctx.dsa_cache_copies.end(), llama_context::CacheCopy{}); +} + +static void openpangu_register_cache_copy( + llama_context & lctx, int il, int slot, ggml_tensor * cpy, size_t step) { + GGML_ASSERT(slot >= 0 && slot < OPENPANGU_CACHE_COPIES_PER_LAYER); + auto & copies = openpangu_cache_copies(lctx); + const size_t idx = (size_t) OPENPANGU_CACHE_COPIES_PER_LAYER*il + slot; + GGML_ASSERT(idx < copies.size()); + copies[idx].cpy = cpy; + copies[idx].step = step; +} + +static uint32_t openpangu_kv_cache_pad(const llama_cparams & cparams) { + return cparams.flash_attn ? 256u : 32u; +} + +static bool openpangu_idx_score_should_chunk(int64_t n_tokens, int64_t chunk) { + return chunk > 0 && n_tokens > 14 && n_tokens > chunk; +} + +static bool openpangu_att_score_should_chunk( + int64_t n_kv_eff, int64_t n_sinks, int64_t n_head, int64_t n_tokens, + int64_t chunk, int64_t cap_mib) { + if (chunk <= 0 || cap_mib <= 0 || n_tokens <= 14 || n_tokens <= chunk) { + return false; + } + + const double full_kq_bytes = + (double) (n_kv_eff + n_sinks) * (double) n_head * (double) n_tokens * (double) sizeof(float); + const double cap_bytes = (double) cap_mib * 1024.0 * 1024.0; + return full_kq_bytes > cap_bytes; +} + +static bool openpangu_dsa_gather_should_engage(int64_t n_kv, int64_t n_tokens, int64_t topk, int64_t pad) { + // Gather must prune at least half the cache to beat get_rows/cast/cont-transpose copy + // overhead; measured 2K regression at ratio ~1 where top_k covers nearly all cache rows. + return n_tokens <= 14 && n_kv >= OPENPANGU_DSA_GATHER_MIN_RATIO*topk + pad + n_tokens; +} + +static bool openpangu_dsa_prefill_gather_should_engage( + int64_t n_kv, int64_t n_tokens, int64_t chunk_start, int64_t chunk_size, int64_t topk, int64_t pad) { + GGML_UNUSED(chunk_size); + if (n_tokens <= 14 || topk <= 0 || chunk_start < 0 || chunk_start >= n_tokens || n_kv < n_tokens) { + return false; + } + + const int64_t min_causal_kv = n_kv - n_tokens + chunk_start; + return min_causal_kv >= OPENPANGU_DSA_GATHER_MIN_RATIO*topk + pad; +} + +static bool openpangu_dsa_gather_rows_fit_cuda(int64_t topk, int64_t n_tokens) { + return topk > 0 && n_tokens > 0 && topk <= OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX / n_tokens; +} + +static int64_t openpangu_dsa_gather_tokens_per_get_rows(int64_t topk) { + GGML_ASSERT(topk > 0 && topk <= OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX); + return std::max(1, OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX / topk); +} + +static ggml_tensor * openpangu_build_v_latent_from_k( + ggml_context * ctx, const llama_kv_cache & kv_self, int il, + int64_t kv_lora_rank, int64_t n_kv_view, int64_t win_off) { + ggml_tensor * kl = kv_self.k_l[il]; + if (ggml_is_quantized(kl->type)) { + ggml_tensor * full_view = ggml_view_2d(ctx, kl, kl->ne[0], n_kv_view, + kl->nb[1], (size_t) win_off*kl->nb[1]); + ggml_tensor * full_f32 = ggml_cast(ctx, full_view, GGML_TYPE_F32); + ggml_tensor * v_src = ggml_view_2d(ctx, full_f32, kv_lora_rank, n_kv_view, + full_f32->nb[1], 0); + return ggml_cont(ctx, ggml_transpose(ctx, v_src)); + } + + ggml_tensor * v_src = ggml_view_2d(ctx, kl, kv_lora_rank, n_kv_view, + kl->nb[1], (size_t) win_off*kl->nb[1]); + return ggml_cont(ctx, ggml_transpose(ctx, v_src)); +} + +static ggml_tensor * openpangu_build_k_latent_for_read( + ggml_context * ctx, const llama_kv_cache & kv_self, int il, + int64_t n_kv_view, int64_t win_off) { + ggml_tensor * kl = kv_self.k_l[il]; + ggml_tensor * k_view = ggml_view_2d(ctx, kl, kl->ne[0], n_kv_view, + kl->nb[1], (size_t) win_off*kl->nb[1]); + if (!ggml_is_quantized(kl->type)) { + return k_view; + } + + return ggml_cast(ctx, k_view, GGML_TYPE_F32); +} + +static ggml_tensor * openpangu_cast_for_latent_cache_write(ggml_context * ctx, ggml_tensor * src, ggml_tensor * kl) { + return ggml_is_quantized(kl->type) && src->type != GGML_TYPE_F32 ? ggml_cast(ctx, src, GGML_TYPE_F32) : src; +} + +static ggml_tensor * openpangu_cast_gathered_latent_for_cache_type(ggml_context * ctx, ggml_tensor * src, ggml_tensor * kl) { + return !ggml_is_quantized(kl->type) && src->type != kl->type ? ggml_cast(ctx, src, kl->type) : src; +} + +static ggml_tensor * openpangu_build_swa_mask_for_graph(llm_build_context & llm, uint32_t window, bool * windowed) { + *windowed = false; + llm.lctx.openpangu_swa_window_view = {}; + + if (window == 0) { + return nullptr; + } + + const uint32_t pad = openpangu_kv_cache_pad(llm.cparams); + const llama_openpangu_swa_window_view view = + llama_openpangu_calc_swa_window_view(llm.n_kv, llm.n_tokens, window, pad); + + if (!view.engaged) { + return llm.build_inp_KQ_mask_swa(); + } + + llm.lctx.openpangu_swa_window_view = { + true, + llm.n_kv, + llm.n_tokens, + window, + pad, + view.w_view, + view.win_off, + }; + *windowed = true; + return llm.build_inp_KQ_mask_swa_win(view.w_view); +} + +// openPangu-2.0-Flash graph. +// +// Attention runs absorbed MLA over a latent KV cache: per position the cache stores only +// the 512-d compressed latent plus the 64-d roped k_pe (k_l, straight layout); +// q_nope is projected into latent space through attn_k_b and the +// attention output is up-projected through attn_v_b after the weighted sum. Per-head K/V +// never materialize. +// +// Pangu-specific pieces implemented here (see vault Stage-2b forward spec): +// - mHC / Hyper-Connections: 4 parallel residual streams mixed per sublayer via a phi +// projection (h_pre combine-in, h_post/h_res scatter-out) with a 20-iter Sinkhorn. +// - MoME: causal depthwise conv (k=3) on the q-lora latent, compressed-kv latent, attn out; +// decode taps come from a recurrent ggml_ssm_conv state slot (openpangu_causal_conv below). +// - param_sink: 128 learned latent-space KV entries per layer, prepended to every query's +// attention span outside the causal/window/top-k masks. +// - DSA + SWA schedule: windowed base layers use the SWA mask; windowless base layers run +// the lightning indexer over a per-position indexer-key cache and restrict +// attention to the top-k scored positions plus the sinks. Schedule-less GGUFs run dense. +// For prompts <= 512 tokens both mechanisms are inert and output is bit-exact to dense. +// - sandwich norms (post_attention / pre_mlp / post_mlp) + block_post on a layer subset. +// - NextN/MTP layers (build_openpangu_mtp below) drive the mtp speculative framework, +// chaining conv state through the same recurrent slot. + +// --- causal depthwise conv1d, kernel=3: out[t] = w0*x[t-2] + w1*x[t-1] + w2*x[t] (per channel) --- +// x: [C, n_tokens]; w: ggml tensor with ne = {3, 1, C} (kernel-major). Returns [C, n_tokens]. +// MOME: out = x + depthwise causal conv1d(k=3). Every Infer call site passes +// residual_connection=1; the tap magnitudes confirm the conv is a small learned +// perturbation on top of the identity, not a standalone filter. +// +// Conv state: `state_all` is this layer's cache_s_l recurrent slot table +// [2*col_ne, qnext_state_slots]. Slot 0 holds the single-sequence decode state, with +// this site's two taps packed at float offset 2*site_off. The buffer is zeroed at cache +// allocation and reset at pos 0, preserving zero history at sequence start (pos-0 graphs +// are discarded from reuse via reset_previous, so the baked reset never runs at pos > 0). +// Speculative rollback snapshots/restores the whole slot via the spec checkpoint. +static ggml_tensor * openpangu_causal_conv(ggml_context * ctx, ggml_cgraph * gf, + ggml_tensor * x, ggml_tensor * w, + ggml_tensor * state_all, int64_t site_off, + ggml_tensor * seq_qnext, bool reset_state) { + const int64_t C = x->ne[0]; + const int64_t T = x->ne[1]; + // weight is stored f16 with ne = {3, C}: per-channel taps contiguous. ggml_ssm_conv + // expects an f32 kernel, so cast once (tiny tensor). + ggml_tensor * wc = ggml_reshape_2d(ctx, ggml_cast(ctx, w, GGML_TYPE_F32), 3, C); + + GGML_ASSERT(state_all != nullptr); + GGML_ASSERT(state_all->type == GGML_TYPE_F32); + GGML_ASSERT(state_all->ne[0] >= 2*(site_off + C)); + GGML_ASSERT(seq_qnext != nullptr); + GGML_ASSERT(seq_qnext->type == GGML_TYPE_I32); + GGML_ASSERT(seq_qnext->ne[0] == 1); + GGML_ASSERT(seq_qnext->ne[1] == T); + + ggml_tensor * state_flat = ggml_view_2d(ctx, state_all, 2*C, 1, state_all->nb[1], + 2*site_off*ggml_element_size(state_all)); + ggml_tensor * state_in = reset_state ? ggml_scale(ctx, state_flat, 0.0f) : state_flat; + ggml_tensor * states = ggml_reshape_3d(ctx, state_in, 2, C, 1); + + ggml_tensor * conv_raw = ggml_ssm_conv(ctx, states, x, wc, seq_qnext, nullptr); + ggml_tensor * conv = ggml_view_2d(ctx, conv_raw, C, T, C*ggml_element_size(conv_raw), 0); + ggml_tensor * out = ggml_add(ctx, x, conv); + + ggml_tensor * new_states = ggml_view_2d(ctx, conv_raw, 2, C, + 3*ggml_element_size(conv_raw), + (1 + C*T)*ggml_element_size(conv_raw)); + ggml_tensor * new_states_cont = ggml_cont(ctx, new_states); + ggml_tensor * new_state_flat = ggml_reshape_2d(ctx, new_states_cont, 2*C, 1); + ggml_build_forward_expand(gf, ggml_cpy(ctx, new_state_flat, state_flat)); + return out; +} + +// --- mHC Sinkhorn: h_res [S*S, T] -> doubly-stochastic per token, 20 iters (ends on col norm) --- +static ggml_tensor * openpangu_sinkhorn(ggml_context * ctx, ggml_tensor * h_res_flat, + int64_t S, int64_t T, int iters, float hc_eps) { + // The flat h_res is torch [r,c] row-major (c fastest), so a bare reshape gives ne0=col. + // Transpose once so ne0=row(S), ne1=col(S): every axis op below then matches the + // reference _mhc_sinkhorn_naive (softmax over col, first norm over row, end on col-sum=1) + // and mhc_post's out[c] = sum_r m[r,c]*residual[r]. + (void) hc_eps; // softmax outputs are strictly positive, so the eps is numerically inert here + ggml_tensor * m = ggml_reshape_3d(ctx, h_res_flat, S, S, T); // ne0=col (bare reshape) + // ref softmaxes h_res over columns; a bare reshape already has ne0=col, so soft_max + // (over ne0) hits the column axis directly -- no pre-permute round-trip needed. + m = ggml_soft_max(ctx, m); // softmax over col + m = ggml_cont(ctx, ggml_permute(ctx, m, 1, 0, 2, 3)); // transpose once -> [row, col, T] + + auto col_norm = [&](ggml_tensor * a) { + ggml_tensor * col_sum = ggml_sum_rows(ctx, a); // sums ne0(row) -> [1, col, T] + return ggml_div(ctx, a, col_sum); // broadcast [1,col,T] over rows + }; + auto row_norm = [&](ggml_tensor * a) { + ggml_tensor * ap = ggml_cont(ctx, ggml_permute(ctx, a, 1, 0, 2, 3)); // [col,row,T] + ggml_tensor * row_sum = ggml_sum_rows(ctx, ap); // [1, row, T] + ggml_tensor * out = ggml_div(ctx, ap, row_sum); + return ggml_cont(ctx, ggml_permute(ctx, out, 1, 0, 2, 3)); // back [row,col,T] + }; + + m = col_norm(m); + for (int i = 0; i < iters - 1; ++i) { + m = row_norm(m); + m = col_norm(m); + } + return m; // [row(S), col(S), T] +} + +// Attention sublayer body, shared by the base layers and the NextN/MTP head. +// x_normed = input-layernormed hidden [n_embd, T]; returns post-o_proj output [n_embd, T]. +// conv_state is the recurrent MoME state slot. seq_qnext is the [1, T] sequence-id input +// used by ggml_ssm_conv and is shared by all three conv sites. +ggml_tensor * llm_build_context::build_openpangu_attention( + ggml_cgraph * gf, const llama_layer & layer, int il, ggml_tensor * x_normed, + ggml_tensor * KQ_mask, ggml_tensor * inp_pos, + ggml_tensor * conv_state, ggml_tensor * seq_qnext, + float kq_scale, bool KQ_mask_swa_windowed) { + const int64_t n_embd_head_qk_rope = hparams.n_rot; // 64 + const int64_t n_embd_head_k = hparams.n_embd_head_k(0); // 192 + const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope; // 128 + const int64_t n_embd_head_v = hparams.n_embd_head_v(0); // 128 + const int64_t kv_lora_rank = hparams.n_lora_kv; // 512 + const int64_t q_lora_rank = hparams.n_lora_q; // 1024 + + // MoME conv-state site offsets within one slot (channels): [qa | compresskv | o] + const int64_t conv_off_qa = 0; + const int64_t conv_off_ckv = q_lora_rank; + const int64_t conv_off_o = q_lora_rank + kv_lora_rank; + const bool reset_conv_state = batch.pos && batch.n_tokens > 0 && batch.pos[0] == 0; + + ggml_tensor * cur = x_normed; + + // --- Q path: q_a -> qa_conv -> q_a_norm -> q_b --- + ggml_tensor * q_lora = ggml_mul_mat(ctx0, layer.wq_a, cur); // [q_lora_rank, T] + q_lora = openpangu_causal_conv(ctx0, gf, q_lora, layer.qa_conv, conv_state, conv_off_qa, seq_qnext, reset_conv_state); + if (il == 0) ggml_set_name(q_lora, "opg0_qlora_conv"); + q_lora = llm_build_norm(ctx0, q_lora, hparams, layer.attn_q_a_norm, NULL, LLM_NORM_RMS, cb, il); + if (il == 0) ggml_set_name(q_lora, "opg0_qlora_norm"); + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, q_lora); // [n_head*192, T] + q = ggml_reshape_3d(ctx0, q, n_embd_head_k, 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); + ggml_tensor * q_rope = ggml_view_3d(ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, q->nb[1], q->nb[2], + n_embd_head_qk_nope*ggml_element_size(q)); + q_rope = ggml_rope_ext(ctx0, ggml_cont(ctx0, q_rope), inp_pos, nullptr, n_rot, rope_type, + n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + // absorbed (MLA-latent) queries: q_nope projected into the 512-latent through the + // load-derived wk_b, then ++ roped q_pe -> one [576, T, H] query against the latent cache + ggml_tensor * qn = ggml_cont(ctx0, ggml_permute(ctx0, ggml_cont(ctx0, q_nope), 0, 2, 1, 3)); // [128, T, H] + ggml_tensor * qp = ggml_cont(ctx0, ggml_permute(ctx0, q_rope, 0, 2, 1, 3)); // [64, T, H] + // wk_b loads 2D [128, H*512] from the converter split (head-major rows) or 3D from the + // load-time derivation; either way the data is [d, r, h] — reshape to the batched form + ggml_tensor * wk_b3 = ggml_reshape_3d(ctx0, layer.wk_b, n_embd_head_qk_nope, kv_lora_rank, n_head); + ggml_tensor * q_lat = ggml_mul_mat(ctx0, wk_b3, qn); // [512, T, H] + ggml_tensor * q_all = ggml_concat(ctx0, q_lat, qp, 0); // [576, T, H] + if (il == 0) ggml_set_name(q_all, "opg0_q_lat"); + + // --- KV path: kv_a -> split -> compresskv_conv -> kv_a_norm -> kv_b --- + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); // [kv_lora+64, T] + ggml_tensor * ckv = ggml_cont(ctx0, ggml_view_2d(ctx0, kv, kv_lora_rank, n_tokens, kv->nb[1], 0)); + ggml_tensor * k_pe = ggml_cont(ctx0, ggml_view_2d(ctx0, kv, n_embd_head_qk_rope, n_tokens, kv->nb[1], + kv_lora_rank*ggml_element_size(kv))); + ckv = openpangu_causal_conv(ctx0, gf, ckv, layer.kv_conv, conv_state, conv_off_ckv, seq_qnext, reset_conv_state); + ckv = llm_build_norm(ctx0, ckv, hparams, layer.attn_kv_a_norm, NULL, LLM_NORM_RMS, cb, il); + if (il == 0) ggml_set_name(ckv, "opg0_ckv_norm"); + // rope k_pe (shared across heads) + k_pe = ggml_reshape_3d(ctx0, k_pe, n_embd_head_qk_rope, 1, n_tokens); + 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); + ggml_tensor * k_pe2d = ggml_reshape_2d(ctx0, k_pe, n_embd_head_qk_rope, n_tokens); + + // ---- latent cache store: per position [ckv 512 | roped k_pe 64] straight into k_l. + // The value-side latent is rebuilt from k_l per graph; per-head K/V are never materialized. + { + ggml_tensor * kl = kv_self.k_l[il]; + ggml_tensor * ckv_store = ckv; + ggml_tensor * kpe_store = k_pe2d; + if (ggml_is_quantized(kl->type)) { + ckv_store = openpangu_cast_for_latent_cache_write(ctx0, ckv_store, kl); + kpe_store = openpangu_cast_for_latent_cache_write(ctx0, kpe_store, kl); + } + ggml_tensor * k_latent = ggml_concat(ctx0, ckv_store, kpe_store, 0); + ggml_tensor * kl_full = ggml_view_2d(ctx0, kl, kv_lora_rank + n_embd_head_qk_rope, + n_tokens, kl->nb[1], kv_head*kl->nb[1]); + ggml_tensor * cpy_kl = ggml_cpy(ctx0, k_latent, kl_full); + openpangu_register_cache_copy(lctx, il, OPENPANGU_COPY_K_CKV, cpy_kl, kl->nb[1]); + ggml_build_forward_expand(gf, cpy_kl); + } + + // ---- DSA lightning indexer: per-query top-k selection mask (DSA layers only) ---- + // Reference (Infer _pangu_torch_calib): q_idx = wq_b(q_lora_normed) [24 heads x 128], + // k_idx = rms(k_norm)(wk(x_normed)) [128, shared across heads], both NEOX-roped on the + // FIRST n_rot channels (opposite order vs the main head's [nope|rope] split); + // score[t,s] = sum_g weights_proj(x)[t,g] * relu(q_idx[t,g]·k_idx[s]) in f32, causal, + // then top-k. Selection only prunes when the causal window exceeds top_k; below that it + // covers everything, so the mask is skipped and the layer is exactly dense. + ggml_tensor * sel_mask = nullptr; // [n_kv, T] additive mask: 0 = selected, -1e30 = pruned + ggml_tensor * sel_idx = nullptr; // [topk, T] per-token top-k rows, reused by gathered DSA + bool dsa_gather_engaged = false; + int64_t dsa_topk = 0; + ggml_tensor * idx_cache = (size_t) il < kv_self.kr_l.size() ? kv_self.kr_l[il] : nullptr; + if (idx_cache && layer.indexer_attn_q_b) { + const int64_t n_ihead = hparams.indexer_n_head; // 24 + const int64_t d_idx = hparams.indexer_head_size; // 128 + int64_t topk = hparams.indexer_top_k; // 2048 + if (lctx.cparams.dsa_top_k > 0) { + topk = lctx.cparams.dsa_top_k; + } + GGML_ASSERT(topk > 0 && topk <= INT_MAX); + dsa_topk = topk; + + const uint32_t pad = openpangu_kv_cache_pad(cparams); + const bool is_base_graph = cparams.mtp_op_type == MTP_OP_NONE; + const bool dsa_gather_predicate = + openpangu_dsa_gather_should_engage(n_kv, n_tokens, topk, pad); + const bool dsa_gather_allowed = + is_base_graph && dsa_gather_predicate && + openpangu_dsa_gather_rows_fit_cuda(topk, n_tokens); + + // indexer keys for this batch -> position-indexed cache (write-before-read holds by + // graph order, same as the kv store; committed columns never change -> rollback-safe) + ggml_tensor * k_idx = ggml_mul_mat(ctx0, layer.indexer_attn_k, x_normed); // [d_idx, T] + k_idx = llm_build_norm(ctx0, k_idx, hparams, layer.indexer_k_norm, layer.indexer_k_norm_b, + layer.indexer_k_norm_b ? LLM_NORM : LLM_NORM_RMS, cb, il); + ggml_tensor * k_idx_rope = ggml_view_2d(ctx0, k_idx, n_embd_head_qk_rope, n_tokens, k_idx->nb[1], 0); + ggml_tensor * k_idx_pass = ggml_view_2d(ctx0, k_idx, d_idx - n_embd_head_qk_rope, n_tokens, k_idx->nb[1], + n_embd_head_qk_rope*ggml_element_size(k_idx)); + k_idx_rope = ggml_rope_ext(ctx0, ggml_reshape_3d(ctx0, ggml_cont(ctx0, k_idx_rope), n_embd_head_qk_rope, 1, n_tokens), + inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + k_idx = ggml_concat(ctx0, ggml_reshape_2d(ctx0, k_idx_rope, n_embd_head_qk_rope, n_tokens), + ggml_cont(ctx0, k_idx_pass), 0); // [d_idx, T] + if (il == 0) ggml_set_name(k_idx, "opg0_idx_k"); + ggml_tensor * idx_w = ggml_view_2d(ctx0, idx_cache, d_idx, n_tokens, idx_cache->nb[1], kv_head*idx_cache->nb[1]); + ggml_tensor * cpy_idx = ggml_cpy(ctx0, k_idx, idx_w); + if ((size_t) il < lctx.dsa_cache_copies.size()) { + lctx.dsa_cache_copies[il].cpy = cpy_idx; + lctx.dsa_cache_copies[il].step = idx_cache->nb[1]; + } + ggml_build_forward_expand(gf, cpy_idx); + + if (n_kv > topk) { + // indexer queries + ggml_tensor * q_idx = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, q_lora); // [n_ihead*d_idx, T] + q_idx = ggml_reshape_3d(ctx0, q_idx, d_idx, n_ihead, n_tokens); + ggml_tensor * q_idx_rope = ggml_view_3d(ctx0, q_idx, n_embd_head_qk_rope, n_ihead, n_tokens, + q_idx->nb[1], q_idx->nb[2], 0); + ggml_tensor * q_idx_pass = ggml_view_3d(ctx0, q_idx, d_idx - n_embd_head_qk_rope, n_ihead, n_tokens, + q_idx->nb[1], q_idx->nb[2], n_embd_head_qk_rope*ggml_element_size(q_idx)); + q_idx_rope = ggml_rope_ext(ctx0, ggml_cont(ctx0, q_idx_rope), inp_pos, nullptr, n_rot, rope_type, + n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + q_idx = ggml_concat(ctx0, q_idx_rope, ggml_cont(ctx0, q_idx_pass), 0); // [d_idx, n_ihead, T] + if (il == 0) ggml_set_name(q_idx, "opg0_idx_q"); + + ggml_tensor * k_all_idx = ggml_view_2d(ctx0, idx_cache, d_idx, n_kv, idx_cache->nb[1], 0); + ggml_tensor * w_idx = ggml_mul_mat(ctx0, layer.indexer_proj, x_normed); // [n_ihead, T] + + const bool chunk_scores = openpangu_idx_score_should_chunk(n_tokens, OPENPANGU_IDX_SCORE_CHUNK); + const bool defer_sel_mask_to_att_chunks = + !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) { + 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); + ggml_tensor * q_idx_c = ggml_view_3d(ctx0, q_idx, d_idx, n_ihead, tc, + q_idx->nb[1], q_idx->nb[2], (size_t) c0*q_idx->nb[2]); + q_idx_c = ggml_cont(ctx0, q_idx_c); + ggml_tensor * w_idx_c = ggml_view_2d(ctx0, w_idx, n_ihead, tc, w_idx->nb[1], (size_t) c0*w_idx->nb[1]); + w_idx_c = ggml_cont(ctx0, w_idx_c); + + // Scores over the whole cache window, causal-masked before selection. + // The chunked path never materializes the legacy [n_kv, n_ihead, T] score tensor. + ggml_tensor * sc_c = ggml_mul_mat(ctx0, k_all_idx, q_idx_c); // [n_kv, n_ihead, Tc] + sc_c = ggml_relu(ctx0, sc_c); + sc_c = ggml_mul(ctx0, sc_c, ggml_reshape_3d(ctx0, w_idx_c, 1, n_ihead, tc)); + sc_c = ggml_cont(ctx0, ggml_permute(ctx0, sc_c, 1, 0, 2, 3)); // [n_ihead, n_kv, Tc] + sc_c = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, sc_c), n_kv, tc); // [n_kv, Tc] + + sc_c = ggml_add(ctx0, sc_c, ggml_cont(ctx0, ggml_view_2d(ctx0, KQ_mask, n_kv, tc, + KQ_mask->nb[1], (size_t) c0*KQ_mask->nb[1]))); + ggml_tensor * sel_idx_c = ggml_top_k(ctx0, sc_c, (int) topk); // [topk, Tc] i32 + sel_idx = sel_idx == nullptr ? sel_idx_c : ggml_concat(ctx0, sel_idx, sel_idx_c, 1); + + if (!dsa_gather_allowed && !defer_sel_mask_to_att_chunks) { + ggml_tensor * base_c = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv, tc); + base_c = ggml_fill(ctx0, base_c, -1e30f); + ggml_tensor * zeros_c = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, topk, tc); + zeros_c = ggml_fill(ctx0, zeros_c, 0.0f); + ggml_tensor * sel_mask_c = ggml_set_rows(ctx0, base_c, zeros_c, sel_idx_c); + sel_mask_c = ggml_reshape_2d(ctx0, sel_mask_c, n_kv, tc); + sel_mask_parts = sel_mask_parts == nullptr ? sel_mask_c : ggml_concat(ctx0, sel_mask_parts, sel_mask_c, 1); + } + } + 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 { + sel_mask = sel_mask_parts; + } + } else { + GGML_ASSERT(n_tokens <= 14 || OPENPANGU_IDX_SCORE_CHUNK == 0 || n_tokens <= OPENPANGU_IDX_SCORE_CHUNK); + ggml_tensor * sc = ggml_mul_mat(ctx0, k_all_idx, q_idx); // [n_kv, n_ihead, T] + sc = ggml_relu(ctx0, sc); + sc = ggml_mul(ctx0, sc, ggml_reshape_3d(ctx0, w_idx, 1, n_ihead, n_tokens)); + sc = ggml_cont(ctx0, ggml_permute(ctx0, sc, 1, 0, 2, 3)); // [n_ihead, n_kv, T] + sc = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, sc), n_kv, n_tokens); // [n_kv, T] + sc = ggml_add(ctx0, sc, ggml_cont(ctx0, ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0))); + if (il == 0) ggml_set_name(sc, "opg0_idx_scores"); + + // Exact top-k -> additive mask: scatter zeros into a -1e30 base at the selected + // positions ([1, n_kv, T] row layout makes set_rows a per-query scatter). + // The legacy scatter path consumes this strided top_k view directly; the gathered + // path below flattens a tiny contiguous copy for one get_rows. + sel_idx = ggml_top_k(ctx0, sc, (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 { + 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); + } + } + } + } + + // ---- param_sink: 128 learned latent-KV entries prepended to the sequence ---- + // The sinks are native latent-space entries: kv_a_norm of the learned latent IS the + // key/value latent; sink_k_pe is used rope-free. Sinks are visible to every query. + const int64_t NS = hparams.param_sink_number; + GGML_ASSERT(layer.param_sink_blk && layer.param_sink_lat_t); + ggml_tensor * sink_blk = layer.param_sink_blk; // [576, NS] + ggml_tensor * s_lat_t = layer.param_sink_lat_t; // [NS, 512] + + // ---- latent attention over [sinks ++ cached tokens] (flash_attn is forced off) ---- + // Keep sinks and cached tokens separate until after KQ so f16 latent caches do not need + // unsupported non-f32 concat along dim1. + const bool use_swa_window = KQ_mask_swa_windowed && lctx.openpangu_swa_window_view.active; + const bool use_dsa_gather = dsa_gather_engaged && sel_idx != nullptr; + const int64_t n_kv_attn = use_dsa_gather ? dsa_topk : + use_swa_window ? lctx.openpangu_swa_window_view.w_view : n_kv; + const int64_t win_off = use_swa_window ? lctx.openpangu_swa_window_view.win_off : 0; + if (sel_mask) { + GGML_ASSERT(!use_swa_window && "openPangu DSA/indexer layers must not use SWA window views"); + } + if (use_dsa_gather) { + GGML_ASSERT(!use_swa_window && "openPangu gathered DSA layers must not use SWA window views"); + GGML_ASSERT(dsa_topk > 0 && n_kv_attn == dsa_topk); + GGML_ASSERT(hparams.f_max_alibi_bias == 0.0f && "maskless gathered DSA cannot carry ALiBi bias"); + } + + ggml_tensor * kl_all = nullptr; + ggml_tensor * vl_all = nullptr; + if (!use_dsa_gather) { + kl_all = openpangu_build_k_latent_for_read(ctx0, kv_self, il, n_kv_attn, win_off); + } + auto get_vl_all = [&]() -> ggml_tensor * { + if (vl_all == nullptr) { + vl_all = openpangu_build_v_latent_from_k(ctx0, kv_self, il, kv_lora_rank, n_kv_attn, win_off); + } + return vl_all; + }; + + ggml_tensor * k_gath = nullptr; + ggml_tensor * kq_cache = nullptr; + if (use_dsa_gather) { + ggml_tensor * kq_sinks = ggml_mul_mat(ctx0, sink_blk, q_all); // [NS, T, H] + if (n_tokens == 1) { + ggml_tensor * kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], kv_lora_rank + n_embd_head_qk_rope, n_kv, + kv_self.k_l[il]->nb[1], 0); + ggml_tensor * sel_idx_flat = ggml_cont_2d(ctx0, sel_idx, dsa_topk, 1); // [topk] i32 + k_gath = ggml_get_rows(ctx0, kl_full, sel_idx_flat); // [576, topk] f32 + k_gath = openpangu_cast_gathered_latent_for_cache_type(ctx0, k_gath, kv_self.k_l[il]); + k_gath = ggml_reshape_3d(ctx0, k_gath, kv_lora_rank + n_embd_head_qk_rope, dsa_topk, 1); + kq_cache = ggml_mul_mat(ctx0, + ggml_reshape_2d(ctx0, k_gath, kv_lora_rank + n_embd_head_qk_rope, dsa_topk), + q_all); // [topk, 1, H] + } else { + ggml_tensor * kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], + kv_lora_rank + n_embd_head_qk_rope, n_kv, + kv_self.k_l[il]->nb[1], 0); + ggml_tensor * sel_idx_flat = ggml_cont_2d(ctx0, sel_idx, dsa_topk*n_tokens, 1); // [topk*T] i32 + k_gath = ggml_get_rows(ctx0, kl_full, sel_idx_flat); // [576, topk*T] f32 + k_gath = openpangu_cast_gathered_latent_for_cache_type(ctx0, k_gath, kv_self.k_l[il]); + k_gath = ggml_reshape_3d(ctx0, k_gath, kv_lora_rank + n_embd_head_qk_rope, + dsa_topk, n_tokens); // [576, topk, T] + ggml_tensor * q_gath = ggml_cont(ctx0, ggml_permute(ctx0, q_all, 0, 2, 1, 3)); // [576, H, T] + kq_cache = ggml_mul_mat(ctx0, k_gath, q_gath); // [topk, H, T] + kq_cache = ggml_cont(ctx0, ggml_permute(ctx0, kq_cache, 0, 2, 1, 3)); // [topk, T, H] + } + ggml_tensor * kq = ggml_concat(ctx0, kq_sinks, kq_cache, 0); // [NS+topk, T, H] + // sel_idx came from scores after adding the causal KQ_mask. The engagement bound gives + // every token at least topk valid positions, so all gathered rows are visible here. + kq = ggml_soft_max_ext(ctx0, kq, nullptr, kq_scale, hparams.f_max_alibi_bias); + + ggml_tensor * kq_s = ggml_view_3d(ctx0, kq, NS, n_tokens, n_head, kq->nb[1], kq->nb[2], 0); + ggml_tensor * kq_c = ggml_view_3d(ctx0, kq, n_kv_attn, n_tokens, n_head, kq->nb[1], kq->nb[2], + NS*ggml_element_size(kq)); + ggml_tensor * kqv_cache = nullptr; + GGML_ASSERT(k_gath != nullptr); + if (n_tokens == 1) { + ggml_tensor * v_gath = ggml_view_2d(ctx0, k_gath, kv_lora_rank, dsa_topk, + k_gath->nb[1], 0); // [512, topk] + ggml_tensor * v_gath_t = ggml_cont(ctx0, ggml_transpose(ctx0, v_gath)); // [topk, 512] + kqv_cache = ggml_mul_mat(ctx0, v_gath_t, kq_c); // [512, 1, H] + } else { + ggml_tensor * v_gath = ggml_view_3d(ctx0, k_gath, kv_lora_rank, dsa_topk, n_tokens, + k_gath->nb[1], k_gath->nb[2], 0); // [512, topk, T] + ggml_tensor * v_gath_t = ggml_cont(ctx0, ggml_permute(ctx0, v_gath, 1, 0, 2, 3)); // [topk, 512, T] + ggml_tensor * kq_c_gath = ggml_cont(ctx0, ggml_permute(ctx0, kq_c, 0, 2, 1, 3)); // [topk, H, T] + kqv_cache = ggml_mul_mat(ctx0, v_gath_t, kq_c_gath); // [512, H, T] + kqv_cache = ggml_cont(ctx0, ggml_permute(ctx0, kqv_cache, 0, 2, 1, 3)); // [512, T, H] + } + ggml_tensor * kqv = ggml_add(ctx0, + ggml_mul_mat(ctx0, s_lat_t, kq_s), + kqv_cache); // [512, T, H] + ggml_tensor * wv_b3 = ggml_reshape_3d(ctx0, layer.wv_b, kv_lora_rank, n_embd_head_v, n_head); + ggml_tensor * out_h = ggml_mul_mat(ctx0, wv_b3, kqv); // [128, T, H] + ggml_tensor * merged = ggml_cont(ctx0, ggml_permute(ctx0, out_h, 0, 2, 1, 3)); // [128, H, T] + cur = ggml_reshape_2d(ctx0, merged, n_embd_head_v * n_head, n_tokens); + } else { + const bool chunk_att = openpangu_att_score_should_chunk(n_kv_attn, NS, n_head, n_tokens, + OPENPANGU_ATT_SCORE_CHUNK, OPENPANGU_ATT_FULL_KQ_MAX_MIB); + const bool use_dsa_sel_idx_mask = sel_idx != nullptr && dsa_topk > 0; + ggml_tensor * kqv = nullptr; + if (chunk_att) { + const bool can_prefill_gather = + use_dsa_sel_idx_mask && !use_swa_window && + hparams.f_max_alibi_bias == 0.0f && + openpangu_dsa_gather_rows_fit_cuda(dsa_topk, 1); + ggml_tensor * kl_full = nullptr; + if (can_prefill_gather) { + kl_full = ggml_view_2d(ctx0, kv_self.k_l[il], kv_lora_rank + n_embd_head_qk_rope, n_kv, + kv_self.k_l[il]->nb[1], 0); + } + + for (int64_t c0 = 0; c0 < n_tokens; c0 += OPENPANGU_ATT_SCORE_CHUNK) { + const int64_t tc = std::min(OPENPANGU_ATT_SCORE_CHUNK, n_tokens - c0); + ggml_tensor * q_all_c = ggml_view_3d(ctx0, q_all, kv_lora_rank + n_embd_head_qk_rope, + tc, n_head, q_all->nb[1], q_all->nb[2], + (size_t) c0*q_all->nb[1]); + q_all_c = ggml_cont(ctx0, q_all_c); + + ggml_tensor * kq_sinks_c = ggml_mul_mat(ctx0, sink_blk, q_all_c); // [NS, Tc, H] + + const bool prefill_gather_chunk = + can_prefill_gather && + openpangu_dsa_prefill_gather_should_engage(n_kv, n_tokens, c0, tc, dsa_topk, + openpangu_kv_cache_pad(cparams)); + ggml_tensor * kqv_c = nullptr; + if (prefill_gather_chunk) { + const int64_t gather_token_chunk = openpangu_dsa_gather_tokens_per_get_rows(dsa_topk); + for (int64_t g0 = 0; g0 < tc; g0 += gather_token_chunk) { + const int64_t tg = std::min(gather_token_chunk, tc - g0); + ggml_tensor * q_all_g = ggml_view_3d(ctx0, q_all_c, + kv_lora_rank + n_embd_head_qk_rope, + tg, n_head, q_all_c->nb[1], q_all_c->nb[2], + (size_t) g0*q_all_c->nb[1]); + q_all_g = ggml_cont(ctx0, q_all_g); + ggml_tensor * kq_sinks_g = ggml_view_3d(ctx0, kq_sinks_c, NS, tg, n_head, + kq_sinks_c->nb[1], kq_sinks_c->nb[2], + (size_t) g0*kq_sinks_c->nb[1]); + ggml_tensor * sel_idx_g = ggml_view_2d(ctx0, sel_idx, dsa_topk, tg, + sel_idx->nb[1], (size_t) (c0 + g0)*sel_idx->nb[1]); + ggml_tensor * sel_idx_flat_g = ggml_cont_2d(ctx0, sel_idx_g, dsa_topk*tg, 1); + ggml_tensor * k_gath_g = ggml_get_rows(ctx0, kl_full, sel_idx_flat_g); // [576, topk*tg] + k_gath_g = openpangu_cast_gathered_latent_for_cache_type(ctx0, k_gath_g, kv_self.k_l[il]); + k_gath_g = ggml_reshape_3d(ctx0, k_gath_g, kv_lora_rank + n_embd_head_qk_rope, + dsa_topk, tg); // [576, topk, tg] + ggml_tensor * q_gath_g = ggml_cont(ctx0, ggml_permute(ctx0, q_all_g, 0, 2, 1, 3)); // [576, H, tg] + ggml_tensor * kq_cache_g = ggml_mul_mat(ctx0, k_gath_g, q_gath_g); // [topk, H, tg] + kq_cache_g = ggml_cont(ctx0, ggml_permute(ctx0, kq_cache_g, 0, 2, 1, 3)); // [topk, tg, H] + ggml_tensor * kq_g_all = ggml_concat(ctx0, kq_sinks_g, kq_cache_g, 0); // [NS+topk, tg, H] + kq_g_all = ggml_soft_max_ext(ctx0, kq_g_all, nullptr, kq_scale, hparams.f_max_alibi_bias); + + ggml_tensor * kq_s_g = ggml_view_3d(ctx0, kq_g_all, NS, tg, n_head, + kq_g_all->nb[1], kq_g_all->nb[2], 0); + ggml_tensor * kq_cache_soft_g = ggml_view_3d(ctx0, kq_g_all, dsa_topk, tg, n_head, + kq_g_all->nb[1], kq_g_all->nb[2], + NS*ggml_element_size(kq_g_all)); + ggml_tensor * v_gath_g = ggml_view_3d(ctx0, k_gath_g, kv_lora_rank, dsa_topk, tg, + k_gath_g->nb[1], k_gath_g->nb[2], 0); + ggml_tensor * v_gath_t_g = ggml_cont(ctx0, ggml_permute(ctx0, v_gath_g, 1, 0, 2, 3)); // [topk, 512, tg] + ggml_tensor * kq_cache_gath_g = ggml_cont(ctx0, ggml_permute(ctx0, kq_cache_soft_g, 0, 2, 1, 3)); // [topk, H, tg] + ggml_tensor * kqv_cache_g = ggml_mul_mat(ctx0, v_gath_t_g, kq_cache_gath_g); // [512, H, tg] + kqv_cache_g = ggml_cont(ctx0, ggml_permute(ctx0, kqv_cache_g, 0, 2, 1, 3)); // [512, tg, H] + ggml_tensor * kqv_g = ggml_add(ctx0, + ggml_mul_mat(ctx0, s_lat_t, kq_s_g), + kqv_cache_g); // [512, tg, H] + kqv_c = kqv_c == nullptr ? kqv_g : ggml_concat(ctx0, kqv_c, kqv_g, 1); + } + } else { + ggml_tensor * kq_cache_c = ggml_mul_mat(ctx0, kl_all, q_all_c); // [n_kv_attn, Tc, H] + ggml_tensor * kq_c_all = ggml_concat(ctx0, kq_sinks_c, kq_cache_c, 0); // [NS+n_kv_attn, Tc, H] + + ggml_tensor * kq_mask_eff = ggml_view_2d(ctx0, KQ_mask, n_kv_attn, tc, + KQ_mask->nb[1], (size_t) c0*KQ_mask->nb[1]); + if (sel_mask) { + ggml_tensor * sel_mask_c = ggml_view_2d(ctx0, sel_mask, n_kv_attn, tc, + sel_mask->nb[1], (size_t) c0*sel_mask->nb[1]); + kq_mask_eff = ggml_add(ctx0, kq_mask_eff, sel_mask_c); + } else if (use_dsa_sel_idx_mask) { + ggml_tensor * sel_idx_c = ggml_view_2d(ctx0, sel_idx, dsa_topk, tc, + sel_idx->nb[1], (size_t) c0*sel_idx->nb[1]); + ggml_tensor * base_src_c = ggml_view_2d(ctx0, kq_cache_c, n_kv_attn, tc, kq_cache_c->nb[1], 0); + ggml_tensor * base_c = ggml_scale_bias(ctx0, base_src_c, 0.0f, -1e30f); + base_c = ggml_reshape_3d(ctx0, base_c, 1, n_kv_attn, tc); + ggml_tensor * zeros_src_c = ggml_view_2d(ctx0, kq_cache_c, dsa_topk, tc, + dsa_topk*ggml_element_size(kq_cache_c), 0); + ggml_tensor * zeros_c = ggml_scale(ctx0, zeros_src_c, 0.0f); + zeros_c = ggml_reshape_3d(ctx0, zeros_c, 1, dsa_topk, tc); + ggml_tensor * sel_mask_c = ggml_set_rows(ctx0, base_c, zeros_c, sel_idx_c); + sel_mask_c = ggml_reshape_2d(ctx0, sel_mask_c, n_kv_attn, tc); + kq_mask_eff = ggml_add(ctx0, kq_mask_eff, sel_mask_c); + } + kq_mask_eff = ggml_cont(ctx0, kq_mask_eff); + ggml_tensor * s_mask0 = ggml_scale(ctx0, + ggml_view_2d(ctx0, kq_c_all, NS, tc, NS*ggml_element_size(kq_c_all), 0), 0.0f); + ggml_tensor * mask_all = ggml_concat(ctx0, s_mask0, kq_mask_eff, 0); // [NS+n_kv_attn, Tc] + kq_c_all = ggml_soft_max_ext(ctx0, kq_c_all, mask_all, kq_scale, hparams.f_max_alibi_bias); + + ggml_tensor * kq_s_c = ggml_view_3d(ctx0, kq_c_all, NS, tc, n_head, + kq_c_all->nb[1], kq_c_all->nb[2], 0); + ggml_tensor * kq_cache_soft_c = ggml_view_3d(ctx0, kq_c_all, n_kv_attn, tc, n_head, + kq_c_all->nb[1], kq_c_all->nb[2], + NS*ggml_element_size(kq_c_all)); + ggml_tensor * kqv_cache_c = ggml_mul_mat(ctx0, get_vl_all(), kq_cache_soft_c); // [512, Tc, H] + kqv_c = ggml_add(ctx0, + ggml_mul_mat(ctx0, s_lat_t, kq_s_c), + kqv_cache_c); // [512, Tc, H] + } + kqv = kqv == nullptr ? kqv_c : ggml_concat(ctx0, kqv, kqv_c, 1); + } + } else { + GGML_ASSERT(!use_dsa_sel_idx_mask || sel_mask != nullptr); + ggml_tensor * kq_sinks = ggml_mul_mat(ctx0, sink_blk, q_all); // [NS, T, H] + kq_cache = ggml_mul_mat(ctx0, kl_all, q_all); // [n_kv_attn, T, H] + ggml_tensor * kq = ggml_concat(ctx0, kq_sinks, kq_cache, 0); // [NS+n_kv_attn, T, H] + + // mask: sinks always visible (0) ++ the causal/SWA KQ_mask (+ the DSA top-k selection + // mask on indexer layers). The zero block is built by scaling finite kq data (KQ_mask + // itself holds -inf, which 0*x would turn into NaN). + ggml_tensor * kq_mask_eff = ggml_view_2d(ctx0, KQ_mask, n_kv_attn, n_tokens, KQ_mask->nb[1], 0); + if (sel_mask) { + kq_mask_eff = ggml_add(ctx0, kq_mask_eff, sel_mask); + } + kq_mask_eff = ggml_cont(ctx0, kq_mask_eff); + ggml_tensor * s_mask0 = ggml_scale(ctx0, ggml_view_2d(ctx0, kq, NS, n_tokens, NS*ggml_element_size(kq), 0), 0.0f); + ggml_tensor * mask_all = ggml_concat(ctx0, s_mask0, kq_mask_eff, 0); // [NS+n_kv, T] + kq = ggml_soft_max_ext(ctx0, kq, mask_all, kq_scale, hparams.f_max_alibi_bias); + + ggml_tensor * kq_s = ggml_view_3d(ctx0, kq, NS, n_tokens, n_head, kq->nb[1], kq->nb[2], 0); + ggml_tensor * kq_c = ggml_view_3d(ctx0, kq, n_kv_attn, n_tokens, n_head, kq->nb[1], kq->nb[2], + NS*ggml_element_size(kq)); + ggml_tensor * kqv_cache = ggml_mul_mat(ctx0, get_vl_all(), kq_c); // [512, T, H] + kqv = ggml_add(ctx0, + ggml_mul_mat(ctx0, s_lat_t, kq_s), + kqv_cache); // [512, T, H] + } + + ggml_tensor * wv_b3 = ggml_reshape_3d(ctx0, layer.wv_b, kv_lora_rank, n_embd_head_v, n_head); + ggml_tensor * out_h = ggml_mul_mat(ctx0, wv_b3, kqv); // [128, T, H] + ggml_tensor * merged = ggml_cont(ctx0, ggml_permute(ctx0, out_h, 0, 2, 1, 3)); // [128, H, T] + cur = ggml_reshape_2d(ctx0, merged, n_embd_head_v * n_head, n_tokens); + } + + // o_conv (MOME on the pre-o_proj attn output), then o_proj + cur = openpangu_causal_conv(ctx0, gf, cur, layer.o_conv, conv_state, conv_off_o, seq_qnext, reset_conv_state); + cur = llm_build_lora_mm(lctx, ctx0, layer.wo, cur); + return cur; +} + +// NextN/MTP head: eh_proj stitching -> one plain-residual Pangu block (sandwich norms, +// NO mHC, no block_post_norm) -> shared head. Mirrors OpenPanguV2MultiTokenPredictorLayer: +// x = eh_proj(cat(enorm(embed(tok)), hnorm(prev_hidden))) +// x = x + post_attn_ln(attn(input_ln(x))) +// x = x + post_mlp_ln(moe(pre_mlp_ln(x))) +// logits = shared_head.head(shared_head.norm(x)) +// The MTP context allocates recurrent conv-state slots for the NextN layers too, so draft +// convs chain real t-1/t-2 taps across warmup and sequential draft steps. +ggml_tensor * llm_build_context::build_openpangu_mtp( + const llama_layer & mtp_layer, + ggml_tensor * prev_embeddings, + ggml_cgraph * gf, + int il, + ggml_tensor * inp_pos, + ggml_tensor * KQ_mask, + ggml_tensor * inp_out_ids, + ggml_tensor * inp_tokens, + ggml_tensor * seq_qnext, + ggml_tensor ** full_hidden_out, + bool select_outputs, + bool build_logits, + bool cache_writes_only, + bool KQ_mask_swa_windowed) { + const float kq_scale = 1.0f / sqrtf(float(hparams.n_embd_head_k(0))); + + // same position-addressing invariant as build_openpangu (worst-case builds exempt) + if (batch.pos && batch.n_tokens > 0) { + GGML_ASSERT((llama_pos) kv_head == batch.pos[0] && + "openPangu KV cache is position-addressed; kv head must equal the first batch position"); + } + + // the batch inputs (tokens, positions, masks, output selection, conv sequence ids) are + // created ONCE by the caller and shared by every head built into the graph: + // llama_set_inputs fills the tensors the lctx.inp_* pointers reference, so per-head + // creation would leave every head but the last reading unwritten memory + ggml_tensor * mtp_embd_weights = mtp_layer.nextn.embed_tokens + ? mtp_layer.nextn.embed_tokens : model.tok_embd; + ggml_tensor * token_emb = ggml_get_rows(ctx0, mtp_embd_weights, inp_tokens); + cb(token_emb, "inp_embd", il); + + ggml_tensor * emb_norm = llm_build_norm(ctx0, token_emb, hparams, mtp_layer.nextn.enorm, NULL, LLM_NORM_RMS, cb, il); + ggml_tensor * hid_norm = llm_build_norm(ctx0, prev_embeddings, hparams, mtp_layer.nextn.hnorm, NULL, LLM_NORM_RMS, cb, il); + + // reference order: cat([inputs_embeds, previous_hidden_states], -1) + ggml_tensor * combined = ggml_concat(ctx0, emb_norm, hid_norm, 0); + ggml_tensor * cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined); + cb(cur, "mtp_eh_proj", il); + + // --- attention sublayer (plain residual) --- + ggml_tensor * inpSA = cur; + cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.attn_norm, NULL, LLM_NORM_RMS, cb, il); + // the MTP context allocates a recurrent conv-state slot for the NextN layers, so the + // draft head chains real t-1/t-2 taps across warmup and sequential draft steps + ggml_tensor * mtp_conv_state = (size_t) il < kv_self.s_l.size() ? kv_self.s_l[il] : nullptr; + cur = build_openpangu_attention(gf, mtp_layer, il, cur, KQ_mask, inp_pos, + mtp_conv_state, seq_qnext, kq_scale, KQ_mask_swa_windowed); + if (cache_writes_only) { + // only this head's latent-cache and conv-slot writes matter at this site (the + // update chain's last head and the draft-time row fill); the FFN, norms, and + // shared head would compute values nobody consumes + cb(cur, "mtp_cache_write_anchor", il); + return cur; + } + cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.attn_post_norm, NULL, LLM_NORM_RMS, cb, il); + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "mtp_ffn_inp", il); + + const bool keep_full_hidden = full_hidden_out != nullptr; + if (select_outputs && inp_out_ids && !keep_full_hidden) { + ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids); + } + + // --- ffn sublayer: MoE (routed + shared expert), plain residual --- + cur = llm_build_norm(ctx0, cur = ffn_inp, hparams, mtp_layer.ffn_norm, NULL, LLM_NORM_RMS, cb, il); + { + ggml_tensor * moe_out = llm_build_moe_ffn(ctx0, lctx, cur, + mtp_layer.ffn_gate_inp, mtp_layer.ffn_up_exps, mtp_layer.ffn_gate_exps, mtp_layer.ffn_down_exps, + mtp_layer.ffn_exp_probs_b, n_expert, n_expert_used, LLM_FFN_SILU, + hparams.expert_weights_norm, true, hparams.expert_weights_scale, + (enum llm_expert_gating_func_type) hparams.expert_gating_func, cb, il, gf, false); + ggml_tensor * shexp = llm_build_ffn(ctx0, lctx, nullptr, cur, + mtp_layer.ffn_up_shexp, NULL, NULL, mtp_layer.ffn_gate_shexp, NULL, NULL, + mtp_layer.ffn_down_shexp, NULL, NULL, NULL, LLM_FFN_SILU, LLM_FFN_PAR, cb, il); + cur = ggml_add(ctx0, moe_out, shexp); + } + cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.ffn_post_norm, NULL, LLM_NORM_RMS, cb, il); + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "mtp_out_resid", il); + + // --- shared head --- + cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, -1); + if (full_hidden_out) { + *full_hidden_out = cur; + } + if (select_outputs && inp_out_ids && keep_full_hidden) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + if (!build_logits) { + cb(cur, "mtp_hidden", il); + return cur; + } + cb(cur, "result_norm", -1); + ggml_tensor * head = mtp_layer.nextn.shared_head_head + ? mtp_layer.nextn.shared_head_head : model.output; + cur = llm_build_lora_mm(lctx, ctx0, head, cur); + cb(cur, "result_output", -1); + return cur; +} + +ggml_cgraph * llm_build_context::build_openpangu() { + ggml_cgraph * gf = new_graph_custom(); + openpangu_clear_cache_copies(lctx); + + // the indexer and latent stores are addressed by absolute position through kv_head; + // enforce head == first batch position on real builds so any future cache + // plumbing that breaks the append-only invariant fails here instead of corrupting + // (worst-case measurement builds pass pos = null and are exempt) + if (batch.pos && batch.n_tokens > 0) { + GGML_ASSERT((llama_pos) kv_head == batch.pos[0] && + "openPangu KV cache is position-addressed; kv head must equal the first batch position"); + } + + const int64_t n_embd_head_k = hparams.n_embd_head_k(0); // 192 + const int64_t S = hparams.mhc_num_stream; // 4 + const int sink_iters = (int) hparams.mhc_recur_norm; // 20 + const float hc_eps = 1e-6f; + const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k)); + + + // NextN/MTP graph (speculative decoding): draft generation selects the + // requested head by depth; warmup/update chains all active heads so their conv slots and + // latent caches hold real committed rows, exposing head 1 as the one-token shortcut. + // Chaining convention (mirrors how head 1 consumes the target's shifted hidden rows): + // head k+1's row at position p consumes head k's output row at position p-1; the p-1 of + // the first batch row lives in the previous warmup/update batch and crosses decodes + // through the inp_mtp_carry input / lctx.mtp_carry storage. + if (cparams.mtp_op_type != MTP_OP_NONE) { + GGML_ASSERT(model.mtp && hparams.nextn_predict_layers > 0 && + "OpenPangu MTP graph requested without NextN layers loaded"); + GGML_ASSERT(batch.token && "openPangu MTP graphs decode token batches"); + + ggml_tensor * hidden_states_from_main_model; + if (cparams.mtp_op_type == MTP_OP_WARMUP || cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) { + hidden_states_from_main_model = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + } else { + hidden_states_from_main_model = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hparams.n_embd); + } + ggml_set_name(hidden_states_from_main_model, "inp_mtp_states"); + ggml_set_input(hidden_states_from_main_model); + lctx.inp_mtp_states = hidden_states_from_main_model; + + // shared batch inputs, created exactly once per graph (see build_openpangu_mtp) + ggml_tensor * inp_pos = build_inp_pos(); + // the NextN/MTP layers are SWA layers with their own window (2048); the mask fill + // uses hparams.n_swa_mtp when the graph is built with an MTP op type + bool KQ_mask_swa_windowed = false; + ggml_tensor * KQ_mask = hparams.n_swa_mtp > 0 && hparams.n_swa > 0 + ? openpangu_build_swa_mask_for_graph(*this, hparams.n_swa_mtp, &KQ_mask_swa_windowed) + : build_inp_KQ_mask(); + ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; + lctx.inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch.n_tokens); + cb(lctx.inp_tokens, "inp_tokens", -1); + ggml_set_input(lctx.inp_tokens); + ggml_tensor * inp_tokens = lctx.inp_tokens; + lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens); + cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1); + ggml_set_input(lctx.inp_s_seq_qnext); + ggml_tensor * seq_qnext = lctx.inp_s_seq_qnext; + + const int il_mtp_first = (int) (hparams.n_layer - hparams.nextn_predict_layers); + const int n_mtp_heads_model = (int) hparams.nextn_predict_layers; + const int n_mtp_heads = lctx.mtp_n_heads > 0 + ? std::max(1, std::min((int) lctx.mtp_n_heads, n_mtp_heads_model)) + : n_mtp_heads_model; + const int step_idx = (int) std::min(std::max(0, lctx.mtp_step_idx), n_mtp_heads - 1); + + // carry input: head k's output at the last committed position (k = 1..n_heads_model-1), + // fixed at model width so the storage layout is head-count independent + ggml_tensor * inp_carry = nullptr; + const bool is_cache_update = cparams.mtp_op_type == MTP_OP_WARMUP || cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED; + const bool needs_carry = n_mtp_heads_model > 1 && + ((is_cache_update && n_mtp_heads > 1) || + (cparams.mtp_op_type == MTP_OP_DRAFT_GEN && step_idx == 1 && n_mtp_heads > 2)); + if (needs_carry) { + inp_carry = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_mtp_heads_model - 1); + ggml_set_name(inp_carry, "inp_mtp_carry"); + ggml_set_input(inp_carry); + lctx.inp_mtp_carry = inp_carry; + } + + ggml_tensor * mtp_out = nullptr; + if (cparams.mtp_op_type == MTP_OP_DRAFT_GEN) { + const int il_mtp = il_mtp_first + step_idx; + mtp_out = build_openpangu_mtp(model.layers[il_mtp], + hidden_states_from_main_model, gf, il_mtp, + inp_pos, KQ_mask, inp_out_ids, inp_tokens, + seq_qnext, + nullptr, true, true, false, + KQ_mask_swa_windowed); + // each draft step runs one head, so a deeper head has no cache row at this + // position when its own decode comes later. Head h first decodes at step h+1 and + // the update batches cover everything below the draft base, which with three + // heads leaves exactly one gap: head 3's row at draft step 2. Pre-write it here + // from the committed carry (head 2's output at the last committed position). + if (step_idx == 1 && n_mtp_heads > 2) { + ggml_tensor * fill_hidden = ggml_view_2d(ctx0, inp_carry, hparams.n_embd, 1, + inp_carry->nb[1], (size_t) 1 * inp_carry->nb[1]); + ggml_tensor * fill = build_openpangu_mtp(model.layers[il_mtp_first + 2], + fill_hidden, gf, il_mtp_first + 2, + inp_pos, KQ_mask, nullptr, inp_tokens, + seq_qnext, + nullptr, false, false, + /*cache_writes_only=*/true, + KQ_mask_swa_windowed); + ggml_build_forward_expand(gf, fill); + } + } else if (n_mtp_heads == 1) { + mtp_out = build_openpangu_mtp(model.layers[il_mtp_first], + hidden_states_from_main_model, gf, il_mtp_first, + inp_pos, KQ_mask, inp_out_ids, inp_tokens, + seq_qnext, + nullptr, true, true, false, + KQ_mask_swa_windowed); + } else { + ggml_tensor * prev_full = hidden_states_from_main_model; + ggml_tensor * head1_hidden = nullptr; + std::vector carry_out_rows; + for (int i = 0; i < n_mtp_heads; ++i) { + const int il_mtp = il_mtp_first + i; + // the last head's block output feeds nothing (no carry, no logits) - only + // its cache writes matter, so skip its FFN and norms + const bool is_last_head = i + 1 == n_mtp_heads; + ggml_tensor * full_hidden = nullptr; + ggml_tensor * out = build_openpangu_mtp(model.layers[il_mtp], + prev_full, gf, il_mtp, + inp_pos, KQ_mask, inp_out_ids, inp_tokens, + seq_qnext, + is_last_head ? nullptr : &full_hidden, + i == 0, + false, + /*cache_writes_only=*/is_last_head && i > 0, + KQ_mask_swa_windowed); + if (i == 0) { + head1_hidden = out; + } else { + ggml_build_forward_expand(gf, out); + } + if (i + 1 < n_mtp_heads) { + // shift: head i+1's row p consumes this head's row p-1; row 0's + // predecessor comes from the carry + ggml_tensor * carry_col = ggml_view_2d(ctx0, inp_carry, hparams.n_embd, 1, + inp_carry->nb[1], (size_t) i * inp_carry->nb[1]); + if (n_tokens > 1) { + ggml_tensor * shifted = ggml_view_2d(ctx0, full_hidden, hparams.n_embd, n_tokens - 1, + full_hidden->nb[1], 0); + prev_full = ggml_concat(ctx0, carry_col, shifted, 1); + } else { + prev_full = carry_col; + } + carry_out_rows.push_back(ggml_view_2d(ctx0, full_hidden, hparams.n_embd, 1, + full_hidden->nb[1], + (size_t) (n_tokens - 1) * full_hidden->nb[1])); + } + } + GGML_ASSERT(head1_hidden != nullptr); + + // committed carries for the next warmup/update and for draft-time row fills + ggml_tensor * carry_out = carry_out_rows[0]; + for (size_t k = 1; k < carry_out_rows.size(); ++k) { + carry_out = ggml_concat(ctx0, carry_out, carry_out_rows[k], 1); + } + carry_out = ggml_cont(ctx0, carry_out); + ggml_set_name(carry_out, "mtp_carry_out"); + ggml_set_output(carry_out); + ggml_build_forward_expand(gf, carry_out); + + const auto & head1_layer = model.layers[il_mtp_first]; + cb(head1_hidden, "result_norm", -1); + ggml_tensor * head = head1_layer.nextn.shared_head_head + ? head1_layer.nextn.shared_head_head : model.output; + mtp_out = llm_build_lora_mm(lctx, ctx0, head, head1_hidden); + cb(mtp_out, "result_output", -1); + } + + ggml_build_forward_expand(gf, mtp_out); + return gf; + } + + ggml_tensor * cur; + ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * KQ_mask = build_inp_KQ_mask(); + // SWA layers get the windowed mask (window 512 base); DSA layers keep the plain causal + // mask and add the indexer's top-k selection inside the attention builder. Absent + // schedule keys (n_swa == 0) keep every layer dense (pre-DSA GGUF fallback). + bool KQ_mask_swa_windowed = false; + ggml_tensor * KQ_mask_swa = hparams.n_swa > 0 ? openpangu_build_swa_mask_for_graph(*this, hparams.n_swa, &KQ_mask_swa_windowed) : nullptr; + lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens); + cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1); + ggml_set_input(lctx.inp_s_seq_qnext); + ggml_tensor * seq_qnext = lctx.inp_s_seq_qnext; + + // mHC entry: repeat the embedding into S residual streams -> R [n_embd, S, n_tokens] + ggml_tensor * R = ggml_repeat(ctx0, ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), + ggml_new_tensor_3d(ctx0, inpL->type, n_embd, S, n_tokens)); + ggml_set_name(R, "opg_R_entry"); + + // mHC pre: combine S streams -> x [n_embd, n_tokens]; also returns h_post [S,T], h_res_mix [S,S,T] + auto mhc_pre = [&](ggml_tensor * Rin, ggml_tensor * phi, ggml_tensor * alpha, + ggml_tensor * beta, ggml_tensor * gamma, + ggml_tensor ** h_post_out, ggml_tensor ** h_res_out) { + ggml_tensor * flat = ggml_reshape_2d(ctx0, ggml_cont(ctx0, Rin), n_embd * S, n_tokens); + ggml_tensor * normed = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps); + normed = ggml_mul(ctx0, normed, gamma); // [S*H, T] + ggml_tensor * mixes = ggml_mul_mat(ctx0, phi, normed); // [(S+2)*S, T] + ggml_tensor * h_pre = ggml_view_2d(ctx0, mixes, S, n_tokens, mixes->nb[1], 0); + ggml_tensor * h_post = ggml_view_2d(ctx0, mixes, S, n_tokens, mixes->nb[1], S*ggml_element_size(mixes)); + ggml_tensor * h_res = ggml_view_2d(ctx0, mixes, S*S, n_tokens, mixes->nb[1], 2*S*ggml_element_size(mixes)); + + // alpha = [a_pre, a_post, a_res]; beta = [b_pre(S), b_post(S), b_res(S*S)] + ggml_tensor * a_pre = ggml_view_1d(ctx0, alpha, 1, 0); + ggml_tensor * b_pre = ggml_view_1d(ctx0, beta, S, 0); + // cont is required: the CUDA broadcast-mul path misreads strided views (h_pre is a + // row-slice of mixes), while CPU handles the strides — token 0 right, tokens 1+ garbage + h_pre = ggml_add(ctx0, ggml_mul(ctx0, ggml_cont(ctx0, h_pre), a_pre), b_pre); // broadcast scalar + [S] + h_pre = ggml_sigmoid(ctx0, h_pre); // [S,T] (+hc_eps omitted, inert) + + // combine: x[h,t] = sum_s h_pre[s,t] * R[h,s,t] + ggml_tensor * hpre3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, h_pre), 1, S, n_tokens); + ggml_tensor * weighted = ggml_mul(ctx0, Rin, hpre3); // [H,S,T] + ggml_tensor * wperm = ggml_cont(ctx0, ggml_permute(ctx0, weighted, 1, 0, 2, 3)); // [S,H,T] + ggml_tensor * x = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, wperm), n_embd, n_tokens); // sum over S + + *h_post_out = ggml_cont(ctx0, h_post); + *h_res_out = ggml_cont(ctx0, h_res); + return x; + }; + + // mHC post: R_new[h,s,t] = h_post[s,t]*y[h,t] + sum_j m[s,j,t]*R[h,j,t] + auto mhc_post = [&](ggml_tensor * y, ggml_tensor * h_post, ggml_tensor * Rin, + ggml_tensor * alpha, ggml_tensor * beta, ggml_tensor * h_res) { + ggml_tensor * a_post = ggml_view_1d(ctx0, alpha, 1, 1*ggml_element_size(alpha)); + ggml_tensor * a_res = ggml_view_1d(ctx0, alpha, 1, 2*ggml_element_size(alpha)); + ggml_tensor * b_post = ggml_view_1d(ctx0, beta, S, S*ggml_element_size(beta)); + ggml_tensor * b_res = ggml_view_1d(ctx0, beta, S*S, 2*S*ggml_element_size(beta)); + + h_post = ggml_add(ctx0, ggml_mul(ctx0, h_post, a_post), b_post); + h_post = ggml_scale(ctx0, ggml_sigmoid(ctx0, h_post), 2.0f); // 2*sigmoid, [S,T] + + ggml_tensor * m = ggml_add(ctx0, ggml_mul(ctx0, h_res, a_res), b_res); // [S*S,T] + m = openpangu_sinkhorn(ctx0, m, S, n_tokens, sink_iters, hc_eps); // [row S, col S, T] + + // term1: h_post[s,t]*y[h,t] -> [H,S,T] + ggml_tensor * y3 = ggml_reshape_3d(ctx0, y, n_embd, 1, n_tokens); + ggml_tensor * hpost3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, h_post), 1, S, n_tokens); + ggml_tensor * term1 = ggml_mul(ctx0, ggml_repeat(ctx0, y3, + ggml_new_tensor_3d(ctx0, y->type, n_embd, S, n_tokens)), hpost3); + + // term2: sum_j m[s,j,t]*R[h,j,t]. For each out-stream s, weight over input streams j. + // Build via: for stream axis, matmul R[H, j, t] with m[j, s, t] batched over t. + // R_perm [j(S), H, T] ; m as [j(S), s(S), T]; batched mul_mat over T -> [H? ] messy. + // Simpler explicit loop over S output streams (S=4, cheap): + ggml_tensor * term2 = nullptr; + for (int64_t s = 0; s < S; ++s) { + // m_s = m[:, s, :] -> weights over input streams j: [S, T] + ggml_tensor * m_s = ggml_cont(ctx0, ggml_view_2d(ctx0, m, S, n_tokens, m->nb[2], s*m->nb[1])); + ggml_tensor * m_s3 = ggml_reshape_3d(ctx0, m_s, 1, S, n_tokens); // [1,S,T] + ggml_tensor * acc = ggml_mul(ctx0, Rin, m_s3); // [H,S,T] + ggml_tensor * accp = ggml_cont(ctx0, ggml_permute(ctx0, acc, 1, 0, 2, 3)); // [S,H,T] + ggml_tensor * summed = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, accp), n_embd, n_tokens); // [H,T] + summed = ggml_reshape_3d(ctx0, summed, n_embd, 1, n_tokens); + term2 = term2 ? ggml_concat(ctx0, term2, summed, 1) : summed; // -> [H,S,T] + } + return ggml_add(ctx0, term1, term2); // [H,S,T] + }; + + // Base generation uses only the transformer layers; the trailing NextN/MTP layers are skipped. + const int n_layer_base = n_layer - (int) hparams.nextn_predict_layers; + for (int il = 0; il < n_layer_base; ++il) { + auto & layer = model.layers[il]; + ggml_tensor * conv_state = (size_t) il < kv_self.s_l.size() ? kv_self.s_l[il] : nullptr; + + // ================= attention sublayer ================= + ggml_tensor * h_post_a, * h_res_a; + ggml_tensor * x = mhc_pre(R, layer.mhc_attn_phi, layer.mhc_attn_alpha, + layer.mhc_attn_beta, layer.mhc_attn_gamma, &h_post_a, &h_res_a); + if (il == 0) ggml_set_name(x, "opg0_attn_mhcpre_x"); + cur = llm_build_norm(ctx0, x, hparams, layer.attn_norm, NULL, LLM_NORM_RMS, cb, il); + if (il == 0) ggml_set_name(cur, "opg0_attn_norm"); + + const bool layer_swa = KQ_mask_swa && hparams.openpangu_window[il] > 0; + ggml_tensor * layer_mask = layer_swa ? KQ_mask_swa : KQ_mask; + cur = build_openpangu_attention(gf, layer, il, cur, layer_mask, inp_pos, + conv_state, seq_qnext, kq_scale, + layer_swa && KQ_mask_swa_windowed); + if (il == 0) ggml_set_name(cur, "opg0_attn_out"); + + cur = llm_build_norm(ctx0, cur, hparams, layer.attn_post_norm, NULL, LLM_NORM_RMS, cb, il); + if (il == 0) ggml_set_name(cur, "opg0_attn_postnorm"); + + // mHC post -> scatter back to S streams + R = mhc_post(cur, h_post_a, R, layer.mhc_attn_alpha, layer.mhc_attn_beta, h_res_a); + if (il == 0) ggml_set_name(R, "opg0_R_attn"); + + // ================= ffn sublayer ================= + ggml_tensor * h_post_m, * h_res_m; + ggml_tensor * xm = mhc_pre(R, layer.mhc_mlp_phi, layer.mhc_mlp_alpha, + layer.mhc_mlp_beta, layer.mhc_mlp_gamma, &h_post_m, &h_res_m); + cur = llm_build_norm(ctx0, xm, hparams, layer.ffn_norm, NULL, LLM_NORM_RMS, cb, il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = llm_build_ffn(ctx0, lctx, nullptr, cur, + layer.ffn_up, NULL, NULL, layer.ffn_gate, NULL, NULL, layer.ffn_down, NULL, NULL, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, cb, il); + } else { + ggml_tensor * moe_out = llm_build_moe_ffn(ctx0, lctx, cur, + layer.ffn_gate_inp, layer.ffn_up_exps, layer.ffn_gate_exps, layer.ffn_down_exps, + layer.ffn_exp_probs_b, n_expert, n_expert_used, LLM_FFN_SILU, + hparams.expert_weights_norm, true, hparams.expert_weights_scale, + (enum llm_expert_gating_func_type) hparams.expert_gating_func, cb, il, gf, false); + ggml_tensor * shexp = llm_build_ffn(ctx0, lctx, nullptr, cur, + layer.ffn_up_shexp, NULL, NULL, layer.ffn_gate_shexp, NULL, NULL, + layer.ffn_down_shexp, NULL, NULL, NULL, LLM_FFN_SILU, LLM_FFN_PAR, cb, il); + cur = ggml_add(ctx0, moe_out, shexp); + } + cur = llm_build_norm(ctx0, cur, hparams, layer.ffn_post_norm, NULL, LLM_NORM_RMS, cb, il); + if (il == 0) ggml_set_name(cur, "opg0_ffn_postnorm"); + + R = mhc_post(cur, h_post_m, R, layer.mhc_mlp_alpha, layer.mhc_mlp_beta, h_res_m); + + // block post-norm on the layer subset (RMSNorm over the concatenated S*H) + if (layer.block_post_norm) { + ggml_tensor * flat = ggml_reshape_2d(ctx0, ggml_cont(ctx0, R), n_embd * S, n_tokens); + flat = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps); + flat = ggml_mul(ctx0, flat, layer.block_post_norm); + R = ggml_reshape_3d(ctx0, flat, n_embd, S, n_tokens); + } + if (il == 0) ggml_set_name(R, "opg0_R_block"); + R = lctx.cvec.apply_to(ctx0, R, il); + } + + // mHC tail merge: collapse S streams -> 1 (pre_only) + { + ggml_tensor * flat = ggml_reshape_2d(ctx0, ggml_cont(ctx0, R), n_embd * S, n_tokens); + ggml_tensor * normed = ggml_mul(ctx0, ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps), model.mhc_merge_gamma); + ggml_tensor * w = ggml_mul_mat(ctx0, model.mhc_merge_phi, normed); // [S, T] + ggml_tensor * a_pre = ggml_view_1d(ctx0, model.mhc_merge_alpha, 1, 0); + w = ggml_sigmoid(ctx0, ggml_add(ctx0, ggml_mul(ctx0, w, a_pre), model.mhc_merge_beta)); // [S,T] + ggml_tensor * w3 = ggml_reshape_3d(ctx0, ggml_cont(ctx0, w), 1, S, n_tokens); + ggml_tensor * weighted = ggml_mul(ctx0, R, w3); // [H,S,T] + ggml_tensor * wperm = ggml_cont(ctx0, ggml_permute(ctx0, weighted, 1, 0, 2, 3)); // [S,H,T] + cur = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, wperm), n_embd, n_tokens); + } + + // select only the output tokens (the framework binds n_outputs rows, not all n_tokens). + // With MTP enabled, keep every token: the speculative framework consumes per-token + // hidden states (result_norm via pooling) to warm up / feed the NextN head. + if (!cparams.mtp) { + ggml_tensor * inp_out_ids = build_inp_out_ids(); + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cur = llm_build_norm(ctx0, cur, hparams, model.output_norm, NULL, LLM_NORM_RMS, cb, -1); + cb(cur, "result_norm", -1); + cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); + cb(cur, "result_output", -1); + ggml_build_forward_expand(gf, cur); + return gf; +} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 46f6882b9..d9f6f8047 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -86,6 +86,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GEMMA4_MTP, "gemma4_mtp" }, { LLM_ARCH_DFLASH_DRAFT, "dflash-draft" }, { LLM_ARCH_GEMMA4_ASSISTANT,"gemma4-assistant" }, + { LLM_ARCH_OPENPANGU, "openpangu" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -264,6 +265,10 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ADAPTER_TYPE, "adapter.type" }, { LLM_KV_ADAPTER_LORA_ALPHA, "adapter.lora.alpha" }, + + { LLM_KV_OPENPANGU_MHC_NUM_STREAM, "%s.mhc_num_stream" }, + { LLM_KV_OPENPANGU_MHC_RECUR_NORM, "%s.mhc_recur_norm" }, + { LLM_KV_OPENPANGU_PARAM_SINK_NUMBER, "%s.param_sink_number" }, }; LLM_KV::LLM_KV(llm_arch arch, const char* suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index 1bb21d12b..45bcbbf99 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -84,6 +84,7 @@ enum llm_arch { LLM_ARCH_GEMMA4_MTP, LLM_ARCH_DFLASH_DRAFT, LLM_ARCH_GEMMA4_ASSISTANT, + LLM_ARCH_OPENPANGU, LLM_ARCH_UNKNOWN, }; @@ -247,6 +248,11 @@ enum llm_kv { LLM_KV_ADAPTER_TYPE, LLM_KV_ADAPTER_LORA_ALPHA, + + // openPangu-2.0 (mHC / Hyper-Connections + param sink) + LLM_KV_OPENPANGU_MHC_NUM_STREAM, + LLM_KV_OPENPANGU_MHC_RECUR_NORM, + LLM_KV_OPENPANGU_PARAM_SINK_NUMBER, }; struct LLM_KV { @@ -381,6 +387,26 @@ enum llm_tensor { LLM_TENSOR_DFLASH_FC, LLM_TENSOR_DFLASH_HIDDEN_NORM, + // openPangu-2.0 + LLM_TENSOR_ATTN_QA_CONV, // MoME causal conv on q-lora latent + LLM_TENSOR_ATTN_KV_CONV, // MoME causal conv on compressed-kv latent + LLM_TENSOR_ATTN_O_CONV, // MoME causal conv on attn output + LLM_TENSOR_ATTN_PARAM_SINK_KV, // learned static latent-kv sink + LLM_TENSOR_ATTN_PARAM_SINK_K_PE, + LLM_TENSOR_BLOCK_POST_NORM, // sandwich whole-block post-norm (layer subset) + LLM_TENSOR_MHC_ATTN_PHI, // mHC / Hyper-Connections (per attn sublayer) + LLM_TENSOR_MHC_ATTN_ALPHA, + LLM_TENSOR_MHC_ATTN_BETA, + LLM_TENSOR_MHC_ATTN_GAMMA, + LLM_TENSOR_MHC_MLP_PHI, // mHC (per mlp sublayer) + LLM_TENSOR_MHC_MLP_ALPHA, + LLM_TENSOR_MHC_MLP_BETA, + LLM_TENSOR_MHC_MLP_GAMMA, + LLM_TENSOR_MHC_MERGE_PHI, // mHC global stream-merge (non-block) + LLM_TENSOR_MHC_MERGE_ALPHA, + LLM_TENSOR_MHC_MERGE_BETA, + LLM_TENSOR_MHC_MERGE_GAMMA, + LLM_TENSOR_UNKNOWN, }; diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 27132606b..117ea2464 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -106,6 +106,8 @@ void llm_build_context::init() { lctx.inp_out_ids = nullptr; lctx.inp_KQ_mask = nullptr; lctx.inp_KQ_mask_swa = nullptr; + lctx.inp_KQ_mask_swa_win = nullptr; + lctx.openpangu_swa_window_view = {}; lctx.inp_K_shift = nullptr; lctx.inp_mean = nullptr; lctx.inp_cls = nullptr; @@ -117,6 +119,7 @@ void llm_build_context::init() { lctx.inp_embd_enc = nullptr; lctx.inp_KQ_mask_cross = nullptr; lctx.inp_dsa_sink = nullptr; + lctx.inp_mtp_carry = nullptr; lctx.dflash.inputs.target_features = nullptr; lctx.dflash.inputs.pos_ctx = nullptr; lctx.dflash.inputs.kq_mask = nullptr; @@ -521,6 +524,25 @@ ggml_tensor * llm_build_context::build_inp_KQ_mask_swa(bool causal) { return flash_attn ? ggml_cast(ctx0, lctx.inp_KQ_mask_swa, GGML_TYPE_F16) : lctx.inp_KQ_mask_swa; } +ggml_tensor * llm_build_context::build_inp_KQ_mask_swa_win(int64_t n_kv_win, bool causal) { + GGML_ASSERT(hparams.n_swa > 0); + GGML_ASSERT(n_kv_win > 0); + if (causal && flash_attn) { + lctx.inp_KQ_mask_swa_win = ggml_new_tensor_2d(ctx0, GGML_TYPE_F16, n_kv_win, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)); + cb(lctx.inp_KQ_mask_swa_win, "KQ_mask_swa_win", -1); + ggml_set_input(lctx.inp_KQ_mask_swa_win); + return lctx.inp_KQ_mask_swa_win; + } + + lctx.inp_KQ_mask_swa_win = causal + ? ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv_win, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)) + : ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tokens, GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)); + cb(lctx.inp_KQ_mask_swa_win, "KQ_mask_swa_win", -1); + ggml_set_input(lctx.inp_KQ_mask_swa_win); + + return flash_attn ? ggml_cast(ctx0, lctx.inp_KQ_mask_swa_win, GGML_TYPE_F16) : lctx.inp_KQ_mask_swa_win; +} + ggml_tensor * llm_build_context::build_inp_mean() { lctx.inp_mean = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tokens, n_tokens); cb(lctx.inp_mean, "inp_mean", -1); @@ -2623,6 +2645,10 @@ ggml_cgraph * llm_build_context::llama_build_graph( { result = llm.build_deepseek2(); } break; + case LLM_ARCH_OPENPANGU: + { + result = llm.build_openpangu(); + } break; case LLM_ARCH_CHATGLM: { result = llm.build_chatglm(); diff --git a/src/llama-build-context.h b/src/llama-build-context.h index 32a667392..854b7b5a8 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -139,6 +139,8 @@ struct llm_build_context { ggml_tensor * build_inp_KQ_mask_swa(bool causal = true); + ggml_tensor * build_inp_KQ_mask_swa_win(int64_t n_kv_win, bool causal = true); + ggml_tensor * build_inp_mean(); ggml_tensor * build_inp_cls(); @@ -273,6 +275,39 @@ struct llm_build_context { ggml_cgraph * build_arctic(); ggml_cgraph * build_deepseek2(); + ggml_cgraph * build_openpangu(); + + // openPangu attention sublayer body (shared by base layers and the NextN/MTP head): + // input is the already-input-normed hidden; returns the post-o_proj attention output. + ggml_tensor * build_openpangu_attention( + ggml_cgraph * gf, + const struct llama_layer & layer, + int il, + ggml_tensor * x_normed, + ggml_tensor * KQ_mask, + ggml_tensor * inp_pos, + ggml_tensor * conv_state, + ggml_tensor * seq_qnext, + float kq_scale, + bool KQ_mask_swa_windowed = false); + + // openPangu NextN/MTP head (plain-residual block, no mHC): eh_proj stitching -> + // attention -> MoE -> shared head. Returns the draft logits tensor. + ggml_tensor * build_openpangu_mtp( + const struct llama_layer & mtp_layer, + ggml_tensor * prev_embeddings, + ggml_cgraph * gf, + int il, + ggml_tensor * inp_pos, + ggml_tensor * KQ_mask, + ggml_tensor * inp_out_ids, + ggml_tensor * inp_tokens, + ggml_tensor * seq_qnext, + ggml_tensor ** full_hidden_out = nullptr, + bool select_outputs = true, + bool build_logits = true, + bool cache_writes_only = false, + bool KQ_mask_swa_windowed = false); ggml_tensor * build_deepseek2_tp_attention( ggml_cgraph * gf, int il, diff --git a/src/llama-context.h b/src/llama-context.h index cd744c969..6c8c01852 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -13,6 +13,28 @@ struct llama_model; #include #include +struct llama_openpangu_swa_window_view { + int64_t w_view = 0; + int64_t win_off = 0; + bool engaged = false; +}; + +static inline llama_openpangu_swa_window_view llama_openpangu_calc_swa_window_view( + int64_t n_kv, int64_t n_tokens, int64_t window, int64_t pad) { + llama_openpangu_swa_window_view result; + if (window <= 0 || n_kv <= 0) { + result.w_view = n_kv; + return result; + } + + const int64_t unpadded = window + pad + n_tokens; + const int64_t overcovered = pad > 1 ? ((unpadded + pad - 1) / pad) * pad : unpadded; + result.w_view = overcovered < n_kv ? overcovered : n_kv; + result.win_off = n_kv - result.w_view; + result.engaged = result.w_view < n_kv; + return result; +} + struct llama_kv_cell { llama_pos pos = -1; llama_pos delta = 0; @@ -42,6 +64,11 @@ struct llama_kv_cache { bool hybrid = false; bool v_trans = true; // the value tensor is transposed + // openPangu s_l holds position-strict MoME conv state, not per-sequence recurrent + // slots; Qwen3Next-style s_l handling (seq ops, state serialization, s_copy) must + // skip it. Speculative rollback snapshots/restores it via the whole-slot spec checkpoint. + bool s_l_position_strict = false; + // Note: The value of head isn't only used to optimize searching // for a free KV slot. llama_decode_internal also uses it, so it // cannot be freely changed after a slot has been allocated. @@ -61,10 +88,10 @@ struct llama_kv_cache { std::vector v_l; std::vector s_l; // per layer recurrent state storage (Qwen3Next) - // DSA lightning-indexer key cache (GLM-5.2 / DeepSeek-V3.2). One per layer, MQA single - // head: [indexer_head_size, kv_size]. Mirrors k_l but stores the (Hadamard-rotated) - // indexer keys so a decoded token scores against ALL past indexer keys, not just the - // current batch. Empty unless the model has the DSA indexer. + // Persistent DSA indexer-key cache. One per indexer layer, MQA single head: + // [indexer_head_size, kv_size]. Stores architecture-specific indexer keys in their + // scoring representation so a decoded token scores against all past indexer keys. + // Empty unless the model has the DSA indexer. std::vector kr_l; // When true, the delta_net graph builder will enable per-step SSM state saves @@ -372,6 +399,7 @@ struct llama_context { struct ggml_tensor * inp_out_ids; // I32 [n_outputs] struct ggml_tensor * inp_KQ_mask; // F32 [kv_size, n_batch] struct ggml_tensor * inp_KQ_mask_swa; // F32 [kv_size, n_batch] + struct ggml_tensor * inp_KQ_mask_swa_win = nullptr; // F32 [openPangu SWA W_view, n_batch] struct ggml_tensor * inp_K_shift; // I32 [kv_size] struct ggml_tensor * inp_mean; // F32 [n_batch, n_batch] struct ggml_tensor * inp_cls; // I32 [n_batch] @@ -384,14 +412,35 @@ struct llama_context { struct ggml_tensor * inp_KQ_mask_cross; // F32 [n_outputs_enc, n_batch] struct ggml_tensor * inp_scale = nullptr; // F32 [n_tokens] struct ggml_tensor * inp_mtp_states = nullptr; + struct ggml_tensor * inp_mtp_carry = nullptr; // F32 [n_embd, nextn-1] per-head hidden at the last committed position struct ggml_tensor * inp_dsa_sink = nullptr; // F32 [n_kv, n_tokens] per-sequence attention-sink boost for DSA indexer top-k struct ggml_tensor * inp_mask_inf = nullptr; + struct openpangu_swa_window_view_state { + bool active = false; + int32_t n_kv = 0; + int32_t n_tokens = 0; + uint32_t window = 0; + uint32_t pad = 0; + int64_t w_view = 0; + int64_t win_off = 0; + } openpangu_swa_window_view; + + // multi-head MTP chaining state: head k's output row at the last committed position, + // written back after each warmup/update decode and fed into the next MTP graph through + // inp_mtp_carry (zeroed when a prompt warmup restarts from position 0). The readback is + // issued async after compute; mtp_carry_pending marks a copy that must be synchronized + // before the host buffer is read or resized. + std::vector mtp_carry; + bool mtp_carry_pending = false; + ggml_backend_t ggml_backend_by_name(const char * name); struct Prev; std::unique_ptr prev; std::unique_ptr prev_mtp; + int32_t mtp_step_idx = 0; + int32_t mtp_n_heads = 0; void reset_scheduler(); bool can_reuse_graph(const llama_batch & u_batch); @@ -409,6 +458,8 @@ struct llama_context { // uninitialized -> wrong block-max-pool/top-k -> degraded/NaN sparse-FA decode). Register the // kr_l cpy per layer here and patch its offset in update_cache_copies(), exactly like K/V. std::vector dsa_cache_copies; + std::vector openpangu_cache_copies; + std::vector openpangu_cache_copies_mtp; bool update_cache_copies(); @@ -418,6 +469,8 @@ struct llama_context { bool prepare_mtp_graph_inputs( struct llama_context & lctx); void set_mtp_op_type(llama_mtp_op_type value); + void set_mtp_step_idx(int32_t value); + void set_mtp_n_heads(int32_t value); int max_nodes(int n_tokens, int n_kv) const; }; diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index de2a2c874..cc32e796b 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -1131,6 +1131,78 @@ void llm_load_hparams( default: model.type = e_model::MODEL_UNKNOWN; } } break; + case LLM_ARCH_OPENPANGU: + { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + + // openPangu routes with a sigmoid gate + e_score_correction bias + hparams.expert_gating_func = LLM_EXPERT_GATING_FUNC_TYPE_NONE; + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + if (hparams.expert_gating_func == LLM_EXPERT_GATING_FUNC_TYPE_NONE) { + hparams.expert_gating_func = LLM_EXPERT_GATING_FUNC_SIGMOID; + } + + // DSA lightning indexer + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false); + + // mHC / Hyper-Connections + learned param sink + ml.get_key(LLM_KV_OPENPANGU_MHC_NUM_STREAM, hparams.mhc_num_stream); + ml.get_key(LLM_KV_OPENPANGU_MHC_RECUR_NORM, hparams.mhc_recur_norm); + ml.get_key(LLM_KV_OPENPANGU_PARAM_SINK_NUMBER, hparams.param_sink_number); + + // NextN / MTP layers are appended at the end and skipped for base generation + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false); + if (hparams.nextn_predict_layers > 0 && hparams.nextn_predict_layers < hparams.n_layer) { + hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers; + } + + // DSA/SWA schedule: openpangu.swa_layers lists the sliding-window layer ids and + // openpangu.sliding_window_list the per-entry window; the remaining base layers + // are DSA (indexer + top-k, no window). The NextN/MTP layers appear in the SWA + // list with their own (larger) window, used by the MTP graphs. Absent keys keep + // every window at 0 = dense fallback (pre-DSA GGUFs keep working). + { + std::vector swa_ids, swa_windows; + const bool have_ids = ml.get_arr("openpangu.swa_layers", swa_ids, false); + const bool have_win = ml.get_arr("openpangu.sliding_window_list", swa_windows, false); + if (have_ids && have_win && swa_ids.size() == swa_windows.size()) { + const uint32_t n_base = hparams.n_layer > hparams.nextn_predict_layers + ? hparams.n_layer - hparams.nextn_predict_layers : hparams.n_layer; + for (size_t i = 0; i < swa_ids.size(); ++i) { + const uint32_t il = swa_ids[i]; + if (il >= hparams.n_layer) { + throw std::runtime_error(format("openpangu.swa_layers contains out-of-range layer %u", il)); + } + hparams.openpangu_window[il] = swa_windows[i]; + if (il < n_base) { + if (hparams.n_swa != 0 && hparams.n_swa != swa_windows[i]) { + throw std::runtime_error("openpangu: non-uniform base sliding windows are not supported"); + } + hparams.n_swa = swa_windows[i]; + } else { + if (hparams.n_swa_mtp != 0 && hparams.n_swa_mtp != swa_windows[i]) { + throw std::runtime_error("openpangu: non-uniform MTP sliding windows are not supported"); + } + hparams.n_swa_mtp = swa_windows[i]; + } + } + } else if (have_ids || have_win) { + LLAMA_LOG_WARN("%s: openpangu SWA schedule keys are inconsistent - keeping dense fallback\n", __func__); + } + } + + + model.type = e_model::MODEL_UNKNOWN; // 92B-A6B (46 + 3 MTP layers) + } break; case LLM_ARCH_CHATGLM: { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 10c910698..c52c842f0 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -131,6 +131,16 @@ struct llama_hparams { // indexer_types metadata if present, else derived from the GLM-5.2 config rule at load time. std::array indexer_is_full = {}; + // openPangu-2.0 (mHC / Hyper-Connections + learned param sink) + uint32_t mhc_num_stream = 1; + uint32_t mhc_recur_norm = 0; + uint32_t param_sink_number = 0; + // openPangu DSA/SWA schedule: per-layer sliding window (0 = DSA layer, full causal + // attention over the indexer's top-k selection). The NextN/MTP layers carry their own + // (larger) window, applied when the graph is built with an MTP op type. + uint32_t n_swa_mtp = 0; + std::array openpangu_window = {}; + // qwen3vl deepstack uint32_t n_deepstack_layers = 0; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 9981b6473..db3937202 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -119,6 +119,7 @@ struct create_tensors_helper : public create_tensors_helper_interface { bool create_arctix_tensors(const LLM_TN & tn); bool create_deepseek2_tensors(const LLM_TN & tn); + bool create_openpangu_tensors(const LLM_TN & tn); bool create_glm_dsa_tensors(const LLM_TN & tn); @@ -2870,6 +2871,151 @@ bool create_tensors_helper::create_glm_dsa_tensors(const LLM_TN & tn) { return use_mmap_buffer; } +// openPangu-2.0-Flash: GLM-DSA-style MLA + MoE base, plus mHC / MoME conv / param-sink / +// sandwich norms / DSA lightning indexer. The graph runs absorbed MLA over a latent KV +// cache from the converter's pre-split attn_k_b/attn_v_b; the indexer tensors feed the +// DSA top-k selection on windowless layers when the GGUF carries a DSA/SWA schedule. +// Conv weights keep their torch [C,1,kernel] layout (ggml ne = {kernel,1,C}). +bool create_tensors_helper::create_openpangu_tensors(const LLM_TN & tn) { + LOADING_PRELUDE + + const int64_t n_embd_head_qk_rope = hparams.n_rot; + const int64_t n_embd_head_qk_nope = hparams.n_embd_head_k(0) - hparams.n_rot; + + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_expert_shared = hparams.n_expert_shared; + + const int64_t S = hparams.mhc_num_stream; // 4 + const int64_t SH = S * n_embd; // concatenated multi-stream width + const int64_t phi_out = (S + 2) * S; // mHC phi output width + const int64_t beta_len = S * (S + 2); // mHC beta length + const int64_t sink_n = hparams.param_sink_number; // 128 + const int64_t kernel = 3; // MoME causal conv width + + model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); + + // output + { + model.output_norm = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}); + model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}); + } + + // global mHC stream-merge module (non-block) + model.mhc_merge_phi = create_tensor(ctx_output, tn(LLM_TENSOR_MHC_MERGE_PHI, "weight"), {SH, S}); + model.mhc_merge_alpha = create_tensor(ctx_output, tn(LLM_TENSOR_MHC_MERGE_ALPHA), {1}); + model.mhc_merge_beta = create_tensor(ctx_output, tn(LLM_TENSOR_MHC_MERGE_BETA), {S}); + model.mhc_merge_gamma = create_tensor(ctx_output, tn(LLM_TENSOR_MHC_MERGE_GAMMA), {SH}); + + for (int i = 0; i < n_layer; ++i) { + const bool is_mtp_layer = hparams.nextn_predict_layers > 0 && + static_cast(i) >= n_layer - hparams.nextn_predict_layers; + + int flags = 0; + if (!model.mtp && is_mtp_layer) { + flags |= llama_model_loader::TENSOR_SKIP | llama_model_loader::TENSOR_NOT_REQUIRED; + } + ggml_context * ctx_layer = ctx_for_layer(i); + ggml_context * ctx_split = ctx_for_layer_split(i); + + auto & layer = model.layers[i]; + + const auto graph_or_attn = (model.split_mode == LLAMA_SPLIT_MODE_GRAPH || + model.split_mode == LLAMA_SPLIT_MODE_ATTN); + auto norm_ctx = graph_or_attn ? ctx_split : ctx_layer; + auto moe_ctx = graph_or_attn ? ctx_split : ctx_layer; + + // --- norms (sandwich): input / post-attn / pre-mlp(ffn) / post-mlp --- + layer.attn_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + layer.attn_post_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, flags); + layer.attn_q_a_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags); + layer.attn_kv_a_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags); + // block_post_norm only present on a layer subset -> optional + // block_post_layernorm is RMSNorm over the concatenated S*H (mhc_num_stream * hidden) + layer.block_post_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_BLOCK_POST_NORM, "weight", i), {SH}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); + + // --- MLA projections (ik-native, pre-split k_b/v_b) --- + layer.wq_a = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags); + layer.wq_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, flags); + layer.wkv_a_mqa = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, flags); + // older GGUFs include the fused kv_b_proj, but the absorbed-MLA graph has no + // consumer for it (attention runs entirely on pre-split k_b/v_b below) -> skip + // when present and allow it to be absent from new conversions + layer.wkv_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_KV_B, "weight", i), {kv_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_v)}, flags | llama_model_loader::TENSOR_SKIP | llama_model_loader::TENSOR_NOT_REQUIRED); + // converter-emitted pre-split k_b/v_b, loaded 2D as written (head-major rows) and + // reshaped in-graph: k_b absorbs q_nope into latent space, v_b up-projects the output + layer.wk_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, n_head * kv_lora_rank}, flags); + layer.wv_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_head * n_embd_head_v}, flags); + layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, flags); + + // --- MoME causal convs (torch [C,1,3] -> gguf squeezes to 2D {3, C}) --- + layer.qa_conv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_QA_CONV, "weight", i), {kernel, q_lora_rank}, flags); + // compresskv_conv acts on the compressed-kv latent only (kv_lora_rank), NOT the k_pe part + layer.kv_conv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_KV_CONV, "weight", i), {kernel, kv_lora_rank}, flags); + layer.o_conv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_O_CONV, "weight", i), {kernel, n_head * n_embd_head_v}, flags); + + // --- learned static param sink (latent-kv prepended to attention) --- + layer.param_sink_kv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_PARAM_SINK_KV, i), {kv_lora_rank, sink_n}, flags); + layer.param_sink_k_pe = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_PARAM_SINK_K_PE, i), {n_embd_head_qk_rope, sink_n}, flags); + + // --- DSA indexer (loaded-but-unused in the dense-fallback graph) --- + layer.indexer_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); + layer.indexer_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, hparams.indexer_n_head}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); + layer.indexer_attn_k = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, hparams.indexer_head_size}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); + layer.indexer_attn_q_b = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, hparams.indexer_n_head * hparams.indexer_head_size}, flags | llama_model_loader::TENSOR_NOT_REQUIRED); + + // --- mHC / Hyper-Connections (per attn + per mlp sublayer, float32) --- + // MTP/NextN layers run WITHOUT mHC (tail_use_mhc=false in the reference; the + // checkpoint has no mhc tensors for them), so only create these for base layers. + if (!is_mtp_layer) { + layer.mhc_attn_phi = create_tensor(norm_ctx, tn(LLM_TENSOR_MHC_ATTN_PHI, "weight", i), {SH, phi_out}, flags); + layer.mhc_attn_alpha = create_tensor(norm_ctx, tn(LLM_TENSOR_MHC_ATTN_ALPHA, i), {3}, flags); + layer.mhc_attn_beta = create_tensor(norm_ctx, tn(LLM_TENSOR_MHC_ATTN_BETA, i), {beta_len}, flags); + layer.mhc_attn_gamma = create_tensor(norm_ctx, tn(LLM_TENSOR_MHC_ATTN_GAMMA, i), {SH}, flags); + layer.mhc_mlp_phi = create_tensor(norm_ctx, tn(LLM_TENSOR_MHC_MLP_PHI, "weight", i), {SH, phi_out}, flags); + layer.mhc_mlp_alpha = create_tensor(norm_ctx, tn(LLM_TENSOR_MHC_MLP_ALPHA, i), {3}, flags); + layer.mhc_mlp_beta = create_tensor(norm_ctx, tn(LLM_TENSOR_MHC_MLP_BETA, i), {beta_len}, flags); + layer.mhc_mlp_gamma = create_tensor(norm_ctx, tn(LLM_TENSOR_MHC_MLP_GAMMA, i), {SH}, flags); + } + + // --- FFN: dense-lead then MoE (routed + shared, sigmoid bias) --- + layer.ffn_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); + layer.ffn_post_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, flags); + + if (i < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags); + layer.ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, flags); + layer.ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags); + } else { + layer.ffn_gate_inp = create_tensor(moe_ctx, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); + layer.ffn_exp_probs_b = create_tensor(moe_ctx, tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags); + + GGML_ASSERT(n_expert > 0); + GGML_ASSERT(n_expert_used > 0); + + layer.ffn_gate_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + + layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, flags); + layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + } + + if (is_mtp_layer) { + layer.nextn.eh_proj = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), { 2 * n_embd, n_embd }, flags); + layer.nextn.enorm = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), { n_embd }, flags); + layer.nextn.embed_tokens = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, flags | llama_model_loader::TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab }, flags | llama_model_loader::TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_norm = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd }, flags); + } + } + return use_mmap_buffer; +} + bool create_tensors_helper::create_glm4_moe_tensors(const LLM_TN & tn) { LOADING_PRELUDE @@ -4514,6 +4660,8 @@ bool create_tensors_helper::create_tensors() { use_mmap_buffer = create_deepseek2_tensors(tn); break; case LLM_ARCH_GLM_DSA: use_mmap_buffer = create_glm_dsa_tensors(tn); break; + case LLM_ARCH_OPENPANGU: + use_mmap_buffer = create_openpangu_tensors(tn); break; case LLM_ARCH_GLM4_MOE: use_mmap_buffer = create_glm4_moe_tensors(tn); break; case LLM_ARCH_BITNET: diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 6d6cf5cf7..e53981887 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1376,3 +1376,4 @@ template std::enable_if::value, bool>::type llama template std::enable_if::value, bool>::type llama_model_loader::get_arr_n(enum llm_kv, unsigned int&, bool); template bool llama_model_loader::get_arr(const std::string &, std::array &, bool); template bool llama_model_loader::get_arr(const std::string &, std::array &, bool); +template bool llama_model_loader::get_arr(const std::string &, std::vector &, bool); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 093427183..f9048b60a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1775,6 +1775,71 @@ static const std::map> LLM_TENSOR_NA }, }, + { + LLM_ARCH_OPENPANGU, + { + { LLM_TENSOR_TOKEN_EMBD, "token_embd" }, + { LLM_TENSOR_OUTPUT_NORM, "output_norm" }, + { LLM_TENSOR_OUTPUT, "output" }, + { LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" }, + { LLM_TENSOR_ATTN_Q_A, "blk.%d.attn_q_a" }, + { LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" }, + { LLM_TENSOR_ATTN_Q_B, "blk.%d.attn_q_b" }, + { LLM_TENSOR_ATTN_KV_A_MQA, "blk.%d.attn_kv_a_mqa" }, + { LLM_TENSOR_ATTN_KV_A_NORM, "blk.%d.attn_kv_a_norm" }, + { LLM_TENSOR_ATTN_KV_B, "blk.%d.attn_kv_b" }, + { LLM_TENSOR_ATTN_K_B, "blk.%d.attn_k_b" }, + { LLM_TENSOR_ATTN_V_B, "blk.%d.attn_v_b" }, + { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, + { LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" }, + // DSA lightning indexer (loaded; unused in the dense-fallback graph) + { LLM_TENSOR_INDEXER_K_NORM, "blk.%d.attn_indexer_k_norm" }, + { LLM_TENSOR_INDEXER_PROJ, "blk.%d.attn_indexer_weights_proj" }, + { LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.attn_indexer_k" }, + { LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.attn_indexer_q_b" }, + // MoME causal convs + param sink + { LLM_TENSOR_ATTN_QA_CONV, "blk.%d.attn_qa_conv" }, + { LLM_TENSOR_ATTN_KV_CONV, "blk.%d.attn_compresskv_conv" }, + { LLM_TENSOR_ATTN_O_CONV, "blk.%d.attn_o_conv" }, + { LLM_TENSOR_ATTN_PARAM_SINK_KV, "blk.%d.attn_param_sink_kv" }, + { LLM_TENSOR_ATTN_PARAM_SINK_K_PE, "blk.%d.attn_param_sink_k_pe" }, + // FFN (dense-lead + MoE, sandwich norms) + { LLM_TENSOR_FFN_NORM, "blk.%d.ffn_norm" }, + { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, + { LLM_TENSOR_FFN_GATE, "blk.%d.ffn_gate" }, + { LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" }, + { LLM_TENSOR_FFN_DOWN, "blk.%d.ffn_down" }, + { LLM_TENSOR_FFN_GATE_INP, "blk.%d.ffn_gate_inp" }, + { LLM_TENSOR_FFN_GATE_EXPS, "blk.%d.ffn_gate_exps" }, + { LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" }, + { LLM_TENSOR_FFN_UP_EXPS, "blk.%d.ffn_up_exps" }, + { LLM_TENSOR_FFN_GATE_SHEXP, "blk.%d.ffn_gate_shexp" }, + { LLM_TENSOR_FFN_DOWN_SHEXP, "blk.%d.ffn_down_shexp" }, + { LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" }, + { LLM_TENSOR_FFN_EXP_PROBS_B, "blk.%d.exp_probs_b" }, + // mHC / Hyper-Connections + { LLM_TENSOR_MHC_ATTN_PHI, "blk.%d.attn_mhc_phi" }, + { LLM_TENSOR_MHC_ATTN_ALPHA, "blk.%d.attn_mhc_alpha" }, + { LLM_TENSOR_MHC_ATTN_BETA, "blk.%d.attn_mhc_beta" }, + { LLM_TENSOR_MHC_ATTN_GAMMA, "blk.%d.attn_mhc_gamma" }, + { LLM_TENSOR_MHC_MLP_PHI, "blk.%d.mlp_mhc_phi" }, + { LLM_TENSOR_MHC_MLP_ALPHA, "blk.%d.mlp_mhc_alpha" }, + { LLM_TENSOR_MHC_MLP_BETA, "blk.%d.mlp_mhc_beta" }, + { LLM_TENSOR_MHC_MLP_GAMMA, "blk.%d.mlp_mhc_gamma" }, + { LLM_TENSOR_MHC_MERGE_PHI, "merge_mhc_phi" }, + { LLM_TENSOR_MHC_MERGE_ALPHA, "merge_mhc_alpha" }, + { LLM_TENSOR_MHC_MERGE_BETA, "merge_mhc_beta" }, + { LLM_TENSOR_MHC_MERGE_GAMMA, "merge_mhc_gamma" }, + { LLM_TENSOR_BLOCK_POST_NORM, "blk.%d.block_post_norm" }, + // NextN / MTP tail + { LLM_TENSOR_NEXTN_EH_PROJ, "blk.%d.nextn.eh_proj" }, + { LLM_TENSOR_NEXTN_EMBED_TOKENS, "blk.%d.nextn.embed_tokens" }, + { LLM_TENSOR_NEXTN_ENORM, "blk.%d.nextn.enorm" }, + { LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" }, + { LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "blk.%d.nextn.shared_head_head" }, + { LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" }, + }, + }, { LLM_ARCH_UNKNOWN, { @@ -2050,6 +2115,10 @@ bool llama_model_has_recurrent(const llama_model * model) { return llm_arch_is_hybrid(model->arch) || llm_arch_is_recurrent(model->arch); } +bool llama_model_is_openpangu(const llama_model * model) { + return model && model->arch == LLM_ARCH_OPENPANGU; +} + bool llama_model_is_gemma4_mtp_assistant(const llama_model * model) { return model && (model->arch == LLM_ARCH_GEMMA4_MTP || model->arch == LLM_ARCH_GEMMA4_ASSISTANT); } @@ -2075,6 +2144,18 @@ bool llama_model_is_split_mode_graph(const struct llama_model * model) { return model && (model->split_mode == LLAMA_SPLIT_MODE_GRAPH || model->split_mode == LLAMA_SPLIT_MODE_ATTN); } +bool llama_model_supports_ctx_shift(const struct llama_model * model) { + // openPangu's latent K rows carry baked-in rope (k_pe) and the DSA indexer cache is + // keyed by absolute position; neither survives K-shift/defrag-style repositioning. + return model && model->arch != LLM_ARCH_OPENPANGU; +} + +bool llama_model_supports_partial_kv_reuse(const struct llama_model * model) { + // openPangu keeps only the current MoME conv state, so a sequence can be extended or + // reset, but rewinding into its decoded middle cannot reconstruct the state at that point. + return model && model->arch != LLM_ARCH_OPENPANGU; +} + llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int il) { auto it = LLM_TENSOR_NAMES.find(arch); if (it == LLM_TENSOR_NAMES.end()) { @@ -2106,12 +2187,24 @@ llm_tensor llm_tensor_type(llm_arch arch, const std::string & tensor_name, int i return LLM_TENSOR_UNKNOWN; } -size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const { +size_t llama_model::cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_type idx_type_k, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const { if (il < 0 || il >= hparams.n_layer) return 0; if (hparams.recurrent_layer_arr[il]) { auto state_sots = std::min(std::max(1, n_seq_max), kv_size); return hparams.n_embd_v_s() * state_sots * sizeof(float); } + if (arch == LLM_ARCH_OPENPANGU) { + // MLA-latent cache: K row [ckv | roped k_pe]. The value-side latent is + // rederived from K per graph. DSA layers also cache one indexer key per + // position. The recurrent conv slot is constant-size and negligible here. + size_t size = ggml_row_size(type_k, hparams.n_lora_kv + hparams.n_rot) * kv_size; + if (hparams.indexer_head_size > 0 && hparams.n_swa > 0 && + il < (int) hparams.n_layer - (int) hparams.nextn_predict_layers && + hparams.openpangu_window[il] == 0) { + size += ggml_row_size(idx_type_k, hparams.indexer_head_size) * kv_size; + } + return size; + } bool is_mla_attn = arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MISTRAL4; if (is_mla_attn && mla_attn) { auto n_embd_head_qk_rope = hparams.n_rot; diff --git a/src/llama-model.h b/src/llama-model.h index a0955cb43..b940d66ee 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -388,10 +388,30 @@ struct llama_layer { struct llama_layer_nextn nextn; + // openPangu-2.0: MoME causal convs + learned static param sink + mHC + block post-norm + struct ggml_tensor * qa_conv = nullptr; + struct ggml_tensor * kv_conv = nullptr; // compresskv_conv + struct ggml_tensor * o_conv = nullptr; + struct ggml_tensor * param_sink_kv = nullptr; + struct ggml_tensor * param_sink_k_pe = nullptr; + struct ggml_tensor * param_sink_blk = nullptr; + struct ggml_tensor * param_sink_lat_t = nullptr; + struct ggml_tensor * block_post_norm = nullptr; + struct ggml_tensor * mhc_attn_phi = nullptr; + struct ggml_tensor * mhc_attn_alpha = nullptr; + struct ggml_tensor * mhc_attn_beta = nullptr; + struct ggml_tensor * mhc_attn_gamma = nullptr; + struct ggml_tensor * mhc_mlp_phi = nullptr; + struct ggml_tensor * mhc_mlp_alpha = nullptr; + struct ggml_tensor * mhc_mlp_beta = nullptr; + struct ggml_tensor * mhc_mlp_gamma = nullptr; + std::unique_ptr computed_wk_b; std::unique_ptr computed_wk_b_pp; std::unique_ptr computed_wv_b; std::unique_ptr computed_wkv_b; + std::unique_ptr computed_param_sink_blk; + std::unique_ptr computed_param_sink_lat_t; // Per-device replicas of computed wk_b/wv_b (-sm graph). Buffers owned via model.bufs. std::vector> computed_wk_b_replicas; @@ -442,6 +462,12 @@ struct llama_model { struct ggml_tensor * output_norm_enc; struct ggml_tensor * output_mtp = nullptr; + // openPangu-2.0: global mHC stream-merge module (non-block) + struct ggml_tensor * mhc_merge_phi = nullptr; + struct ggml_tensor * mhc_merge_alpha = nullptr; + struct ggml_tensor * mhc_merge_beta = nullptr; + struct ggml_tensor * mhc_merge_gamma = nullptr; + std::unique_ptr output_mtp_ptr; llama_split_tensor split_output; @@ -545,7 +571,7 @@ struct llama_model { return hadamard_size(hparams.n_embd_head_v(il)); } - size_t cache_size(int il, ggml_type type_k, ggml_type type_v, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const; + size_t cache_size(int il, ggml_type type_k, ggml_type type_v, ggml_type idx_type_k, uint32_t kv_size, int mla_attn, int n_seq_max, bool flash_attn) const; void set_tensor_overrides(const llama_model_params& params); diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 3329668f1..edcd125bf 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2028,6 +2028,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "qwen2" || tokenizer_pre == "deepseek-r1-qwen" || tokenizer_pre == "kormo" || + tokenizer_pre == "openpangu" || tokenizer_pre == "f2llmv2") { pre_type = LLAMA_VOCAB_PRE_TYPE_QWEN2; clean_spaces = false; diff --git a/src/llama.cpp b/src/llama.cpp index fac695d4e..c06b7194a 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -28,6 +28,8 @@ #include "ggml-backend.h" void llama_set_mtp_target_context(struct llama_context * ctx, struct llama_context * target_ctx); +void llama_set_mtp_step_idx(struct llama_context * ctx, int32_t mtp_step_idx); +void llama_set_mtp_n_heads(struct llama_context * ctx, int32_t mtp_n_heads); // TODO: fix these includes #include "iqk/iqk_quantize.h" @@ -96,6 +98,7 @@ void llama_set_mtp_target_context(struct llama_context * ctx, struct llama_conte #include #include #include +#include #include #include #include @@ -556,6 +559,8 @@ struct llama_context::Prev { int save_per_step_ssm; int per_step_max_allocated; llama_mtp_op_type mtp_op_type; + int32_t mtp_step_idx; + int32_t mtp_n_heads; ggml_cgraph * graph; }; @@ -578,6 +583,8 @@ static void why_not_reuse_previous(const llama_batch & u_batch, const llama_cont if (ctx.n_outputs != the_prev->n_outputs) { printf(" n_outputs is not the same\n"); return; } if (u_batch.n_tokens != the_prev->n_tokens) { printf(" n_tokens is not the same\n"); return; } if (ctx.cparams.mtp_op_type != the_prev->mtp_op_type) { printf(" mtp_op_type is not the same\n"); return; } + if (ctx.mtp_step_idx != the_prev->mtp_step_idx) { printf(" mtp_step_idx is not the same\n"); return; } + if (ctx.mtp_n_heads != the_prev->mtp_n_heads) { printf(" mtp_n_heads is not the same\n"); return; } printf(" update_cache_copies() must have failed\n"); } @@ -596,6 +603,8 @@ bool llama_context::can_reuse_graph(const llama_batch & u_batch) { n_outputs == the_prev->n_outputs && u_batch.n_tokens == the_prev->n_tokens && cparams.mtp_op_type == the_prev->mtp_op_type && + mtp_step_idx == the_prev->mtp_step_idx && + mtp_n_heads == the_prev->mtp_n_heads && update_cache_copies(); if (false && !result) { printf("%s(%d):", __func__, cparams.mtp_op_type); @@ -606,6 +615,48 @@ bool llama_context::can_reuse_graph(const llama_batch & u_batch) { bool llama_context::update_cache_copies() { if (model.arch == LLM_ARCH_GEMMA4_MTP || model.arch == LLM_ARCH_GEMMA4_ASSISTANT) return true; + auto patch_dsa_cache_copies = [&]() -> bool { + // DSA indexer-key cache: patch the kr_l write offset for reused graphs. Each + // registered cpy writes this ubatch's index keys into kr_l at the kv_head slot; + // like the K/V copies, its baked view_offs must be re-pointed to the current + // kv_head or reused graphs keep writing to the first ubatch's slot. + // step = kr_l->nb[1]. No-op when DSA is off or entries are null. + for (size_t il = 0; il < dsa_cache_copies.size(); ++il) { + auto & c = dsa_cache_copies[il]; + if (!c.cpy) continue; + // The registered cpy must still be rooted at this layer's kr_l cache. If + // the graph was rebuilt with a different shape, refuse reuse and rebuild. + if (c.cpy->op != GGML_OP_CPY || + il >= kv_self.kr_l.size() || kv_self.kr_l[il] == nullptr || + c.cpy->view_src != kv_self.kr_l[il]) { + return false; + } + c.cpy->view_offs = kv_self.head * c.step; + c.cpy->src[1]->data = (char *) kv_self.kr_l[il]->data + c.cpy->view_offs; + c.cpy->data = c.cpy->src[1]->data; + } + return true; + }; + if (model.arch == LLM_ARCH_OPENPANGU) { + auto & copies = cparams.mtp_op_type == MTP_OP_NONE ? openpangu_cache_copies : openpangu_cache_copies_mtp; + bool any = false; + for (auto & c : copies) { + if (!c.cpy) { + continue; + } + any = true; + if (c.cpy->op != GGML_OP_CPY || c.cpy->view_src == nullptr || c.cpy->src[1] == nullptr) { + return false; + } + c.cpy->view_offs = kv_self.head*c.step; + c.cpy->src[1]->data = (char *)c.cpy->view_src->data + c.cpy->view_offs; + c.cpy->data = c.cpy->src[1]->data; + } + if (!patch_dsa_cache_copies()) { + return false; + } + return any; + } const int n_layer = model.mtp && cparams.mtp_op_type != MTP_OP_NONE ? model.hparams.n_layer : model.hparams.n_layer - model.hparams.nextn_predict_layers; //cache_copies.size()/2; auto layer_has_attention_kv = [&](int il) { @@ -672,29 +723,7 @@ bool llama_context::update_cache_copies() { } } } - // GLM-DSA lightning indexer: patch the indexer-key (kr_l) cache write offset for the reused - // graph. Each registered cpy writes this ubatch's index keys into kr_l at the kv_head slot; - // like the K/V copies above, its baked view_offs must be re-pointed to the CURRENT kv_head, - // else a reused graph (FA pad-256, constant n_kv) keeps writing to the first ubatch's slot and - // the recent indexer-key cells read uninitialized (scattered selection -> degraded/NaN decode). - // step = kr_l->nb[1] (one index-key row = head_size * F16). No-op when DSA is off (entries null). - for (size_t il = 0; il < dsa_cache_copies.size(); ++il) { - auto & c = dsa_cache_copies[il]; - if (!c.cpy) continue; - // Sanity guard (the MSA fix omitted this): the registered cpy must still be a CPY whose - // destination view is rooted at this layer's kr_l cache (mirrors the K/V guard above, - // which checks c.cpy->view_src). If the graph was rebuilt with a different shape these no - // longer match, so refuse reuse and force a rebuild. - if (c.cpy->op != GGML_OP_CPY || - il >= kv_self.kr_l.size() || kv_self.kr_l[il] == nullptr || - c.cpy->view_src != kv_self.kr_l[il]) { - return false; - } - c.cpy->view_offs = kv_self.head * c.step; - c.cpy->src[1]->data = (char *) kv_self.kr_l[il]->data + c.cpy->view_offs; - c.cpy->data = c.cpy->src[1]->data; - } - return true; + return patch_dsa_cache_copies(); } static std::vector & llama_all_contexts() { @@ -710,9 +739,11 @@ llama_context::llama_context(const llama_model & model) } else { cache_copies.resize(2*hparams.n_layer); } - // GLM-DSA lightning indexer: one indexer-key (kr_l) cache copy per layer. Entries stay null - // for non-DSA models / non-indexer layers, so update_cache_copies() is a no-op when DSA is off. + // DSA indexer-key cache copy. Entries stay null for non-DSA models and non-indexer layers, + // so update_cache_copies() is a no-op when DSA is off. dsa_cache_copies.resize(hparams.n_layer); + openpangu_cache_copies.resize(hparams.n_layer); + openpangu_cache_copies_mtp.resize(hparams.n_layer); llama_all_contexts().push_back(this); } @@ -722,6 +753,18 @@ void llama_context::set_mtp_op_type(llama_mtp_op_type value) { cparams.mtp_op_type = value; } +void llama_context::set_mtp_step_idx(int32_t value) { + LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); + + mtp_step_idx = std::max(0, value); +} + +void llama_context::set_mtp_n_heads(int32_t value) { + LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); + + mtp_n_heads = std::max(0, value); +} + llama_context::~llama_context() { if (dflash.kv.cache_sched != nullptr) { ggml_backend_sched_free(dflash.kv.cache_sched); @@ -744,8 +787,11 @@ llama_context::~llama_context() { } } +static int llama_openpangu_chunked_graph_nodes(const llama_model & model, const llama_cparams & cparams, int n_tokens, int n_kv); + int llama_context::max_nodes(int n_tokens, int n_kv) const { int max_nodes = model.max_nodes(n_tokens); + max_nodes += llama_openpangu_chunked_graph_nodes(model, cparams, n_tokens, n_kv); if (model.is_mla_model() && cparams.mla_attn > 1 && n_tokens >= 128 && @@ -787,6 +833,11 @@ static inline uint32_t llama_kv_v_row_embd( const llama_model & model, const llama_hparams & hparams, uint32_t il) { + // openPangu caches the 512-d compressed latent as its (transposed) V, not per-head values + if (model.arch == LLM_ARCH_OPENPANGU) { + return hparams.n_lora_kv; + } + // qwen3next recurrent state is stored in a dedicated V-cache tail (per sequence), // so per-token V rows include only attention values. if (llm_arch_is_hybrid(model.arch)) { @@ -796,6 +847,158 @@ static inline uint32_t llama_kv_v_row_embd( return hparams.n_embd_v_gqa(il) + hparams.n_embd_v_s(); } +// per-token K-cache row width in elements: openPangu stores [ckv 512 | roped k_pe 64] +// latents instead of per-head keys +static inline uint32_t llama_kv_k_row_embd( + const llama_model & model, + const llama_hparams & hparams, + uint32_t il) { + if (model.arch == LLM_ARCH_OPENPANGU) { + return hparams.n_lora_kv + hparams.n_rot; + } + return hparams.n_embd_k_gqa(il) + hparams.n_embd_k_s(); +} + +static int64_t llama_div_ceil_i64(int64_t num, int64_t denom) { + GGML_ASSERT(num >= 0); + GGML_ASSERT(denom > 0); + return (num + denom - 1) / denom; +} + +static int llama_openpangu_chunked_graph_nodes(const llama_model & model, const llama_cparams & cparams, int n_tokens, int n_kv) { + const llama_hparams & hparams = model.hparams; + if (model.arch != LLM_ARCH_OPENPANGU || n_tokens <= 14 || n_kv <= 0) { + return 0; + } + + // Keep these fixed constants in sync with build_openpangu.cpp. This budget + // estimator mirrors that file's chunk loops, including dense-fallback attention + // chunks when the DSA schedule/indexer is absent, plus a small fixed margin below + // for graph-node bookkeeping drift. + static constexpr int64_t OPENPANGU_IDX_SCORE_CHUNK = 256; + static constexpr int64_t OPENPANGU_ATT_SCORE_CHUNK = 256; + static constexpr int64_t OPENPANGU_ATT_FULL_KQ_MAX_MIB = 1024; + static constexpr int64_t OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX = 65535; + static constexpr int64_t OPENPANGU_DSA_GATHER_MIN_RATIO = 2; + + const int64_t idx_chunk = OPENPANGU_IDX_SCORE_CHUNK; + const int64_t att_chunk = OPENPANGU_ATT_SCORE_CHUNK; + const int64_t att_cap_mib = OPENPANGU_ATT_FULL_KQ_MAX_MIB; + + const int64_t topk = cparams.dsa_top_k > 0 ? cparams.dsa_top_k : hparams.indexer_top_k; + const bool has_dsa_indexer = topk > 0 && hparams.indexer_head_size > 0 && hparams.n_swa > 0; + + const int64_t n_layer_base = hparams.n_layer > hparams.nextn_predict_layers ? + hparams.n_layer - hparams.nextn_predict_layers : hparams.n_layer; + const int64_t n_heads = hparams.n_head(0); + const int64_t n_sinks = hparams.param_sink_number; + const int64_t pad = cparams.flash_attn ? 256 : 32; + const int64_t idx_chunks = idx_chunk > 0 && n_tokens > idx_chunk ? llama_div_ceil_i64(n_tokens, idx_chunk) : 0; + + int64_t extra_nodes = 0; + for (int64_t il = 0; il < n_layer_base; ++il) { + const bool is_swa_layer = hparams.n_swa > 0 && hparams.openpangu_window[il] > 0; + const bool is_dsa_layer = has_dsa_indexer && hparams.openpangu_window[il] == 0; + const int64_t n_kv_eff = !is_swa_layer ? n_kv : + llama_openpangu_calc_swa_window_view(n_kv, n_tokens, hparams.openpangu_window[il], pad).w_view; + const double full_kq_bytes = + (double) (n_kv_eff + n_sinks) * (double) n_heads * (double) n_tokens * (double) sizeof(float); + const bool att_chunks = + att_chunk > 0 && att_cap_mib > 0 && n_tokens > att_chunk && + full_kq_bytes > (double) att_cap_mib * 1024.0 * 1024.0; + + if (!is_dsa_layer) { + // SWA and dense-fallback layers only use the outer attention chunk loop when + // the full score tensor would exceed the cap; there is no top-k gather subchunk loop. + if (att_chunks) { + extra_nodes += llama_div_ceil_i64(n_tokens, att_chunk) * 32; + } + continue; + } + + // DSA layers build one top-k indexer chain per token chunk. In the deep prefill + // path, attention then adds an outer attention chunk loop plus a CUDA-grid-safe + // get_rows subchunk loop where topk * subchunk_tokens <= 65535. + extra_nodes += idx_chunks * 24; + if (att_chunks) { + const bool can_prefill_gather = + hparams.f_max_alibi_bias == 0.0f && + topk <= OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX; + const int64_t gather_tokens = can_prefill_gather ? + std::max(1, OPENPANGU_CUDA_GET_ROWS_GRID_Y_MAX / topk) : 1; + + for (int64_t c0 = 0; c0 < n_tokens; c0 += att_chunk) { + const int64_t tc = std::min(att_chunk, n_tokens - c0); + const bool prefill_gather = + can_prefill_gather && + n_kv >= n_tokens && + n_kv - n_tokens + c0 >= OPENPANGU_DSA_GATHER_MIN_RATIO*topk + pad; + if (prefill_gather) { + extra_nodes += 32 + llama_div_ceil_i64(tc, gather_tokens) * 32; + } else { + extra_nodes += 48; + } + } + } + } + + if (extra_nodes > 0) { + extra_nodes += 64 + extra_nodes / 20; + } + + return extra_nodes > (int64_t) INT_MAX ? INT_MAX : (int) extra_nodes; +} + +static inline bool llama_openpangu_latent_k_cache_type_supported(ggml_type type) { + return type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_Q8_0; +} + +static inline bool llama_openpangu_latent_v_cache_type_supported(ggml_type type) { + return type == GGML_TYPE_F32 || type == GGML_TYPE_F16; +} + +static std::string llama_openpangu_latent_cache_type_error(ggml_type type_k, ggml_type type_v) { + return format("OpenPangu latent K cache supports only f32, f16, and q8_0; latent V cache supports only f32 and f16 (requested %s/%s); use -ctk q8_0 -ctv f16", + ggml_type_name(type_k), ggml_type_name(type_v)); +} + +static bool llama_openpangu_validate_latent_cache_types( + ggml_type type_k, + ggml_type type_v, + std::string * error_msg) { + if (!llama_openpangu_latent_k_cache_type_supported(type_k) || + !llama_openpangu_latent_v_cache_type_supported(type_v)) { + if (error_msg) { + *error_msg = llama_openpangu_latent_cache_type_error(type_k, type_v); + } + return false; + } + + return true; +} + +static inline bool llama_openpangu_indexer_cache_type_supported(ggml_type type) { + return type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16 || type == GGML_TYPE_Q8_0; +} + +static std::string llama_openpangu_indexer_cache_type_error(ggml_type type) { + return format("OpenPangu indexer K cache supports only f32, f16, bf16, and q8_0 (requested %s); use -ictk q8_0", + ggml_type_name(type)); +} + +static bool llama_openpangu_validate_indexer_cache_type( + ggml_type idx_type_k, + std::string * error_msg) { + if (!llama_openpangu_indexer_cache_type_supported(idx_type_k)) { + if (error_msg) { + *error_msg = llama_openpangu_indexer_cache_type_error(idx_type_k); + } + return false; + } + + return true; +} + static inline uint32_t llama_qwen3next_state_slots(const llama_cparams & cparams, uint32_t kv_size) { return std::min(std::max(1, cparams.n_seq_max), kv_size); } @@ -820,6 +1023,11 @@ static inline uint32_t llama_kv_qnext_state_slots(const llama_kv_cache & cache) } static inline bool llama_kv_has_qnext_state_storage(const llama_kv_cache & cache) { + // openPangu s_l is position-strict conv state, not qnext per-sequence state; keep it + // out of qnext seq-copy and state serialization (rollback rides the spec checkpoint). + if (cache.s_l_position_strict) { + return false; + } return llama_kv_qnext_state_slots(cache) > 0; } @@ -969,10 +1177,12 @@ static bool llama_kv_cache_init( if (needs_v_cache) cache.v_l.reserve(n_layer); cache.s_l.resize(n_layer, nullptr); - // DSA lightning-indexer key cache: one [indexer_head_size, kv_size] tensor per layer. - // Allocated below (in the MLA branch) only when the model carries the indexer tensors. - const bool has_dsa_indexer = model.arch == LLM_ARCH_GLM_DSA && hparams.indexer_head_size > 0; - if (has_dsa_indexer) { + // DSA indexer-key cache: one [indexer_head_size, kv_size] tensor per indexer layer. + // Allocated below only when the model carries persistent DSA indexer tensors. + const bool has_glm_dsa_indexer = model.arch == LLM_ARCH_GLM_DSA && hparams.indexer_head_size > 0; + const bool has_openpangu_dsa_indexer = + model.arch == LLM_ARCH_OPENPANGU && hparams.indexer_head_size > 0 && hparams.n_swa > 0; + if (has_glm_dsa_indexer || has_openpangu_dsa_indexer) { cache.kr_l.resize(n_layer, nullptr); } @@ -991,7 +1201,8 @@ static bool llama_kv_cache_init( // For MTP-only context, skip KV allocation for non-MTP layers if (cparams.mtp_op_type != MTP_OP_NONE && i < n_mtp_first_layer) { cache.k_l.push_back(nullptr); - if (!is_mla_attn || !cparams.mla_attn || (cparams.mla_attn == 1 && !cparams.flash_attn)) { + if (model.arch != LLM_ARCH_OPENPANGU && + (!is_mla_attn || !cparams.mla_attn || (cparams.mla_attn == 1 && !cparams.flash_attn))) { cache.v_l.push_back(nullptr); } continue; @@ -1023,7 +1234,7 @@ static bool llama_kv_cache_init( cache.k_l.push_back(kv); // DSA lightning-indexer key cache (MQA, single head). Store the Hadamard-rotated // indexer keys in F16 so a decoded token can score against ALL past keys. - if (has_dsa_indexer && model.layers[i].indexer_attn_k && hparams.indexer_is_full[i] && !is_mtp_tail_layer) { + if (has_glm_dsa_indexer && model.layers[i].indexer_attn_k && hparams.indexer_is_full[i] && !is_mtp_tail_layer) { ggml_tensor * kr = ggml_new_tensor_2d(ctx, idx_type_k, hparams.indexer_head_size, kv_size); ggml_format_name(kr, "cache_kr_l%d", i); cache.kr_l[i] = kr; @@ -1110,34 +1321,67 @@ static bool llama_kv_cache_init( } int n_embd_head_v = hparams.n_embd_head_v(i); auto this_type_k = type_k; - if (type_k_first != type_k && n_k_first > 0 && i < n_k_first) { + if (model.arch != LLM_ARCH_OPENPANGU && type_k_first != type_k && n_k_first > 0 && i < n_k_first) { this_type_k = type_k_first; } - if (type_k_last != type_k && n_k_last > 0 && i >= n_layer - n_k_last) { + if (model.arch != LLM_ARCH_OPENPANGU && type_k_last != type_k && n_k_last > 0 && i >= n_layer - n_k_last) { this_type_k = type_k_last; } if (this_type_k != type_k) { LLAMA_LOG_INFO("================= Setting K-cache type in layer %2d to %s\n", i, ggml_type_name(this_type_k)); } - k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*kv_size); - int64_t v_ne = int64_t(n_embd_v_row)*kv_size; auto this_type_v = type_v; - if (type_v_first != type_v && n_v_first > 0 && i < n_v_first) { + if (model.arch != LLM_ARCH_OPENPANGU && type_v_first != type_v && n_v_first > 0 && i < n_v_first) { this_type_v = type_v_first; } - if (type_v_last != type_v && n_v_last > 0 && i >= n_layer - n_v_last) { + if (model.arch != LLM_ARCH_OPENPANGU && type_v_last != type_v && n_v_last > 0 && i >= n_layer - n_v_last) { this_type_v = type_v_last; } if (this_type_v != type_v) { LLAMA_LOG_INFO("================= Setting V-cache type in layer %2d to %s\n", i, ggml_type_name(this_type_v)); } - v = ggml_new_tensor_1d(ctx, this_type_v, v_ne); + + if (model.arch == LLM_ARCH_OPENPANGU) { + // MLA-latent cache: k_l holds [ckv_norm 512 | roped k_pe 64] per position. + // The value-side latent is rederived from k_l per graph; no persistent V store. + const int64_t n_lat = (int64_t) hparams.n_lora_kv + hparams.n_rot; // 576 + k = ggml_new_tensor_2d(ctx, this_type_k, n_lat, kv_size); + } else { + k = ggml_new_tensor_2d(ctx, this_type_k, n_embd_head_k, n_head_kv*kv_size); + v = ggml_new_tensor_1d(ctx, this_type_v, v_ne); + } auto k_name = std::string{"cache_k_l"} + std::to_string(i); auto v_name = std::string{"cache_v_l"} + std::to_string(i); ggml_set_name(k, k_name.c_str()); - ggml_set_name(v, v_name.c_str()); + if (v) { + ggml_set_name(v, v_name.c_str()); + } + + if (model.arch == LLM_ARCH_OPENPANGU) { + // MoME conv state for ggml_ssm_conv. Each qnext-style slot packs the three + // conv sites as two tap-contiguous floats per channel: + // [qa 2*n_lora_q | compresskv 2*n_lora_kv | o 2*n_head*v_dim]. + // s_l_position_strict stays true so qnext seq ops and state serialization + // skip this slot; speculative rollback snapshots/restores it via the + // whole-slot spec checkpoint. + const int64_t conv_col_ne = hparams.n_lora_q + hparams.n_lora_kv + + (int64_t) hparams.n_head(i)*hparams.n_embd_head_v(i); + ggml_tensor * s_conv = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2*conv_col_ne, qnext_state_slots); + ggml_format_name(s_conv, "cache_s_l%d", i); + cache.s_l[i] = s_conv; + cache.s_l_position_strict = true; + + // DSA layers (window == 0, indexer present) also cache the per-position + // 128-d indexer key. Position-indexed like everything else in the cache, so + // the same rollback invariant applies: committed columns never change. + if (has_openpangu_dsa_indexer && i < n_mtp_first_layer && hparams.openpangu_window[i] == 0) { + ggml_tensor * idxk = ggml_new_tensor_2d(ctx, idx_type_k, hparams.indexer_head_size, kv_size); + ggml_format_name(idxk, "cache_kr_l%d", i); + cache.kr_l[i] = idxk; + } + } if (split_cache_i) { bool use_V_for_K = model.layers[i].attn_k_norm && model.layers[i].attn_k_norm->ne[0] == K->ne[1] ? true : false; @@ -1178,7 +1422,9 @@ static bool llama_kv_cache_init( v->extra = (void *)&split_v_l.ggml; } cache.k_l.push_back(k); - cache.v_l.push_back(v); + if (model.arch != LLM_ARCH_OPENPANGU) { + cache.v_l.push_back(v); + } } } if (is_mla_attn && cparams.mla_attn && n_mla < n_kv_active_layers && n_mla > 0) { @@ -1222,7 +1468,7 @@ static bool llama_kv_cache_init( } printf("\n"); } - if (cache.v_l[il]->extra) { + if (cache.v_l.size() > (size_t) il && cache.v_l[il] && cache.v_l[il]->extra) { printf("Layer %2d, V-buffer: %p:", il, (void *)cache.v_l[il]->buffer); auto split_vl = (ggml_split_tensor_t *)cache.v_l[il]->extra; for (int id = 0; id < split_vl->n_device; ++id) { @@ -1961,6 +2207,11 @@ static void llama_kv_cache_seq_add( llama_pos p0, llama_pos p1, llama_pos delta) { + // position-strict caches (openPangu latent rows and indexer keys) cannot be + // re-positioned: rope is baked into the cached k_pe rows and the side state is keyed + // by absolute position. Fail loudly instead of corrupting. + GGML_ASSERT(!cache.s_l_position_strict && "K-shift/context shift is not supported for this model's position-indexed KV cache"); + uint32_t new_head = cache.size; if (p0 < 0) p0 = 0; @@ -2009,6 +2260,9 @@ static void llama_kv_cache_seq_div( llama_pos p0, llama_pos p1, int d) { + // see llama_kv_cache_seq_add: position-strict caches cannot be re-positioned + GGML_ASSERT(!cache.s_l_position_strict && "self-extend/position division is not supported for this model's position-indexed KV cache"); + if (p0 < 0) p0 = 0; if (p1 < 0) p1 = std::numeric_limits::max(); // If there is no range then return early to avoid looping over the cache. @@ -2064,6 +2318,12 @@ static llama_pos llama_kv_cache_seq_pos_min(struct llama_kv_cache & cache, llama } static void llama_kv_cache_defrag(struct llama_kv_cache & cache) { + if (cache.s_l_position_strict) { + // defrag moves cells, which would break the cell-index == position mapping the + // openPangu indexer cache relies on; skipping is safe (defrag is an optimization) + LLAMA_LOG_WARN("%s: defrag is not supported for this model's position-indexed KV cache - skipping\n", __func__); + return; + } cache.do_defrag = true; } @@ -2344,7 +2604,9 @@ static void llm_prepare_mla(llama_model & model, int mla) { const int n_layer = model.layers.size(); int n_to_compute = 0; for (auto& l : model.layers) { - if (!l.wk_b) ++n_to_compute; + // layers without attention weights (e.g. NextN/MTP heads whose tensors were + // skipped at load) have nothing to derive from + if (!l.wk_b && l.wkv_b) ++n_to_compute; } if (mla > 0 && n_to_compute > 0) { // Prepare wk_b tensors to enable MLA usage also for model files that do not include @@ -2374,7 +2636,7 @@ static void llm_prepare_mla(llama_model & model, int mla) { size_t max_wkv_size = 0; size_t max_wk_size = 0; for (auto& l : model.layers) { - if (!l.wk_b) { + if (!l.wk_b && l.wkv_b) { auto new_type = ggml_is_quantized(l.wkv_b->type) ? GGML_TYPE_Q8_0 : l.wkv_b->type; auto size = ggml_row_size(new_type, n_embd_head_qk_nope)*kv_lora_rank*n_head; max_wk_size = std::max(max_wk_size, size); @@ -2400,7 +2662,7 @@ static void llm_prepare_mla(llama_model & model, int mla) { std::vector tensor_data(2*n_embd_head_qk_nope*kv_lora_rank*n_head*sizeof(float) + 2*max_wk_size); for (int il = 0; il < n_layer; ++il) { auto& l = model.layers[il]; - if (l.wk_b) continue; + if (l.wk_b || !l.wkv_b) continue; auto wkv_b = *l.wkv_b; if (!ggml_backend_buffer_is_host(l.wkv_b->buffer)) { ggml_backend_tensor_get(l.wkv_b, wkv_buffer.data(), 0, ggml_nbytes(l.wkv_b)); @@ -2909,6 +3171,156 @@ static void llm_prepare_mla(llama_model & model, int mla) { ggml_free(ctx); } +static void llm_prepare_openpangu_param_sinks(llama_model & model) { + if (model.arch != LLM_ARCH_OPENPANGU) return; + + const auto & hparams = model.hparams; + const int n_layer = model.layers.size(); + const int64_t n_sink = hparams.param_sink_number; + if (n_sink <= 0) return; + + int n_to_compute = 0; + size_t max_kv_size = 0; + size_t max_kpe_size = 0; + size_t max_norm_size = 0; + size_t max_tensor_data = 0; + + for (auto & l : model.layers) { + if (l.param_sink_blk && l.param_sink_lat_t) continue; + if (!l.param_sink_kv || !l.param_sink_k_pe || !l.attn_kv_a_norm) continue; + + ++n_to_compute; + max_kv_size = std::max(max_kv_size, ggml_nbytes(l.param_sink_kv)); + max_kpe_size = std::max(max_kpe_size, ggml_nbytes(l.param_sink_k_pe)); + max_norm_size = std::max(max_norm_size, ggml_nbytes(l.attn_kv_a_norm)); + + const int64_t n_kv = l.param_sink_kv->ne[0]; + const int64_t n_kpe = l.param_sink_k_pe->ne[0]; + const size_t s_ckv_size = n_kv*n_sink*sizeof(float); + const size_t s_kpe_size = n_kpe*n_sink*sizeof(float); + const size_t sink_blk_size = (n_kv + n_kpe)*n_sink*sizeof(float); + const size_t s_lat_t_size = n_kv*n_sink*sizeof(float); + max_tensor_data = std::max(max_tensor_data, s_ckv_size + s_kpe_size + sink_blk_size + s_lat_t_size); + } + if (n_to_compute == 0) return; + + LLAMA_LOG_INFO("============ %s: need to compute %d openPangu param_sink tensors\n", __func__, 2*n_to_compute); + + const size_t context_size = ggml_tensor_overhead()*32*n_layer + 4*MiB; + ggml_init_params params{context_size, nullptr, true}; + auto ctx = ggml_init(params); + auto graph = ggml_new_graph_custom(ctx, 16, false); + + std::vector work_data; + std::vector kv_buffer(max_kv_size); + std::vector kpe_buffer(max_kpe_size); + std::vector norm_buffer(max_norm_size); + std::vector tensor_data(max_tensor_data); + + auto make_host_tensor = [](ggml_tensor * source, std::vector & buffer) { + ggml_tensor result = *source; + if (source->buffer && !ggml_backend_buffer_is_host(source->buffer)) { + const size_t nbytes = ggml_nbytes(source); + GGML_ASSERT(buffer.size() >= nbytes); + ggml_backend_tensor_get(source, buffer.data(), 0, nbytes); + result.data = buffer.data(); + } + return result; + }; + + auto materialize = [&model](ggml_tensor * source, + std::unique_ptr & computed, + ggml_backend_buffer_type_t buft, + const std::string & name) -> ggml_tensor * { + computed = std::make_unique(*source); + computed->buffer = ggml_backend_buft_alloc_buffer(buft, ggml_nbytes(source)); + if (!computed->buffer) { + throw std::runtime_error("Failed to allocate buffer for " + name); + } + model.bufs.push_back(computed->buffer); + computed->data = ggml_backend_buffer_get_base(computed->buffer); + computed->op = GGML_OP_NONE; + for (int j = 0; j < GGML_MAX_SRC; ++j) computed->src[j] = nullptr; + computed->view_src = nullptr; + computed->view_offs = 0; + computed->extra = nullptr; + ggml_set_name(computed.get(), name.c_str()); + ggml_backend_buffer_set_usage(computed->buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_tensor_set(computed.get(), source->data, 0, ggml_nbytes(source)); + if (ggml_backend_buffer_is_host(computed->buffer)) { + iqk_modify_tensor(computed.get()); + } + model.tensors_by_name.push_back(std::make_pair(name, computed.get())); + return computed.get(); + }; + + for (int il = 0; il < n_layer; ++il) { + auto & l = model.layers[il]; + if (l.param_sink_blk && l.param_sink_lat_t) continue; + if (!l.param_sink_kv || !l.param_sink_k_pe || !l.attn_kv_a_norm) continue; + + GGML_ASSERT(l.param_sink_kv->ne[1] == n_sink); + GGML_ASSERT(l.param_sink_k_pe->ne[1] == n_sink); + GGML_ASSERT(l.attn_kv_a_norm->ne[0] == l.param_sink_kv->ne[0]); + + auto param_sink_kv = make_host_tensor(l.param_sink_kv, kv_buffer); + auto param_sink_k_pe = make_host_tensor(l.param_sink_k_pe, kpe_buffer); + auto attn_kv_a_norm = make_host_tensor(l.attn_kv_a_norm, norm_buffer); + + char * ptr = tensor_data.data(); + char * const end = tensor_data.data() + tensor_data.size(); + + ggml_tensor * s_ckv = ggml_fused_rms_norm(ctx, ¶m_sink_kv, &attn_kv_a_norm, hparams.f_norm_rms_eps); + s_ckv->data = ptr; + ptr += ggml_nbytes(s_ckv); + + ggml_tensor * s_kpe = ¶m_sink_k_pe; + if (param_sink_k_pe.type != GGML_TYPE_F32) { + s_kpe = ggml_cast(ctx, ¶m_sink_k_pe, GGML_TYPE_F32); + s_kpe->data = ptr; + ptr += ggml_nbytes(s_kpe); + } + + ggml_tensor * sink_blk = ggml_concat(ctx, s_ckv, s_kpe, 0); + sink_blk->data = ptr; + ptr += ggml_nbytes(sink_blk); + + ggml_tensor * s_lat_t = ggml_cont(ctx, ggml_transpose(ctx, s_ckv)); + s_lat_t->data = ptr; + ptr += ggml_nbytes(s_lat_t); + + GGML_ASSERT(ptr <= end); + + ggml_build_forward_expand(graph, sink_blk); + ggml_build_forward_expand(graph, s_lat_t); + + auto plan = ggml_graph_plan(graph, std::thread::hardware_concurrency()/2); + if (plan.work_size > work_data.size()) work_data.resize(plan.work_size); + plan.work_data = work_data.data(); + + auto status = ggml_graph_compute(graph, &plan); + if (status != GGML_STATUS_SUCCESS) throw std::runtime_error("Failed to compute openPangu param_sink tensors"); + + auto buft = ggml_backend_buffer_get_type(l.param_sink_kv->buffer); + auto sink_blk_name = std::string{"blk."} + std::to_string(il) + ".attn_param_sink_blk"; + auto s_lat_t_name = std::string{"blk."} + std::to_string(il) + ".attn_param_sink_lat_t"; + + l.param_sink_blk = materialize(sink_blk, l.computed_param_sink_blk, buft, sink_blk_name); + l.param_sink_lat_t = materialize(s_lat_t, l.computed_param_sink_lat_t, buft, s_lat_t_name); + + LLAMA_LOG_INFO("Computed %s as %d x %d of type %s and stored in buffer %s\n", + sink_blk_name.c_str(), (int)sink_blk->ne[0], (int)sink_blk->ne[1], + ggml_type_name(sink_blk->type), ggml_backend_buffer_name(l.param_sink_blk->buffer)); + LLAMA_LOG_INFO("Computed %s as %d x %d of type %s and stored in buffer %s\n", + s_lat_t_name.c_str(), (int)s_lat_t->ne[0], (int)s_lat_t->ne[1], + ggml_type_name(s_lat_t->type), ggml_backend_buffer_name(l.param_sink_lat_t->buffer)); + + ggml_graph_clear(graph); + } + + ggml_free(ctx); +} + // Fold the 64-block Hadamard into wv_b/wk_b_pp at init; build_deepseek2.cpp then // skips the runtime cache_nope un-Hadamard. Math identity by H^T H = I. static void llm_apply_khad_pretransform(llama_model & model) { @@ -3159,7 +3571,7 @@ struct expert_tensors { }; static std::pair, double> get_layer_sizes(const llama_model_loader & ml, const llama_model & model, - ggml_type cache_type_k, ggml_type cache_type_v, uint32_t max_ctx_size, int mla_attn, int n_seq_max, int n_ubatch, + ggml_type cache_type_k, ggml_type cache_type_v, ggml_type idx_type_k, uint32_t max_ctx_size, int mla_attn, int n_seq_max, int n_ubatch, int amb, int worst_case_tokens, bool flash_attn, std::vector & experts) { int n_layer = model.hparams.n_layer; @@ -3367,7 +3779,7 @@ static std::pair, double> get_layer_sizes(const llama_model_ LLAMA_LOG_INFO("------------------- Layer sizes:\n"); double tot_model = 0, tot_cache = 0, max_compute = 0; for (int il = 0; il < n_layer; ++il) { - auto kv_size = model.cache_size(il, cache_type_k, cache_type_v, max_ctx_size, mla_attn, n_seq_max, flash_attn); + auto kv_size = model.cache_size(il, cache_type_k, cache_type_v, idx_type_k, max_ctx_size, mla_attn, n_seq_max, flash_attn); LLAMA_LOG_INFO("Layer %2d: %9.2f, %9.2f, %9.2f %9.2f MiB\n", il, result[il]/1024./1024., kv_size/1024./1024., (result[il] + kv_size)/1024./1024., compute[il]/1024./1024.); max_compute = std::max(max_compute, compute[il]); tot_model += result[il]; @@ -3398,6 +3810,7 @@ static bool llm_load_tensors( const float * tensor_split, ggml_type cache_type_k, ggml_type cache_type_v, + ggml_type idx_type_k, ggml_type extra_output_type, uint32_t max_ctx_size, int n_seq_max, @@ -3604,7 +4017,7 @@ static bool llm_load_tensors( // placement is already handled above, so skipping this block is safe. if (device_count > 0 && !model.devices.empty()) { std::vector experts; - auto [layer_sizes, max_compute] = get_layer_sizes(ml, model, cache_type_k, cache_type_v, max_ctx_size, mla_attn, n_seq_max, n_ubatch, + auto [layer_sizes, max_compute] = get_layer_sizes(ml, model, cache_type_k, cache_type_v, idx_type_k, max_ctx_size, mla_attn, n_seq_max, n_ubatch, amb, worst_case_tokens, flash_attn, experts); size_t required_mem = 0; for (int i = 0; i <= n_layer; ++i) { @@ -4072,6 +4485,9 @@ static bool llm_load_tensors( llm_prepare_mla(model, mla_attn); } } + if (!dry_run && model.arch == LLM_ARCH_OPENPANGU) { + llm_prepare_openpangu_param_sinks(model); + } if (model.arch == LLM_ARCH_GEMMA4) { llm_scale_gate_inp_s(model, use_mmap_buffer); } @@ -4190,6 +4606,15 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam } catch(const std::exception & e) { throw std::runtime_error("error loading model hyperparameters: " + std::string(e.what())); } + if (model.arch == LLM_ARCH_OPENPANGU) { + std::string error_msg; + if (!llama_openpangu_validate_latent_cache_types(params.type_k, params.type_v, &error_msg)) { + throw std::runtime_error(error_msg); + } + if (!llama_openpangu_validate_indexer_cache_type(params.idx_type_k, &error_msg)) { + throw std::runtime_error(error_msg); + } + } if (params.defer_experts && params.use_mmap) { #ifdef __linux__ ml.build_expert_tensor_index(model.hparams); @@ -4220,7 +4645,7 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam if (!llm_load_tensors( ml, model, params.n_gpu_layers, params.mla, params.split_mode, params.main_gpu, params.max_gpu, params.tensor_split, - params.type_k, params.type_v, params.extra_output_type, + params.type_k, params.type_v, params.idx_type_k, params.extra_output_type, params.max_ctx_size, params.n_seq_max, params.n_ubatch, params.amb, params.fit_margin, params.fit_margin_array, params.worst_graph_tokens, params.flash_attn, params.use_mlock, params.validate_quants, params.mtp, params.fit, params.dry_run, @@ -4394,6 +4819,23 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { #endif } + if (lctx.inp_mtp_carry) { + const size_t n_floats = (size_t) ggml_nelements(lctx.inp_mtp_carry); + if (lctx.mtp_carry_pending) { + // retire the async readback issued after the previous warmup/update before + // touching the host buffer + ggml_backend_sched_synchronize(lctx.sched); + lctx.mtp_carry_pending = false; + } + if (lctx.cparams.mtp_op_type == MTP_OP_WARMUP && batch.pos && batch.pos[0] == 0) { + // fresh prompt warmup: no committed history for the deeper heads yet + lctx.mtp_carry.assign(n_floats, 0.0f); + } else if (lctx.mtp_carry.size() < n_floats) { + lctx.mtp_carry.resize(n_floats, 0.0f); + } + ggml_backend_tensor_set(lctx.inp_mtp_carry, lctx.mtp_carry.data(), 0, n_floats*sizeof(float)); + } + if (lctx.inp_pos && lctx.inp_scale) { #if IK_PRINT_TIMING == 2 auto tim1 = ggml_time_us(); @@ -4458,7 +4900,7 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { "causal attention is not supported by this model" ); - if (lctx.inp_KQ_mask || lctx.inp_KQ_mask_swa) { + if (lctx.inp_KQ_mask || lctx.inp_KQ_mask_swa || lctx.inp_KQ_mask_swa_win) { #if IK_PRINT_TIMING == 2 auto tim1 = ggml_time_us(); #endif @@ -4469,14 +4911,21 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { ? lctx.mtp_target_ctx->kv_self : kv_self; const int64_t n_kv = mask_kv_self.n; + // openPangu: the NextN/MTP layers carry their own sliding window, and MTP graphs + // run with mtp_op_type set, so the effective window is per-context. + const uint32_t n_swa_eff = lctx.model.arch == LLM_ARCH_OPENPANGU && + cparams.mtp_op_type != MTP_OP_NONE && hparams.n_swa_mtp > 0 + ? hparams.n_swa_mtp : hparams.n_swa; if (cparams.causal_attn && !lctx.is_encoding) { const int64_t n_tokens = batch.n_tokens; float * data = nullptr; float * data_swa = nullptr; + float * data_swa_win = nullptr; ggml_half * data_f16 = nullptr; ggml_half * data_swa_f16 = nullptr; + ggml_half * data_swa_win_f16 = nullptr; if (lctx.inp_KQ_mask) { GGML_ASSERT(ggml_backend_buffer_is_host(lctx.inp_KQ_mask->buffer)); @@ -4496,7 +4945,76 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { } } - auto noalibi_f16 = [&mask_kv_self, &hparams, n_kv, data_f16, data_swa_f16] (int j, llama_pos pos, llama_seq_id seq_id, int first, int last) { + if (lctx.inp_KQ_mask_swa_win) { + GGML_ASSERT(ggml_backend_buffer_is_host(lctx.inp_KQ_mask_swa_win->buffer)); + if (cparams.flash_attn) { + data_swa_win_f16 = (ggml_half *) lctx.inp_KQ_mask_swa_win->data; + } else { + data_swa_win = (float *) lctx.inp_KQ_mask_swa_win->data; + } + } + + if (data_swa_win || data_swa_win_f16) { + const auto & built = lctx.openpangu_swa_window_view; + const uint32_t pad = llama_kv_cache_get_padding(cparams); + const llama_openpangu_swa_window_view view = + llama_openpangu_calc_swa_window_view(n_kv, n_tokens, built.window, pad); + GGML_ASSERT(lctx.model.arch == LLM_ARCH_OPENPANGU && + "windowed SWA mask input is only valid for OpenPangu"); + GGML_ASSERT(built.active && view.engaged && + "openPangu SWA window view must be engaged when KQ_mask_swa_win is present"); + GGML_ASSERT(built.n_kv == n_kv && built.n_tokens == n_tokens && + built.pad == pad && built.w_view == view.w_view && + built.win_off == view.win_off && + "openPangu SWA window view reuse-key mismatch"); + + const int64_t W_view = built.w_view; + const int64_t win_off = built.win_off; + for (int j = 0; j < n_tokens; ++j) { + const llama_pos pos = batch.pos[j]; + const llama_seq_id seq_id = batch.seq_id[j][0]; + for (int64_t c = 0; c < W_view; ++c) { + const int64_t i = win_off + c; + float f; + if (!mask_kv_self.cells[i].has_seq_id(seq_id) || mask_kv_self.cells[i].pos > pos) { + f = -INFINITY; + } else { + f = hparams.use_alibi ? -std::abs(mask_kv_self.cells[i].pos - pos) : 0.0f; + } + + if (f > -INFINITY) { + if (hparams.n_attn_chunk) { + llama_pos pos_chunk_start = (pos / hparams.n_attn_chunk) * hparams.n_attn_chunk; + if (mask_kv_self.cells[i].pos < pos_chunk_start || pos < pos_chunk_start) { + f = -INFINITY; + } + } else if (pos - mask_kv_self.cells[i].pos >= (int32_t) built.window) { + f = -INFINITY; + } + } + + if (data_swa_win) { + data_swa_win[j*W_view + c] = f; + } + if (data_swa_win_f16) { + data_swa_win_f16[j*W_view + c] = ggml_fp32_to_fp16(f); + } + } + } + + const int64_t n_tokens_padded = GGML_PAD(n_tokens, GGML_KQ_MASK_PAD); + if (n_tokens_padded > n_tokens) { + if (data_swa_win) { + std::fill(data_swa_win + int64_t(n_tokens)*W_view, data_swa_win + n_tokens_padded*W_view, -INFINITY); + } + if (data_swa_win_f16) { + const ggml_half h_inf = ggml_fp32_to_fp16(-INFINITY); + std::fill(data_swa_win_f16 + int64_t(n_tokens)*W_view, data_swa_win_f16 + n_tokens_padded*W_view, h_inf); + } + } + } + + auto noalibi_f16 = [&mask_kv_self, &hparams, n_kv, n_swa_eff, data_f16, data_swa_f16] (int j, llama_pos pos, llama_seq_id seq_id, int first, int last) { ggml_half h_inf = ggml_fp32_to_fp16(-INFINITY); ggml_half h_zero = ggml_fp32_to_fp16(0.f); for (int i = first; i < last; ++i) { @@ -4510,7 +5028,7 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { h = h_inf; } } else { - if (pos - mask_kv_self.cells[i].pos >= (int32_t)hparams.n_swa) { + if (pos - mask_kv_self.cells[i].pos >= (int32_t)n_swa_eff) { h = h_inf; } } @@ -4523,7 +5041,7 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { if (n_kv >= 1024 && n_tokens >= 32) { int n_thread = std::max(1, int(std::thread::hardware_concurrency()/2)); int npt = (n_kv + n_thread - 1)/n_thread; - auto compute = [&batch, &mask_kv_self, &hparams, &cparams, &noalibi_f16, n_tokens, n_kv, npt, data, data_swa, data_f16, data_swa_f16] (int ith) { + auto compute = [&batch, &mask_kv_self, &hparams, &cparams, &noalibi_f16, n_tokens, n_kv, n_swa_eff, npt, data, data_swa, data_f16, data_swa_f16] (int ith) { int first = ith * npt; int last = std::min(int(n_kv), first + npt); if (last <= first) return; @@ -4564,7 +5082,7 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { f = -INFINITY; } } else { - if (pos - mask_kv_self.cells[i].pos >= (int32_t)hparams.n_swa) { + if (pos - mask_kv_self.cells[i].pos >= (int32_t)n_swa_eff) { f = -INFINITY; } } @@ -4644,7 +5162,7 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { f = -INFINITY; } } else { - if (pos - mask_kv_self.cells[i].pos >= (int32_t)hparams.n_swa) { + if (pos - mask_kv_self.cells[i].pos >= (int32_t)n_swa_eff) { f = -INFINITY; } } @@ -4734,7 +5252,7 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { if (mask_kv_self.cells[i].pos < pos_chunk_start || pos < pos_chunk_start) { h = h_inf; } - } else if (pos - mask_kv_self.cells[i].pos >= (int32_t)hparams.n_swa) { + } else if (pos - mask_kv_self.cells[i].pos >= (int32_t)n_swa_eff) { h = h_inf; } } @@ -4957,7 +5475,7 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) { int32_t * data = (int32_t *) lctx.inp_s_seq_qnext->data; for (int64_t j = 0; j < n_tokens; ++j) { - // qwen3next linear-attention path uses a single local recurrent state slot. + // qwen3next and openPangu use a single local recurrent state slot. data[j] = 0; } } @@ -5526,7 +6044,7 @@ static int llama_decode_internal( (int)u_batch.all_seq_id, (int)lctx.n_outputs, (int)kv_self_used.n, (int)u_batch.n_tokens, kv_self_used.save_per_step_ssm, kv_self_used.ckpt.per_step_max_allocated, - cparams.mtp_op_type, gf}); + cparams.mtp_op_type, lctx.mtp_step_idx, lctx.mtp_n_heads, gf}); } } else { //printf("Reusing graph with type = %d, n_kv = %d, n_tokens = %d\n", cparams.mtp_op_type, (int)prev->n_kv, (int)prev->n_tokens); @@ -5617,7 +6135,7 @@ static int llama_decode_internal( bool reset_previous = false; // update the kv ring buffer { - if (llama_model_has_recurrent(&lctx.model) && kv_self.head == 0) { + if ((llama_model_has_recurrent(&lctx.model) || llama_model_is_openpangu(&lctx.model)) && kv_self.head == 0) { reset_previous = true; } kv_self.head += n_tokens; @@ -5745,14 +6263,36 @@ static int llama_decode_internal( printf("get_embedding(...): %d us\n", int(tim2-tim1)); #endif } + // persist the multi-head MTP carries (each head's output at the last committed row) + // for the next warmup/update graph; must happen per ubatch so a multi-ubatch prompt + // warmup chains its heads across ubatch boundaries. The copy is issued async on the + // backend stream (ordered before any later graph can overwrite the source buffer) + // and synchronized lazily when the host buffer is next consumed, so the update + // decode does not stall on a device readback. + if (cparams.mtp_op_type == MTP_OP_WARMUP || cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) { + if (ggml_tensor * carry_out = ggml_graph_get_tensor(gf, "mtp_carry_out"); carry_out != nullptr) { + const size_t n_floats = (size_t) ggml_nelements(carry_out); + if (lctx.mtp_carry_pending) { + ggml_backend_sched_synchronize(lctx.sched); + lctx.mtp_carry_pending = false; + } + if (lctx.mtp_carry.size() < n_floats) { + lctx.mtp_carry.resize(n_floats, 0.0f); + } + ggml_backend_t backend_carry = ggml_backend_sched_get_tensor_backend(lctx.sched, carry_out); + GGML_ASSERT(backend_carry != nullptr); + ggml_backend_tensor_get_async(backend_carry, carry_out, lctx.mtp_carry.data(), 0, n_floats*sizeof(float)); + lctx.mtp_carry_pending = true; + } + } + n_outputs_prev += lctx.n_outputs; n_outputs_prev_embd += (has_mtp && embd) ? embd->ne[1] : lctx.n_outputs; cur_token += n_tokens; if (reset_previous) { - // We need to discard this graph. Otherwise, iwith CUDA graphs enabled, the graph will get resused and this will reset the - // recurrent state for each new token. This is probably not very relevant in practice because we basically never run TG with - // empty context, but for the sake of correctness let's just do it. - lctx.prev.reset(); + // The graph was built with a position-zero state reset. Reusing it at later + // positions would reset the recurrent state again, so discard the active slot. + prev.reset(); } if (stop_internal_decode) { return -3; @@ -6969,12 +7509,40 @@ struct llama_context * llama_init_from_model( params.flash_attn = false; } + if (params.flash_attn && model->arch == LLM_ARCH_OPENPANGU) { + // openPangu attention prepends 128 learned param_sink K/V entries, which the manual + // (soft_max) attention path handles; the FA kernel has no way to include them. + LLAMA_LOG_WARN("%s: flash_attn is not compatible with OpenPangu param_sink attention - forcing off\n", __func__); + params.flash_attn = false; + } + + if (params.n_seq_max > 1 && model->arch == LLM_ARCH_OPENPANGU) { + // The MoME conv slot and the DSA indexer cache are single-sequence (one position + // stream); parallel sequences would silently share and corrupt them. + LLAMA_LOG_ERROR("%s: OpenPangu supports a single sequence only (requested n_seq_max = %u); run with -np 1\n", + __func__, params.n_seq_max); + return nullptr; + } + + if (model->arch == LLM_ARCH_OPENPANGU) { + std::string error_msg; + if (!llama_openpangu_validate_latent_cache_types(params.type_k, params.type_v, &error_msg)) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, error_msg.c_str()); + return nullptr; + } + if (!llama_openpangu_validate_indexer_cache_type(params.idx_type_k, &error_msg)) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, error_msg.c_str()); + return nullptr; + } + } + //if (params.flash_attn && model->hparams.n_embd_head_k != model->hparams.n_embd_head_v) { // LLAMA_LOG_WARN("%s: flash_attn requires n_embd_head_k == n_embd_head_v - forcing off\n", __func__); // params.flash_attn = false; //} - if (params.type_v != GGML_TYPE_F16 && params.type_v != GGML_TYPE_BF16 && !params.flash_attn) { + if (model->arch != LLM_ARCH_OPENPANGU && + params.type_v != GGML_TYPE_F16 && params.type_v != GGML_TYPE_BF16 && !params.flash_attn) { LLAMA_LOG_ERROR("%s: V cache quantization requires flash_attn\n", __func__); return nullptr; } @@ -7176,6 +7744,7 @@ struct llama_context * llama_init_from_model( model->arch != LLM_ARCH_QWEN35MOE && model->arch != LLM_ARCH_GEMMA4 && model->arch != LLM_ARCH_GEMMA4_MTP && model->arch != LLM_ARCH_GLM_DSA && model->arch != LLM_ARCH_GEMMA4_ASSISTANT && + model->arch != LLM_ARCH_OPENPANGU && cparams.mtp != 0) { cparams.mtp = 0; } @@ -7419,7 +7988,6 @@ struct llama_context * llama_init_from_model( memory_size_k_indexer += ggml_nbytes(k); } } - for (auto & v : ctx->kv_self.v_l) { if (v) { memory_size_v += ggml_nbytes(v); @@ -7693,6 +8261,7 @@ enum llama_rope_type llama_rope_type(const struct llama_model * model) { return LLAMA_ROPE_TYPE_NORM; // the pairs of head values are offset by n_rot/2 + case LLM_ARCH_OPENPANGU: // rope_interleave=false -> rotate_half (Infer: is_neox_style = not rope_interleave) case LLM_ARCH_FALCON: case LLM_ARCH_GROK: case LLM_ARCH_DBRX: @@ -8103,6 +8672,15 @@ void llama_kv_cache_clear(struct llama_context * ctx) { // Unified speculative-checkpoint static bool spec_ckpt_try_per_step(llama_kv_cache & kv, const llama_model & model, int max_tokens) { + // openPangu carries only a conv state (no SSM recurrent term), so the per-step + // path - which sizes itself from the ssm_* hparams (ssm_dt_rank etc, all zero + // here) - does not apply. Decline it so the checkpoint resolves to the whole-slot + // shadow (gpu-fallback), which is arch-agnostic. + if (model.arch == LLM_ARCH_OPENPANGU) { + kv.save_per_step_ssm = false; + return false; + } + // Split recurrent tensors are supported as long as each layer exposes // concrete backend buffers for the per-step tensors. CPU-only and mixed // CPU/GPU recurrent placement are also allowed. @@ -8567,7 +9145,7 @@ struct llama_data_write { // Iterate and write all the keys first, each row is a cell // Get whole range at a time for (uint32_t il = 0; il < n_layer; ++il) { - const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il) + hparams.n_embd_k_s(); + const uint32_t n_embd_k_gqa = llama_kv_k_row_embd(ctx->model, hparams, il); const uint32_t n_embd_head_qk_rope = hparams.n_rot; const uint32_t kv_lora_rank = hparams.n_lora_kv; const bool has_k_cache = kv_self.k_l[il] != nullptr && need_kv; @@ -8990,7 +9568,7 @@ struct llama_data_read { // For each layer, read the keys for each cell, one row is one cell, read as one contiguous block for (uint32_t il = 0; il < n_layer; ++il) { - const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il) + hparams.n_embd_k_s(); + const uint32_t n_embd_k_gqa = llama_kv_k_row_embd(ctx->model, hparams, il); const uint32_t n_embd_head_qk_rope = hparams.n_rot; const uint32_t kv_lora_rank = hparams.n_lora_kv; const bool has_k_cache = kv_self.k_l[il] != nullptr && need_kv; @@ -9491,6 +10069,18 @@ struct llama_data_read_file : llama_data_read { } }; +// openPangu: refuse state I/O outright until the side state is part of the format. K/V +// rows alone are not a complete snapshot; the MoME conv slot (s_l) and the DSA +// indexer cache are needed to resume a sequence, and restoring without them diverges +// silently instead of failing. +static bool llama_state_io_supported(const struct llama_context * ctx, const char * func) { + if (ctx->model.arch == LLM_ARCH_OPENPANGU) { + LLAMA_LOG_ERROR("%s: state save/restore is not supported for openPangu (conv slot and indexer cache are not serialized)\n", func); + return false; + } + return true; +} + /** copy state data into either a buffer or file depending on the passed in context * * file context: @@ -9505,6 +10095,9 @@ struct llama_data_read_file : llama_data_read { * */ static size_t llama_state_get_data_internal(struct llama_context * ctx, llama_data_write & data_ctx) { + if (!llama_state_io_supported(ctx, __func__)) { + return 0; + } llama_synchronize(ctx); data_ctx.write_model_info(ctx); @@ -9544,6 +10137,9 @@ size_t llama_state_get_size(struct llama_context * ctx) { } static size_t llama_state_set_data_internal(struct llama_context * ctx, llama_data_read & data_ctx) { + if (!llama_state_io_supported(ctx, __func__)) { + return SIZE_MAX; + } llama_synchronize(ctx); data_ctx.read_model_info(ctx); @@ -9573,6 +10169,9 @@ size_t llama_state_set_data(struct llama_context * ctx, const uint8_t * src, siz } static bool llama_state_load_file_internal(struct llama_context * ctx, const char * path_session, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { + if (!llama_state_io_supported(ctx, __func__)) { + return false; + } llama_file file(path_session, "rb"); // sanity checks @@ -9624,6 +10223,9 @@ bool llama_state_load_file(struct llama_context * ctx, const char * path_session } static bool llama_state_save_file_internal(struct llama_context * ctx, const char * path_session, const llama_token * tokens, size_t n_token_count) { + if (!llama_state_io_supported(ctx, __func__)) { + return false; + } llama_file file(path_session, "wb"); file.write_u32(LLAMA_SESSION_MAGIC); @@ -9650,6 +10252,9 @@ bool llama_state_save_file(struct llama_context * ctx, const char * path_session } static size_t llama_state_seq_get_data_internal(struct llama_context * ctx, llama_data_write & data_ctx, llama_seq_id seq_id, llama_state_seq_flags flags) { + if (!llama_state_io_supported(ctx, __func__)) { + return 0; + } llama_synchronize(ctx); data_ctx.write_kv_cache(ctx, seq_id, flags); @@ -9673,6 +10278,9 @@ size_t llama_state_seq_get_data(struct llama_context * ctx, uint8_t * dst, size_ } static size_t llama_state_seq_set_data_internal(struct llama_context * ctx, llama_data_read & data_ctx, llama_seq_id dest_seq_id, llama_state_seq_flags flags) { + if (!llama_state_io_supported(ctx, __func__)) { + return SIZE_MAX; + } llama_synchronize(ctx); data_ctx.read_kv_cache(ctx, dest_seq_id, flags); @@ -9691,6 +10299,9 @@ size_t llama_state_seq_set_data(struct llama_context * ctx, const uint8_t * src, } static size_t llama_state_seq_save_file_internal(struct llama_context * ctx, const char * filepath, llama_seq_id seq_id, const llama_token * tokens, size_t n_token_count) { + if (!llama_state_io_supported(ctx, __func__)) { + return 0; + } llama_file file(filepath, "wb"); file.write_u32(LLAMA_STATE_SEQ_MAGIC); @@ -9710,6 +10321,9 @@ static size_t llama_state_seq_save_file_internal(struct llama_context * ctx, con } static size_t llama_state_seq_load_file_internal(struct llama_context * ctx, const char * filepath, llama_seq_id dest_seq_id, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) { + if (!llama_state_io_supported(ctx, __func__)) { + return 0; + } llama_file file(filepath, "rb"); // version checks @@ -9877,6 +10491,14 @@ void llama_set_mtp_op_type(llama_context * ctx, llama_mtp_op_type mtp_op_type) { ctx->set_mtp_op_type(mtp_op_type); } +void llama_set_mtp_step_idx(llama_context * ctx, int32_t mtp_step_idx) { + ctx->set_mtp_step_idx(mtp_step_idx); +} + +void llama_set_mtp_n_heads(llama_context * ctx, int32_t mtp_n_heads) { + ctx->set_mtp_n_heads(mtp_n_heads); +} + void llama_synchronize(struct llama_context * ctx) { ggml_backend_sched_synchronize(ctx->sched);