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

deepseek2 : GLM-DSA sparse attention (lightning indexer), --dsa off by default (#2045)

Browse Source 
* Add GLM-5.2/DeepSeek-V3.2 DSA lightning indexer (batch-local, single-seq prefill)

Implements the sparse top-k "lightning indexer" attention for LLM_ARCH_GLM_DSA
in build_deepseek2_layer_attention (ik's deepseek2 graph).

What it does (per layer, gated on model.arch==GLM_DSA && indexer_attn_q_b):
- indexer_q = indexer_attn_q_b(q_lora latent), split rope(64)/nope(64), NEOX-rope
  the pe part, concat. indexer_k = indexer_attn_k(attn_norm out), LayerNorm w/ bias,
  same rope/concat (single key head, MQA).
- scores = relu(indexer_k . indexer_q), scaled per-head weights (indexer_proj),
  summed over heads, + base causal mask, then ggml_top_k(min(top_k, n_tokens)).
- sparse mask: ggml_fill(-inf) -> ggml_set_rows(0) at top_k positions -> + causal,
  used in the soft_max_ext attention path (-mla 1 -fa 0) instead of KQ_mask.

Simplifications (intentional, proven sound):
- Batch-local: no indexer KV-cache. Indexer keys are the current batch tokens.
- Walsh-Hadamard transform omitted: orthonormal rotation, (Hq).(Hk)==q.k, no score change.

Validation (GLM-5.2-UD-IQ2_M, 3x P100, -mla 1 -fa 0):
- Compiles clean (CUDA sm_60); loads and runs.
- c512 -b512 (n_seq=1) PPL = 2.7760, byte-identical to dense baseline (indexer
  disabled) = 2.7760, all 8 chunks match -> indexer is an exact no-op when
  top_k>=n_tokens. Proves correctness-preservation.
- 3105-token prompt completion (top_k=2048 < 3105 -> indexer ACTIVELY masks):
  prompt-eval produces coherent, accurate continuation, identical to dense for the
  prompt+early-gen tokens. No NaN/crash. Confirms the masking path works in prefill.

Known limitations (documented follow-ups, NOT handled):
- Single-sequence prefill only. Multi-sequence batches (n_seq>1, e.g. perplexity
  default n_batch>n_ctx) and kv_head>0 (decode) break the batch-local key->slot
  mapping. n_seq>1 -> NaN (use n_batch==n_ctx). Decode (kv_head>0): each generated
  token sees only itself as an indexer key, so generation degenerates into repetition
  after the prompt (dense A/B stays coherent) -- this is the decode-cache stub, the
  documented next step.
- Flash-attn path (-fa 1, F16 mask) still uses dense KQ_mask (soft_max path only).
- Decode indexer KV-cache + Hadamard cached-K storage not implemented.

Runtime gate: DSA_INDEXER_DISABLE=1 falls back to dense attention (for A/B).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-5.2 DSA indexer: decode-correct via persistent indexer-K cache

Make the lightning-indexer correct for DECODE (not just prefill). Previously the
indexer was batch-local, so a generated token only scored against itself and
generation degenerated. Now the indexer keys are cached across the full context.

Changes
- llama_kv_cache: add per-layer indexer-key cache `kr_l` [indexer_head_size, kv_size]
  (F16, MQA single head), allocated alongside the MLA latent cache for GLM_DSA.
- build_deepseek2_dsa_indexer: write the batch's (Hadamard-rotated) indexer keys to
  kr_l at kv_head, read back the full [128, n_kv] cached keys, and score the indexer
  queries against ALL past keys. Returns the full descending argsort of the scores.
- Walsh-Hadamard rotation of indexer q/k (cparams.dsa_indexer_hadamard, default on;
  filled in llama_set_inputs). Score-preserving; improves cached-K F16 precision.
- build_deepseek2_dsa_sparse_mask: rank-based full-coverage scatter (write a 0/-BIG
  penalty into EVERY key slot keyed by rank) instead of partial set_rows into a -inf
  fill — the CUDA in-place set_rows does not preserve an un-written base, which had
  corrupted decode when n_kv > top_k.
- Attention-sink force-inclusion (DSA_SINK, default 1): boost the first key(s) so the
  sink always survives top-k. The IQ2_M-quantized indexer under-ranks the sink, and
  masking it collapsed decode; with the boost, top_k=2048 over n_kv>2048 stays coherent.

ggml backend fixes (needed by the indexer)
- CUDA argsort: report unsupported when padded ncols > 1024 (one-thread-per-column
  bitonic launch limit) so the scheduler falls back to the CPU argsort. Fixes
  "invalid configuration argument" for top_k over a large n_kv.
- CUDA cpy/dup: support I32 -> I32 (top_k index copies / cross-backend moves).

Validation (GLM-5.2-UD-IQ2_M, 3xP100 + --cpu-moe, -mla 1 -fa 0)
- c512 PPL = 2.0743, byte-identical to dense (all 8 chunks): no-op path exact.
- Short-context decode (300 tok): coherent, identical to dense.
- Long-context decode (2521-tok prompt, n_kv>top_k, real masking of ~474 keys,
  120+ tok generated): coherent with the sink boost; dense A/B also coherent.

Gated behind arch==GLM_DSA + indexer tensors + kr_l cache; DSA_INDEXER_DISABLE=1
forces dense. Remaining: FA path still uses the dense KQ_mask; multi-sequence
(n_seq>1) batches; deepseek32 arch wiring.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-5.2 DSA indexer: wire sparse mask into the flash-attention path (-fa 1)

The DSA sparse top-k mask is now applied on the -fa 1 path (our serving config),
not just -fa 0 soft_max. c512 PPL on -fa 1 = 2.0743, byte-identical to dense
(no regression, indexer no-op exact at n_kv <= top_k). Gated arch==GLM_DSA with
DSA_INDEXER_DISABLE escape; -fa 0 path unchanged.

Long-context -fa 1 decode coherence (n_kv > top_k, mask actually biting) validation
is still running at commit time; the FA mask reuses the same full-coverage scatter
proven coherent on the -fa 0 decode path, so it should hold, but confirm before
relying on long-context -fa 1.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-5.2 DSA indexer: UPDATE 4 — MLA-FA fix merged, FA path validated, multi-seq characterized

Document the re-validation after cherry-picking the MLA-FA vec-decode fix (5f18dcc0):

- FA path is ALIVE. Long-ctx -fa 1 decode (2521-tok prompt > top_k, mask actively
  biting) is now COHERENT at -mla 1 and -mla 3, vs the pre-fix degeneration into
  "0.0.0.0..." repetition. Matches dense (DSA_INDEXER_DISABLE) and -fa 0 controls.
- c512 -fa 1 PPL: indexer-ON == dense == 2.0854, byte-identical all 8 chunks (exact
  no-op when n_kv <= top_k; no regression). The 2.0743->2.0854 shift is the MLA-FA
  fix changing V accumulation, not an indexer artifact (ON==dense proves it).
- Indexer is feature-complete + validated for single-seq prefill+decode on both
  -fa 0 and -fa 1, at -mla 1 and -mla 3 (the R740 serving target).

Remaining PR gaps, characterized honestly:
- Multi-seq (n_seq>1) with active mask is BROKEN (n_seq=2 c4096 PPL 62.6 vs dense
  multi-seq 2.54 and single-seq indexer 3.05). No NaN/crash anymore. Root cause:
  the indexer uses a single scalar kv_head/n_kv for the whole ubatch; multi-seq
  needs per-sequence cache writes + per-sequence top-k. Fix deferred (structural).
- deepseek32 arch: N/A in this fork. DSA lives entirely under LLM_ARCH_GLM_DSA;
  there is no LLM_ARCH_DEEPSEEK32 enum. Documented the steps to add one if a real
  deepseek32 GGUF is ever served.

Also commit DSA_REFERENCE.md (verbatim mainline deepseek32/glm-dsa source, the port
reference), trimmed of a stray agent-handoff footer.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-5.2 DSA indexer: per-sequence attention sink — fix multi-seq (n_seq>1)

UPDATE 5. The DSA lightning indexer was numerically broken for multi-sequence
batches once the top-k mask bites (n_kv > top_k): c4096 n_seq=2 PPL 62.6 vs
dense 2.54, while single-seq was fine. Root cause: the attention-sink
force-include boosted the GLOBAL key range [0, n_sink) by +1e20, which only
protects sequence 0's sink. With several sequences packed contiguously into one
ubatch (seq 0 at cells [0,n0), seq 1 at [n0,n1), ...), every non-first
sequence's sink lives at cell n0.. (not cell 0), got no boost, and was dropped
from top-k once the mask bites — collapsing that sequence (chunk[2]=61.2 while
chunk[1]=2.33).

The cache write and score/argsort were already per-sequence correct: tokens are
placed contiguously like the main K cache, and the base KQ_mask (filled from
kv_self.cells[i].has_seq_id) already drives cross-seq keys to -inf before
argsort. Only the sink was anchored at the wrong (global) cell.

Fix: replace the global arange sink boost with a per-graph input tensor
inp_dsa_sink {n_kv, n_tokens} (F32), filled on the CPU in llama_set_inputs from
kv_self.cells exactly like the KQ_mask:
  inp_dsa_sink[j,i] = 1e20 iff cell[i].pos in [0,n_sink) AND
                              cell[i].has_seq_id(seq_of_query_j), else 0
so each query force-includes only its OWN sequence's sink. For a single
contiguous sequence from pos 0 this is exactly the old "cell index < n_sink"
set with the same magnitude, so n_seq==1 is byte-identical.

Validation (3x P100, -ngl 99 --cpu-moe -mla 3 -fa 1, wikitext-2):
- c4096 n_seq=2 indexer chunk[2]: 61.2 -> 3.07 (== single-seq 3.05).
- c2048 topk=1024 (mask bites): n_seq=4 == n_seq=1 chunk-for-chunk
  (2.5005/2.6080/2.7759/3.1137 vs .../3.1138) -> multi-seq is numerically
  identical to processing each sequence alone.
- c512 n_seq=1 indexer ON == dense, all 4 chunks byte-identical (no regression).

n_seq=4 at full c4096 (n_kv=16384) OOMs the P100 compute buffer (capacity, not
correctness; n_seq=4 proven correct at c2048/n_kv=8192).

GLM-5.2 DSA indexer is now sequence-correct for n_seq>=1, prefill+decode,
soft_max+FA, -mla 1/-mla 3. Fully general and PR-ready.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-5.2 DSA indexer: UPDATE 6 — serving-correctness (kr_l maintained across shift/defrag/seq-ops; per-seq sink on first-present pos)

An adversarial review found the indexer was proven on the perplexity path but
not the serving path: the persistent indexer-K cache kr_l was written/read but
never *maintained* by the KV-cache mutators, and the attention sink anchored on
absolute pos<n_sink (wrong after multi-turn seq_rm). This closes those gaps and
pins down what is actually reachable on the MLA model.

kr_l maintenance:
- build_k_shift (llama-build-context.cpp): rotate the indexer keys by the same
  per-cell delta as the main K. The cached key is H*concat(RoPE(k_pe,pos),k_nope),
  so un-Hadamard (H sym/orthonormal => H*H=I) -> RoPE-delta the pe sub-block ->
  re-Hadamard. Exact because GLM-DSA has no rope-scaling metadata (ext_factor=0,
  attn_factor=1, freq_scale=1), so NEOX RoPE is pure/composable. Params mirror the
  forward indexer RoPE exactly (rope_factors=nullptr); no DEEPSEEK2 yarn-shift leak.
  Non-in-place (cont->rope->concat->re-Had->cpy), no aliasing. K-shift Hadamard
  input filled in llama_set_k_shift with the identical Sylvester construction.
- build_defrag: kr_l row-move mirrors the k_l move (defrag never changes pos, so
  no re-RoPE). max_moves divisor 6->9 *n_layer when the indexer cache is present.
- seq_rm/seq_cp/seq_keep are metadata-only (verified) so kr_l rows stay matched to
  cells; seq_add/seq_div set has_shift and route through K-shift. No seq-op change.

Per-seq sink (llama.cpp llama_set_inputs): anchor on each sequence's FIRST PRESENT
pos (min present pos over the scored n_kv span), not absolute pos<n_sink. After
multi-turn seq_rm drops a sequence's early tokens its earliest survivor has
pos>=n_sink; the absolute test would protect nothing. Fresh seq at pos 0 => min=0
=> byte-identical to the old behaviour.

Serving-shift finding (the whole point): a RoPE context-shift on this model is
REFUSED BY THE ENGINE. get_can_shift() returns false for all MLA models
(is_mla_model() includes GLM_DSA); llama_kv_cache_update returns 1 ->
"main : failed to eval". Reproduced AND isolated with a dense control
(DSA_INDEXER_DISABLE=1): dense fails identically at the same token. The failure is
pre-existing MLA engine behaviour, independent of the indexer. On the MLA path the
shift never happens, so the indexer's kr_l can never desync via K-shift; the
build_k_shift kr_l block is correct-and-dormant (documented loudly in code).

Validation (3x P100, -ngl 99 --cpu-moe -mla 3 -fa 1, GGML_CUDA_NO_PINNED=1,
numactl --interleave=all, wikitext-2):
- No regression: c512 n_seq=1 indexer ON == dense == 2.1957 +/- 0.12031,
  byte-identical all 4 chunks (2.2770/2.8741/2.3956/2.1957).
- Multi-seq: c4096 n_seq=2 chunk[1]=2.33 chunk[2]=3.07 healthy (== UPDATE 5;
  per-seq sink change did not regress).
- Serving shift: engine-refused for MLA, dense control fails identically.
- Independent adversarial review: GO, no correctness defect in the diff.
- Build clean (llama-cli, llama-perplexity, sm_60).

Comments updated (build_deepseek2.cpp): multi-seq+FA no longer limitations; sink
description matches per-seq min-pos anchoring; BIG=1e30 masks on both soft_max and
FA paths.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-5.2 DSA indexer: UPDATE 7 — FIX latent graph-reuse cache-fixup omission for the kr_l indexer cache

update_cache_copies() re-points the K/V cache writes to the current kv_head whenever a
compute graph is REUSED (can_reuse_graph reuses iff kv_self.n == prev->n_kv). The persistent
indexer-key cache write (kr_l) is a separate ggml_cpy whose destination view bakes kv_head at
graph-build time, and it was NEVER registered for that fixup. Under FA the cache pads to 256,
so consecutive single-token decode ubatches share the same padded n_kv and the graph IS reused;
without the fixup the kr_l write keeps landing in the first ubatch's slot and later ubatches
never populate their own recent index-key cells (those cells stay at the alloc-zeroed 0.0).
Structurally identical to the MiniMax MSA bug (fork commit 133d14c9).

Fix (mirrors the K/V cache_copies fixup, same shape as MSA 133d14c9):
- llama-context.h: new std::vector<CacheCopy> dsa_cache_copies.
- llama.cpp ctor: resize dsa_cache_copies to n_layer (null entries -> no-op when DSA off).
- build_deepseek2.cpp: register the kr_l ggml_cpy as dsa_cache_copies[il] = {kr_cpy, kr->nb[1]}.
- llama.cpp update_cache_copies(): re-point each registered cpy view_offs = kv_head*step and
  patch src[1]->data/data, exactly like K/V, with the c.cpy->view_src == kv_self.kr_l[il]
  (+ null/op) guard the MSA fix omitted. soft_max / non-DSA paths byte-identical.

Validation (GLM-5.2-UD-IQ2_M, 3x P100 -ngl 99 --cpu-moe -t 32, NO_PINNED, P2P-disable patch
re-applied to get a working multi-GPU baseline — see UPDATE 7.3; that patch was lost in the
upstream rebase and is required separately):
- c512 -fa1 -mla3 indexer ON: 2.1983 (== prior baseline; build healthy).
- Long-ctx FA decode, 2735-tok recall prompt, -mla3 -fa1 temp0, reuse ON (default): coherent,
  correct deep-context recall ("Dr. Mariana Velasquez ... Daniel Okonkwo") on BOTH the fixed and
  the unfixed binary.
- ub128 PPL -fa1 -mla3 reuse ON, unfixed: 1.7239/1.8211/2.1888/2.4517, healthy (no inflation).

Honest scope: the bug is real in code but LATENT for GLM-DSA at its configured top_k=2048
(permissive selection keeps the genuinely-attended recent blocks even when reuse leaves some
recent index-key cells stale), unlike MSA's tighter top-k where it inflated PPL ~2x. The fix is
correct and prevents the latent corruption from biting at any tighter top_k / longer ctx /
future serving config. The pre-P2P-patch "nan" seen at ub128 was P2P corruption, not this bug.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-DSA: convert sparse-attention control from env vars to CLI args (off by default)

Implements ikawrakow's direction from discussion #2040: the DSA sparse
indexer must be controllable via command-line argument (not environment
variables), and must be OFF by default for now.

Control surface, before -> after:
  DSA_INDEXER_DISABLE (env, inverted: on-by-default)  -> --dsa / -dsa
      (cparams.dsa, default false; opt-in, dense-by-default)
  DSA_TOPK_OVERRIDE   (env)                            -> --dsa-top-k N / -dsatk N
      (cparams.dsa_top_k, default -1 == model's configured indexer_top_k)
  DSA_HADAMARD_DISABLE, DSA_SINK (env)                 -> kept as DEBUG-ONLY env
      knobs (clearly commented; no CLI surface, not system on/off controls)

Plumbing mirrors existing boolean/int feature flags (-mla, -khad):
  include/llama.h        llama_context_params {bool dsa; int dsa_top_k;}
  src/llama.cpp          default_params (false / -1); cparams assignment
  src/llama-cparams.h    llama_cparams {bool dsa=false; int dsa_top_k=-1;}
  common/common.h        gpt_params {bool dsa=false; int dsa_top_k=-1;}
  common/common.cpp      arg parse + help text + cparams copy
  src/graphs/build_deepseek2.cpp  gate now checks cparams.dsa instead of
      getenv; top-k override reads cparams.dsa_top_k. Stays arch-gated to
      LLM_ARCH_GLM_DSA. When --dsa is off (default) the indexer function is
      never called -> existing dense MLA path, byte-identical to no-feature.

Validation (GLM-5.2-UD-IQ2_M, 3x P100, -ngl 99 --cpu-moe -mla 3 -fa 1,
wikitext-2, 4 chunks @ c2560):
  --dsa OFF (default, dense):              PPL 2.4151  (graph nodes 4166)
  --dsa ON, default top_k=2048:            PPL 2.4697  (graph nodes 8846)
  --dsa ON, --dsa-top-k 1024:              PPL 3.5107
Off-by-default runs the dense path; ON activates the indexer (node count
jumps, PPL shifts as the top-k mask bites once n_kv > top_k). No env var
is consulted for the primary on/off or the top-k knob.

Graph-parallel (-sm graph) interaction (the item ikawrakow flagged):
Under -sm graph the MLA layers are TP-split (wo->extra) and route to
build_deepseek2_tp_attention(), which contains NO indexer code. So --dsa
is silently a NO-OP under -sm graph: it does not error or crash, it runs
dense. Empirically, --dsa --dsa-top-k 1024 under -sm graph gives
PPL 2.4308 (chunks 1.6967/1.7906/2.1664/2.4308) -- the dense baseline
(2.4151), NOT the DSA top_k=1024 numbers (3.5107). The 0.016 delta is
f16 TP-reduce numerics, not DSA. Conclusion: DSA "works under deepseek2"
only on the non-TP (layer) path; serving DSA with -sm graph would require
wiring the indexer into the TP attention path (or a dedicated DSA arch).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-DSA: warn that --dsa is inactive under -sm graph/attn (TP path runs dense MLA)

The DSA lightning indexer is built only in the layer-mode (non-TP) attention
path. Under -sm graph / -sm attn the tensor-parallel attention path has no
indexer, so --dsa would silently run dense MLA. Emit a clear one-time
LLAMA_LOG_WARN at context creation instead of degrading silently.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-DSA: drop in-tree dev reference docs from the PR branch

DSA_REFERENCE.md and the R740 progress note are development scratch, not
part of the submission. Remove them so the PR diff is code-only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* GLM-DSA: fix CPU-only crashes in the sparse-attention path

PR #2045 adds GLM-DSA sparse attention but was validated on CUDA (--cpu-moe).
A CPU-only build (-ngl 0 --dsa) crashes in four spots where the CUDA backend
tolerates something the CPU backend does not. These make GLM-5.2 --dsa run
coherently on CPU; with --dsa off they are no-ops (DSA CPU path only).

1. set_rows into an F32 dest segfaults (ggml.c set_rows_f32):
   type_traits[F32].from_float is NULL, so the DSA sparse-mask scatter calls a
   NULL fn (segfault at ip=0). memcpy when the dest is F32. CUDA has a real F32
   set_rows path, so this only bit the CPU build.

2. ggml_add(F32 score, F16 mask) aborts on CPU (build_deepseek2_dsa_indexer and
   build_deepseek2_dsa_sparse_mask): under -fa 1 the dense KQ_mask is F16 and CPU
   add only accepts F32+F16 when src0 is F16. Cast the causal mask view to F32.
   CUDA's add accepts the mixed types.

3. dsa_fa_mask dim-1 concat must be F32 on CPU (build_deepseek2_dsa_fa_mask):
   CPU ggml_concat only supports F16 along dim 0; do the row (dim-1) concat in
   F32 then cast the result to F16. CUDA supports the F16 dim-1 concat.

4. indexer k_norm epsilon is 0 -> ggml_norm aborts (llama-hparams.cpp): the
   lightning-indexer k_norm is a non-RMS LayerNorm using f_norm_eps, but the
   GLM-DSA GGUF only carries the RMS eps so f_norm_eps stays 0
   (GGML_ASSERT(eps > 0)). Mirror the RMS eps. CUDA's norm doesn't assert on eps=0.

Validated: GLM-5.2 UD-Q4_K_M, single-socket Xeon w7-2475X, CPU-only (-ngl 0 --dsa)
- coherent at 49K+ ctx, correct 30K needle retrieval, prefill flat with length
(~32 tok/s, the O(L) DSA signature) vs the dense build's O(L^2) decline.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* DSA: loop over attention heads + use builtin Hadamard

* DSA: ggml_blend

* DSA: remove a bunch of unnecessary ggml_cont

* DSA: fix CUDA blend - but something is still wrong

* DSA: use ggml_top_k instead of ggml_argsort when FA is ON

* CUDA: add CUB based argsort

* DSA: avoid graph leaves

* Various

* GLM-5.2 DSA: IndexShare (shared layers reuse full-layer top-k)

GLM-5.2's indexer_types marks 21 'full' layers that compute their own
lightning-indexer top-k and 57 'shared' layers that reuse the previous
full layer's top-k. This port computed an independent top-k on every
layer, which mis-selects keys on the 57 shared layers (the transformers
reference sets indexer=None on shared layers and reuses prev_topk).

Shared layers now reuse the most-recent full layer's selection. Full/
shared map derived from the config rule (full iff il<=1 or il%4==2),
which reproduces indexer_types exactly; loader can later override from
GGUF metadata. Built on #2063's tree; head-loop/ggml_hadamard/ggml_blend/
argsort/FA-mask unchanged.

4K PPL (unsloth IQ2_M, top_k 2048, CPU): DSA-on 3.1922 -> 2.7111, dense
2.6972 (~97% of the gap). top_k>=n_kv reproduces dense exactly. Single-
seq and 4x8 parallel decode coherent.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Apply suggestion from @ikawrakow

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: mgkwill <168222+mgkwill@users.noreply.github.com>
Co-authored-by: Kawrakow <iwankawrakow@gmail.com>
This commit is contained in:
mb8565
2026-07-02 09:36:49 +02:00
committed by GitHub
co-authored by Claude Opus 4.8 mgkwill Kawrakow
parent 068b173649
commit 9da563d09d
21 changed files with 1190 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/common.cpp
+13
View File
@@ -1892,6 +1892,15 @@ bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_pa
params.mla_attn = std::stoi(argv[i]);
return true;
}
if (arg == "-dsa" || arg == "--dsa") {
params.dsa = true;
return true;
}
if (arg == "-dsatk" || arg == "--dsa-top-k") {
CHECK_ARG
params.dsa_top_k = std::stoi(argv[i]);
return true;
}
if (arg == "-amb" || arg == "--attention-max-batch") {
CHECK_ARG
params.attn_max_batch = std::stoi(argv[i]);
@@ -3011,6 +3020,8 @@ void gpt_params_print_usage(int /*argc*/, char ** argv, const gpt_params & param
options.push_back({ "*", "-no-fa, --no-flash-attn", "disable Flash Attention (default: %s)", params.flash_attn ? "enabled" : "disabled" });
options.push_back({ "*", "-fa, --flash-attn (auto|on|off|0|1)", "set Flash Attention (default: %s)", params.flash_attn ? "on" : "off" });
options.push_back({ "*", "-mla, --mla-use", "enable MLA (default: %d)", params.mla_attn });
options.push_back({ "*", "-dsa, --dsa", "enable GLM DSA sparse attention (GLM-DSA arch only; default: %s)", params.dsa ? "enabled" : "disabled" });
options.push_back({ "*", "-dsatk, --dsa-top-k", "DSA top-k override; <0 uses the model's configured indexer_top_k (default: %d)", params.dsa_top_k });
options.push_back({ "*", "-amb, --attention-max-batch", "max batch size for attention computations (default: %d)", params.attn_max_batch});
options.push_back({ "*", "-no-fmoe, --no-fused-moe", "disable fused MoE (default: %s)", params.fused_moe_up_gate ? "enabled" : "disabled" });
options.push_back({ "*", "-ger, --grouped-expert-routing", "enable grouped expert routing (default: %s)", params.grouped_expert_routing ? "enabled" : "disabled" });
@@ -4256,6 +4267,8 @@ struct llama_context_params common_context_params_to_llama(const gpt_params & pa
cparams.fused_mmad = params.fused_mmad;
cparams.rope_cache = params.rope_cache;
cparams.graph_reuse = params.graph_reuse;
cparams.dsa = params.dsa;
cparams.dsa_top_k = params.dsa_top_k;
cparams.k_cache_hadamard = params.k_cache_hadamard;
cparams.v_cache_hadamard = params.v_cache_hadamard;
cparams.split_mode_graph_scheduling = params.split_mode_graph_scheduling;
common/common.h
+2
View File
@@ -417,6 +417,8 @@ struct gpt_params {
bool grouped_expert_routing = false; // if to use grouped expert routing (BailingMoeV2 arch)
bool rope_cache = false; // if to use RoPE cache (for supported models)
bool graph_reuse = true; // if to reuse compute graphs
bool dsa = false; // enable GLM DSA sparse attention (off by default; opt-in via --dsa)
int dsa_top_k = -1; // DSA top-k override (<0 => use the model's configured indexer_top_k)
int min_experts = -1;
float thresh_experts = 0;
ggml/include/ggml.h
+9
View File
@@ -703,6 +703,7 @@ extern "C" {
GGML_OP_FAKE_CPY,
GGML_OP_FUSED_NORM,
GGML_OP_FUSED_RMS_RMS_ADD,
GGML_OP_BLEND,
GGML_OP_COUNT,
};
@@ -2393,6 +2394,14 @@ extern "C" {
struct ggml_tensor * a,
float c);
// Overwrite values in a with c for the indeces stored in b
GGML_API struct ggml_tensor * ggml_blend(
struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b,
float c);
// sort rows
enum ggml_sort_order {
GGML_SORT_ORDER_ASC,
ggml/src/ggml-cuda.cu
+23 -2
View File
@@ -56,6 +56,7 @@
#include "ggml-cuda/reduce.cuh"
#include "ggml-cuda/tri.cuh"
#include "ggml-cuda/delta-net.cuh"
#include "ggml-cuda/blend.cuh"
#include <algorithm>
#include <array>
@@ -3640,6 +3641,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
case GGML_OP_REDUCE:
ggml_cuda_op_reduce(ctx, dst);
break;
case GGML_OP_BLEND:
ggml_cuda_op_blend(ctx, dst);
break;
case GGML_OP_FAKE_CPY:
break;
case GGML_OP_ARGMAX:
@@ -4865,6 +4869,10 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) {
return true;
}
if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_I32) {
// DSA lightning-indexer top_k indices (I32) copy/cont.
return true;
}
if (ggml_is_quantized(src0_type) && (src1_type == GGML_TYPE_F16 || src1_type == GGML_TYPE_F32)) {
return true;
}
@@ -4879,6 +4887,7 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
return false;
} break;
case GGML_OP_REDUCE:
case GGML_OP_BLEND:
case GGML_OP_FAKE_CPY:
case GGML_OP_ARGMAX:
return true;
@@ -4960,11 +4969,23 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
return ggml_is_contiguous(op->src[0]);
//case GGML_OP_ROPE:
// return ggml_is_contiguous(op->src[0]);
case GGML_OP_ARGSORT:
return true;
case GGML_OP_ARGSORT_THRESH:
// The CUDA bitonic argsort launches one thread per (padded) column, so the
// row width rounded up to a power of 2 must fit in a single CUDA block (<=1024
// threads). Wider rows (e.g. the DSA lightning-indexer scoring over a large
// n_kv) would fail at launch with "invalid configuration argument"; report them
// as unsupported so the scheduler falls back to the CPU argsort (no size limit).
{
int64_t ncols = op->src[0]->ne[0];
int64_t ncols_pad = 1;
while (ncols_pad < ncols) ncols_pad *= 2;
return ncols_pad <= 1024;
}
case GGML_OP_IM2COL:
case GGML_OP_POOL_2D:
case GGML_OP_SUM_ROWS:
case GGML_OP_ARGSORT:
case GGML_OP_ARGSORT_THRESH:
case GGML_OP_GROUPED_TOPK:
case GGML_OP_ACC:
case GGML_OP_GROUP_NORM:
ggml/src/ggml-cuda/argsort.cu
+157
View File
@@ -426,6 +426,151 @@ static void argsort_openai_f32_f32_i32_cuda(const float * x, float * weights, in
}
}
#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
# define GGML_CUDA_USE_CUB
#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
#ifdef GGML_CUDA_USE_CUB
# include <cub/cub.cuh>
# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1)
# define STRIDED_ITERATOR_AVAILABLE
# include <cuda/iterator>
# endif
using namespace cub;
#endif // GGML_CUDA_USE_CUB
#ifndef STRIDED_ITERATOR_AVAILABLE
static __global__ void init_offsets(int * offsets, const int ncols, const int nrows) {
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx <= nrows) {
offsets[idx] = idx * ncols;
}
}
#endif // STRIDED_ITERATOR_AVAILABLE
#ifdef GGML_CUDA_USE_CUB
static __global__ void init_indices(int * indices, const int ncols, const int nrows) {
const int col = blockIdx.x * blockDim.x + threadIdx.x;
const int row = blockIdx.y;
if (col < ncols && row < nrows) {
indices[row * ncols + col] = col;
}
}
void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
const float * x,
int * dst,
const int ncols,
const int nrows,
ggml_sort_order order,
cudaStream_t stream) {
ggml_cuda_pool_alloc<int> temp_indices_alloc(pool, ncols * nrows);
ggml_cuda_pool_alloc<float> temp_keys_alloc(pool, ncols * nrows);
int * temp_indices = temp_indices_alloc.get();
float * temp_keys = temp_keys_alloc.get();
static const int block_size = 256;
const dim3 grid_size((ncols + block_size - 1) / block_size, nrows);
init_indices<<<grid_size, block_size, 0, stream>>>(temp_indices, ncols, nrows);
#ifdef STRIDED_ITERATOR_AVAILABLE
auto offset_iterator = cuda::make_strided_iterator(cuda::make_counting_iterator(0), ncols);
#else
// offset_iterator needs to populate nrows + 1 elements, so we also have to ceildiv nrows + 1 by block_size
const int nrows_offset = nrows + 1;
ggml_cuda_pool_alloc<int> offsets_alloc(pool, nrows_offset);
int * offset_iterator = offsets_alloc.get();
const dim3 offset_grid((nrows_offset + block_size - 1) / block_size);
init_offsets<<<offset_grid, block_size, 0, stream>>>(offset_iterator, ncols, nrows);
#endif
CUDA_CHECK(cudaMemcpyAsync(temp_keys, x, ncols * nrows * sizeof(float), cudaMemcpyDeviceToDevice, stream));
size_t temp_storage_bytes = 0;
bool is_capturing = false;
#ifdef USE_CUDA_GRAPH
// Currently (confirmed for CCCL <= 3.2) DeviceSegmentedSort does not support stream capture, while DeviceSegmentedRadixSort does.
// See https://github.com/NVIDIA/cccl/issues/5661#issuecomment-3229037149
// TODO: constrain this to the CCCL versions that have this issue once it's resolved in a future CCCL release.
cudaStreamCaptureStatus capture_status;
CUDA_CHECK(cudaStreamIsCapturing(stream, &capture_status));
is_capturing = (capture_status != cudaStreamCaptureStatusNone);
#endif // USE_CUDA_GRAPH
if (order == GGML_SORT_ORDER_ASC) {
if (nrows == 1) {
CUDA_CHECK(DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place)
temp_indices, dst, // values (indices)
ncols, 0, sizeof(float) * 8, stream));
} else if (is_capturing) {
CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(
nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place)
temp_indices, dst, // values (indices)
ncols * nrows, nrows, // num items, num segments
offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream));
} else {
CUDA_CHECK(DeviceSegmentedSort::SortPairs(nullptr, temp_storage_bytes, temp_keys,
temp_keys, // keys (in-place)
temp_indices, dst, // values (indices)
ncols * nrows, nrows, // num items, num segments
offset_iterator, offset_iterator + 1, stream));
}
} else {
if (nrows == 1) {
CUDA_CHECK(DeviceRadixSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys,
temp_keys, // keys (in-place)
temp_indices, dst, // values (indices)
ncols, 0, sizeof(float) * 8, stream));
} else if (is_capturing) {
CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending(
nullptr, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows,
offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream));
} else {
CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(nullptr, temp_storage_bytes, temp_keys, temp_keys,
temp_indices, dst, ncols * nrows, nrows,
offset_iterator, offset_iterator + 1, stream));
}
}
ggml_cuda_pool_alloc<uint8_t> temp_storage_alloc(pool, temp_storage_bytes);
void * d_temp_storage = temp_storage_alloc.get();
if (order == GGML_SORT_ORDER_ASC) {
if (nrows == 1) {
CUDA_CHECK(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys,
temp_keys, // keys (in-place)
temp_indices, dst, // values (indices)
ncols, 0, sizeof(float) * 8, stream));
} else if (is_capturing) {
CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys,
temp_indices, dst, ncols * nrows, nrows, offset_iterator,
offset_iterator + 1, 0, sizeof(float) * 8, stream));
} else {
CUDA_CHECK(DeviceSegmentedSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys,
temp_indices, dst, ncols * nrows, nrows, offset_iterator,
offset_iterator + 1, stream));
}
} else {
if (nrows == 1) {
CUDA_CHECK(DeviceRadixSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys,
temp_keys, // keys (in-place)
temp_indices, dst, // values (indices)
ncols, 0, sizeof(float) * 8, stream));
} else if (is_capturing) {
CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending(
d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows,
offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream));
} else {
CUDA_CHECK(DeviceSegmentedSort::SortPairsDescending(d_temp_storage, temp_storage_bytes, temp_keys,
temp_keys, temp_indices, dst, ncols * nrows, nrows,
offset_iterator, offset_iterator + 1, stream));
}
}
}
#endif // GGML_CUDA_USE_CUB
void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const float * src0_d = (const float *)src0->data;
@@ -439,8 +584,20 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const int64_t ncols = src0->ne[0];
const int64_t nrows = ggml_nrows(src0);
#ifdef GGML_CUDA_USE_CUB
const int ncols_pad = next_power_of_2(ncols);
const size_t shared_mem = ncols_pad * sizeof(int);
const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb;
enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0];
if (shared_mem > max_shared_mem || ncols > 1024) {
ggml_cuda_pool & pool = ctx.pool();
argsort_f32_i32_cuda_cub(pool, src0_d, (int *) dst_d, ncols, nrows, order, stream);
return;
}
#endif
argsort_f32_T_cuda(src0_d, (int *)dst_d, ncols, nrows, ncols, order, -1, 0.f, stream);
}
ggml/src/ggml-cuda/blend.cu
+101
View File
@@ -0,0 +1,101 @@
#include "blend.cuh"
#define CUDA_BLEND_BLOCK_SIZE 256
template <typename Data, typename Idx>
static __global__ void kernel_blend(int n, int nidx, const Data * x, const Idx * idx, Data * y, float c,
int ne1, int ne2,
size_t nb01, size_t nb02, size_t nb03,
size_t nb11, size_t nb12, size_t nb13,
size_t nb1, size_t nb2, size_t nb3) {
Data b;
if constexpr (std::is_same_v<Data, nv_bfloat16>) {
b = __float2bfloat16(c);
} else {
b = (Data)c;
}
int ii = blockIdx.x;
int i3 = ii / (ne1*ne2); ii -= i3*ne1*ne2;
int i2 = ii / ne1;
int i1 = ii - i2*ne1;
auto x_row = x + i1*nb01 + i2*nb02 + i3*nb03;
auto y_row = y + i1*nb1 + i2*nb2 + i3*nb3;
auto idx_row = idx + i1*nb11 + i2*nb12 + i3*nb13;
if (x_row != y_row) {
for (int i = threadIdx.x; i < n; i += blockDim.x) {
y_row[i] = x_row[i];
}
__syncthreads();
}
for (int i = threadIdx.x; i < nidx; i += blockDim.x) {
y_row[idx_row[i]] = b;
}
}
void ggml_cuda_op_blend(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const ggml_tensor * src1 = dst->src[1];
GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16);
GGML_ASSERT(src0->type == dst->type);
GGML_ASSERT(ggml_are_same_shape(src0, dst));
GGML_ASSERT(src1->type == GGML_TYPE_I32 || src1->type == GGML_TYPE_I64);
GGML_ASSERT(src1->ne[1] == src0->ne[1] && src1->ne[2] == src0->ne[2] && src1->ne[3] == src0->ne[3]);
GGML_ASSERT(src1->ne[0] <= src0->ne[0]);
float c;
memcpy(&c, dst->op_params, sizeof(c));
auto nrows = ggml_nrows(dst);
dim3 grid_dims(nrows, 1, 1);
dim3 block_size(CUDA_BLEND_BLOCK_SIZE, 1, 1);
if (src1->type == GGML_TYPE_I32) {
auto idx = (const int32_t *)src1->data;
if (src0->type == GGML_TYPE_F32) {
kernel_blend<<<grid_dims, block_size, 0, ctx.stream()>>>(src0->ne[0], src1->ne[0],
(const float *)src0->data, idx, (float *)dst->data, c, src0->ne[1], src0->ne[2],
src0->nb[1]/sizeof(float), src0->nb[2]/sizeof(float), src0->nb[3]/sizeof(float),
src1->nb[1]/sizeof(int32_t), src1->nb[2]/sizeof(int32_t), src1->nb[3]/sizeof(int32_t),
dst->nb[1]/sizeof(float), dst->nb[2]/sizeof(float), dst->nb[3]/sizeof(float));
}
else if (src0->type == GGML_TYPE_F16) {
kernel_blend<<<grid_dims, block_size, 0, ctx.stream()>>>(src0->ne[0], src1->ne[0],
(const half *)src0->data, idx, (half *)dst->data, c, src0->ne[1], src0->ne[2],
src0->nb[1]/sizeof(half), src0->nb[2]/sizeof(half), src0->nb[3]/sizeof(half),
src1->nb[1]/sizeof(int32_t), src1->nb[2]/sizeof(int32_t), src1->nb[3]/sizeof(int32_t),
dst->nb[1]/sizeof(half), dst->nb[2]/sizeof(half), dst->nb[3]/sizeof(half));
}
else {
kernel_blend<<<grid_dims, block_size, 0, ctx.stream()>>>(src0->ne[0], src1->ne[0],
(const nv_bfloat16 *)src0->data, idx, (nv_bfloat16 *)dst->data, c, src0->ne[1], src0->ne[2],
src0->nb[1]/sizeof(nv_bfloat16), src0->nb[2]/sizeof(nv_bfloat16), src0->nb[3]/sizeof(nv_bfloat16),
src1->nb[1]/sizeof(int32_t), src1->nb[2]/sizeof(int32_t), src1->nb[3]/sizeof(int32_t),
dst->nb[1]/sizeof(nv_bfloat16), dst->nb[2]/sizeof(nv_bfloat16), dst->nb[3]/sizeof(nv_bfloat16));
}
} else {
auto idx = (const int64_t *)src1->data;
if (src0->type == GGML_TYPE_F32) {
kernel_blend<<<grid_dims, block_size, 0, ctx.stream()>>>(src0->ne[0], src1->ne[0],
(const float *)src0->data, idx, (float *)dst->data, c, src0->ne[1], src0->ne[2],
src0->nb[1]/sizeof(float), src0->nb[2]/sizeof(float), src0->nb[3]/sizeof(float),
src1->nb[1]/sizeof(int64_t), src1->nb[2]/sizeof(int64_t), src1->nb[3]/sizeof(int64_t),
dst->nb[1]/sizeof(float), dst->nb[2]/sizeof(float), dst->nb[3]/sizeof(float));
}
else if (src0->type == GGML_TYPE_F16) {
kernel_blend<<<grid_dims, block_size, 0, ctx.stream()>>>(src0->ne[0], src1->ne[0],
(const half *)src0->data, idx, (half *)dst->data, c, src0->ne[1], src0->ne[2],
src0->nb[1]/sizeof(half), src0->nb[2]/sizeof(half), src0->nb[3]/sizeof(half),
src1->nb[1]/sizeof(int64_t), src1->nb[2]/sizeof(int64_t), src1->nb[3]/sizeof(int64_t),
dst->nb[1]/sizeof(half), dst->nb[2]/sizeof(half), dst->nb[3]/sizeof(half));
}
else {
kernel_blend<<<grid_dims, block_size, 0, ctx.stream()>>>(src0->ne[0], src1->ne[0],
(const nv_bfloat16 *)src0->data, idx, (nv_bfloat16 *)dst->data, c, src0->ne[1], src0->ne[2],
src0->nb[1]/sizeof(nv_bfloat16), src0->nb[2]/sizeof(nv_bfloat16), src0->nb[3]/sizeof(nv_bfloat16),
src1->nb[1]/sizeof(int64_t), src1->nb[2]/sizeof(int64_t), src1->nb[3]/sizeof(int64_t),
dst->nb[1]/sizeof(nv_bfloat16), dst->nb[2]/sizeof(nv_bfloat16), dst->nb[3]/sizeof(nv_bfloat16));
}
}
}
ggml/src/ggml-cuda/blend.cuh
+3
View File
@@ -0,0 +1,3 @@
#include "common.cuh"
void ggml_cuda_op_blend(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
ggml/src/ggml-cuda/cpy.cu
+4
View File
@@ -643,6 +643,10 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg
ggml_cpy_flt_cuda<float, int32_t> (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream, dest_ptrs_d, graph_cpynode_index);
} else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_F32) {
ggml_cpy_flt_cuda<int32_t, float> (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream, dest_ptrs_d, graph_cpynode_index);
} else if (src0->type == GGML_TYPE_I32 && src1->type == GGML_TYPE_I32) {
// Needed for the DSA lightning-indexer: top_k (argsort) produces I32 indices that may be
// cont'd / copied / moved across backends (CPU argsort fallback for large n_kv).
ggml_cpy_flt_cuda<int32_t, int32_t> (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream, dest_ptrs_d, graph_cpynode_index);
} else if (ggml_are_same_shape(src0, src1) && src0->type == GGML_TYPE_Q8_0 && src1->type == GGML_TYPE_Q8_0) {
// This is needed for MLA with mla=2 when using q8_0 cache.
transpose_q8_0(ctx, src0, src1);
ggml/src/ggml.c
+67 -14
View File
@@ -4320,9 +4320,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
"FAKE_CPY",
"FUSED_NORM",
"FUSED_RMS_RMS_ADD",
"BLEND",
};
static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103");
static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"none",
@@ -4440,10 +4441,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
"fake_cpy(x,y)",
"norm(x,y)",
"rms(x1)+rms(x2)",
"blend(a,b,c)",
};
static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103");
static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
@@ -10137,6 +10139,28 @@ struct ggml_tensor * ggml_fill_inplace(
return ggml_fill_impl(ctx, a, c, true);
}
// ggml blend
struct ggml_tensor * ggml_blend(
struct ggml_context * ctx,
struct ggml_tensor * a,
struct ggml_tensor * b,
float c) {
GGML_ASSERT(a->type == GGML_TYPE_F32 || a->type == GGML_TYPE_F16 || a->type == GGML_TYPE_BF16);
GGML_ASSERT(b->type == GGML_TYPE_I32 || b->type == GGML_TYPE_I64);
GGML_ASSERT(b->ne[0] <= a->ne[0]);
for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) {
GGML_ASSERT(a->ne[dim] == b->ne[dim]);
}
struct ggml_tensor * result = ggml_dup_tensor(ctx, a);
result->src[0] = a;
result->src[1] = b;
memcpy(result->op_params, &c, sizeof(c));
result->op = GGML_OP_BLEND;
return result;
}
// ggml_argsort
struct ggml_tensor * ggml_argsort(
@@ -14974,8 +14998,6 @@ static void ggml_compute_forward_concat_any(
GGML_ASSERT(src0->type == src1->type && src0->type == dst->type);
const int32_t dim = ggml_get_op_params_i32(dst, 0);
// Let's do it for dim = 0 only for now
GGML_ASSERT(dim == 0);
int ith = params->ith;
int nth = params->nth;
@@ -14983,21 +15005,46 @@ static void ggml_compute_forward_concat_any(
int64_t nrows = ggml_nrows(dst);
int64_t nrows_per_thread = (nrows + nth - 1)/nth;
int64_t first_row = ith*nrows_per_thread;
if (first_row >= nrows) return;
int64_t last_row = MIN(first_row + nrows_per_thread, nrows);
if (first_row >= last_row) return;
int64_t src0_row_size = ggml_row_size(src0->type, src0->ne[0]);
int64_t src1_row_size = ggml_row_size(src1->type, src1->ne[0]);
for (int64_t row = first_row; row < last_row; ++row) {
int64_t i3 = row/(dst->ne[1]*dst->ne[2]);
int64_t i2 = (row - i3*dst->ne[1]*dst->ne[2])/dst->ne[1];
int64_t i1 = row - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1];
char * y = (char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3];
const char * x0 = (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3];
const char * x1 = (const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3];
memcpy(y, x0, src0_row_size);
memcpy(y + src0_row_size, x1, src1_row_size);
if (dim == 0) {
for (int64_t row = first_row; row < last_row; ++row) {
int64_t i3 = row/(dst->ne[1]*dst->ne[2]);
int64_t i2 = (row - i3*dst->ne[1]*dst->ne[2])/dst->ne[1];
int64_t i1 = row - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1];
char * y = (char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3];
const char * x0 = (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3];
const char * x1 = (const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3];
memcpy(y, x0, src0_row_size);
memcpy(y + src0_row_size, x1, src1_row_size);
}
}
else {
GGML_ASSERT(src0_row_size == src1_row_size);
for (int64_t row = first_row; row < last_row; ++row) {
int64_t i3 = row/(dst->ne[1]*dst->ne[2]);
int64_t i2 = (row - i3*dst->ne[1]*dst->ne[2])/dst->ne[1];
int64_t i1 = row - i3*dst->ne[1]*dst->ne[2] - i2*dst->ne[1];
char * y = (char *)dst->data + i1*dst->nb[1] + i2*dst->nb[2] + i3*dst->nb[3];
const char * x;
if (dim == 1) {
x = i1 < src0->ne[1] ? (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]
: (const char *)src1->data + (i1 - src0->ne[1])*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3];
}
else if (dim == 2) {
x = i2 < src0->ne[2] ? (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]
: (const char *)src1->data + i1*src1->nb[1] + (i2 - src0->ne[2])*src1->nb[2] + i3*src1->nb[3];
}
else {
x = i3 < src0->ne[3] ? (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3]
: (const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + (i3 - src0->ne[3])*src1->nb[3];
}
memcpy(y, x, src0_row_size);
}
}
}
@@ -24404,6 +24451,10 @@ static int ggml_compute_forward(struct ggml_compute_params * params, struct ggml
{
iqk_rms_rms_add(tensor, params->ith, params->nth);
} break;
case GGML_OP_BLEND:
{
iqk_blend(tensor, params->ith, params->nth);
} break;
case GGML_OP_FUSED_NORM:
{
ggml_compute_forward_fused_norm(params, tensor);
@@ -25243,6 +25294,7 @@ static void ggml_compute_backward(struct ggml_context * ctx, struct ggml_tensor
case GGML_OP_FUSED_RMS_NORM:
case GGML_OP_FUSED_RMS_RMS_ADD:
case GGML_OP_FUSED_NORM:
case GGML_OP_BLEND:
{
GGML_ABORT("fatal error"); // TODO: not implemented
}
@@ -26443,6 +26495,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
case GGML_OP_RMS_NORM:
case GGML_OP_FUSED_RMS_NORM:
case GGML_OP_FUSED_RMS_RMS_ADD:
case GGML_OP_BLEND:
case GGML_OP_FUSED_NORM:
case GGML_OP_RMS_NORM_BACK:
case GGML_OP_GROUP_NORM:
ggml/src/iqk/iqk_cpu_ops.cpp
+70
View File
@@ -974,3 +974,73 @@ void iqk_rms_rms_add(struct ggml_tensor * dst, int ith, int nth) {
}
}
}
namespace {
template <typename Data, typename Idx>
inline void iqk_blend_row(int n, int nidx, const Data * x, const Idx * idx, Data * y, float c) {
Data b;
if constexpr (std::is_same_v<Data, float>) {
b = c;
}
else if constexpr (std::is_same_v<Data, ggml_half>) {
b = GGML_FP32_TO_FP16(c);
}
else {
b = GGML_FP32_TO_BF16(c);
}
if (y != x) {
for (int j = 0; j < n; ++j) y[j] = x[j];
}
for (int j = 0; j < nidx; ++j) y[idx[j]] = b;
}
}
void iqk_blend(struct ggml_tensor * dst, int ith, int nth) {
auto src0 = dst->src[0];
auto src1 = dst->src[1];
GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16);
GGML_ASSERT(src1->type == GGML_TYPE_I32 || src1->type == GGML_TYPE_I64);
GGML_ASSERT(src1->ne[0] <= src0->ne[0]);
GGML_ASSERT(src0->type == dst->type);
GGML_ASSERT(src0->ne[0] == dst->ne[0]);
for (int dim = 1; dim < GGML_MAX_DIMS; ++dim) {
GGML_ASSERT(src0->ne[dim] == src1->ne[dim]);
GGML_ASSERT(src0->ne[dim] == dst->ne[dim]);
}
float c;
std::memcpy(&c, dst->op_params, sizeof(c));
int nrows = ggml_nrows(src0);
int npt = (nrows + nth - 1)/nth;
int first = ith*npt;
int last = std::min(nrows, first + npt);
for (int ir = first; ir < last; ++ir) {
int ii = ir;
int i3 = ii/(src0->ne[1]*src0->ne[2]); ii -= i3*src0->ne[1]*src0->ne[2];
int i2 = ii/(src0->ne[1]); ii -= i2*src0->ne[1];
int i1 = ii;
auto x = (const char *)src0->data + i1*src0->nb[1] + i2*src0->nb[2] + i3*src0->nb[3];
auto y = ( char *) dst->data + i1* dst->nb[1] + i2* dst->nb[2] + i3* dst->nb[3];
auto idx = (const char *)src1->data + i1*src1->nb[1] + i2*src1->nb[2] + i3*src1->nb[3];
if (src1->type == GGML_TYPE_I32) {
if (src0->type == GGML_TYPE_F32) {
iqk_blend_row(src0->ne[0], src1->ne[0], (const float *)x, (const int32_t *)idx, (float *)y, c);
}
else if (src0->type == GGML_TYPE_F16) {
iqk_blend_row(src0->ne[0], src1->ne[0], (const ggml_half *)x, (const int32_t *)idx, (ggml_half *)y, c);
}
else {
iqk_blend_row(src0->ne[0], src1->ne[0], (const ggml_bf16_t *)x, (const int32_t *)idx, (ggml_bf16_t *)y, c);
}
} else {
if (src0->type == GGML_TYPE_F32) {
iqk_blend_row(src0->ne[0], src1->ne[0], (const float *)x, (const int64_t *)idx, (float *)y, c);
}
else if (src0->type == GGML_TYPE_F16) {
iqk_blend_row(src0->ne[0], src1->ne[0], (const ggml_half *)x, (const int64_t *)idx, (ggml_half *)y, c);
}
else {
iqk_blend_row(src0->ne[0], src1->ne[0], (const ggml_bf16_t *)x, (const int64_t *)idx, (ggml_bf16_t *)y, c);
}
}
}
}
ggml/src/iqk/iqk_cpu_ops.h
+2
View File
@@ -41,6 +41,8 @@ bool iqk_ssm_conv4(int nr, int nc, int nt,
void iqk_rms_rms_add(struct ggml_tensor * dst, int ith, int nth);
void iqk_blend(struct ggml_tensor * dst, int ith, int nth);
#ifdef __cplusplus
}
#endif
include/llama.h
+2
View File
@@ -489,6 +489,8 @@ extern "C" {
bool fused_mmad; // whether to use fused mul+multi_add op [EXPERIMENTAL]
bool rope_cache; // whether to use RoPE cache [EXPERIMENTAL]
bool graph_reuse; // whether to reuse graphs when possible [EXPERIMENTAL]
bool dsa; // enable GLM DSA sparse attention (off by default) [EXPERIMENTAL]
int dsa_top_k; // DSA top-k override (<0 => model's configured indexer_top_k) [EXPERIMENTAL]
int min_experts;
float thresh_experts;
bool only_active_experts;
src/graphs/build_deepseek2.cpp
+424 -4
View File
@@ -323,6 +323,370 @@ ggml_tensor * llm_build_context::build_deepseek2_tp_attention(
return combined;
}
// DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). CACHE-BACKED: the batch's indexer keys are
// (Hadamard-rotated and) written to a persistent per-layer indexer-key cache (kv_self.kr_l[il]) at
// kv_head, then the FULL [head_size, n_kv] cached key set is read back and scored against the current
// batch's indexer queries. This makes DECODE correct: a generated token (kv_head>0) scores against
// ALL past indexer keys, not just itself. Returns top_k [n_top_k, n_tokens] over the n_kv key axis.
//
// The Walsh-Hadamard rotation H (orthonormal, H^2==I) is applied to both indexer_q and indexer_k.
// (H q)*(H k) == q*k so it is score-preserving; its purpose is to improve the precision of the keys
// we store in the F16 indexer cache (matches the reference). Gated by cparams.dsa_indexer_hadamard.
//
// MULTI-SEQUENCE & FA: both are now handled.
// - Multi-sequence batches (n_seq>1) are correct: each token is written to its own cache cell at
// kv_self.head+i (per-sequence), the base block-diagonal KQ_mask drives cross-sequence keys to
// -inf before argsort, and the attention sink is anchored per-sequence (inp_dsa_sink). Validated
// n_seq=4 == n_seq=1 chunk-for-chunk (UPDATE 5).
// - The FA path (-fa 1) consumes the sparse top-k mask too: build_deepseek2_dsa_fa_mask converts
// the F32 sparse mask into the padded F16 mask the flash-attention kernel reads (UPDATE 4). It
// does NOT fall back to the dense KQ_mask.
// - Serving (context-shift / defrag / multi-turn seq_rm): the persistent indexer-key cache kr_l is
// maintained by build_k_shift (delta-RoPE around the Hadamard), build_defrag (row move), and the
// seq ops are metadata-only so kr_l rows stay matched to their cells (UPDATE 6).
ggml_tensor * llm_build_context::build_deepseek2_dsa_indexer(
ggml_cgraph * gf,
int il,
ggml_tensor * qr,
ggml_tensor * cur,
ggml_tensor * KQ_mask,
ggml_tensor * inp_pos) {
const auto & layer = model.layers[il];
const int64_t n_ihead = hparams.indexer_n_head;
const int64_t head_size = hparams.indexer_head_size;
const int64_t rope_dim = n_rot; // n_embd_head_qk_rope
const int64_t nope_dim = head_size - rope_dim;
// ---- indexer_q : {head_size * n_ihead, n_tokens} ----
ggml_tensor * indexer_q = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr);
cb(indexer_q, "dsa_indexer_q", il);
// split rope/nope along dim0, per head
ggml_tensor * indexer_q_pe = ggml_view_3d(ctx0, indexer_q, rope_dim, n_ihead, n_tokens,
ggml_row_size(indexer_q->type, head_size),
ggml_row_size(indexer_q->type, head_size) * n_ihead, 0);
ggml_tensor * indexer_q_nope = ggml_view_3d(ctx0, indexer_q, nope_dim, n_ihead, n_tokens,
ggml_row_size(indexer_q->type, head_size),
ggml_row_size(indexer_q->type, head_size) * n_ihead,
ggml_row_size(indexer_q->type, rope_dim));
indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot,
LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
// {head_size, n_ihead, n_tokens}
indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0);
cb(indexer_q, "dsa_indexer_q_cat", il);
// ---- indexer_k : {head_size, n_tokens} (single key head, MQA) ----
ggml_tensor * indexer_k = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur);
// LayerNorm (with weight + bias) over head_size
indexer_k = llm_build_norm(ctx0, indexer_k, hparams, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, cb, il);
cb(indexer_k, "dsa_indexer_k", il);
ggml_tensor * indexer_k_pe = ggml_view_3d(ctx0, indexer_k, rope_dim, 1, n_tokens,
ggml_row_size(indexer_k->type, head_size),
ggml_row_size(indexer_k->type, head_size), 0);
ggml_tensor * indexer_k_nope = ggml_view_3d(ctx0, indexer_k, nope_dim, 1, n_tokens,
ggml_row_size(indexer_k->type, head_size),
ggml_row_size(indexer_k->type, head_size),
ggml_row_size(indexer_k->type, rope_dim));
indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot,
LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow);
// {head_size, 1, n_tokens}
indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0);
cb(indexer_k, "dsa_indexer_k_cat", il);
// ---- Walsh-Hadamard rotation (score-preserving; improves cached-K F16 precision) ----
// nrot = largest power of 2 dividing head_size (== head_size for head_size = 128).
// DSA_HADAMARD_DISABLE: DEBUG-ONLY env knob (no CLI surface). Default = rotation enabled.
static const bool dsa_had_disable = getenv("DSA_HADAMARD_DISABLE") != nullptr;
if (lctx.cparams.dsa_indexer_hadamard && !dsa_had_disable) {
GGML_ASSERT((head_size & ~(head_size - 1)) == head_size);
indexer_q = ggml_hadamard(ctx0, indexer_q, head_size);
indexer_k = ggml_hadamard(ctx0, indexer_k, head_size);
}
// ---- write the batch's indexer keys into the persistent indexer-key cache at kv_head ----
// kr_l[il] is [head_size, kv_size] (F16, MQA single head). Store {head_size, n_tokens}.
ggml_tensor * kr_cache = kv_self.kr_l[il];
GGML_ASSERT(kr_cache && "DSA indexer key cache not allocated");
{
ggml_tensor * indexer_k_2d = ggml_reshape_2d(ctx0, indexer_k, head_size, n_tokens);
ggml_tensor * kr_view = ggml_view_2d(ctx0, kr_cache, head_size, n_tokens,
ggml_row_size(kr_cache->type, head_size),
ggml_row_size(kr_cache->type, head_size) * kv_head);
ggml_tensor * kr_cpy = ggml_cpy(ctx0, indexer_k_2d, kr_view);
// GRAPH-REUSE FIXUP REGISTRATION: the K/V cache_copies fixup in update_cache_copies()
// re-points the K/V cache writes to the current kv_head when a graph is reused, but it
// does NOT touch this indexer-key (kr_l) write, whose view bakes kv_head at build time.
// Under FA the cache pads to 256, so consecutive decode ubatches keep the SAME n_kv and
// can_reuse_graph() reuses the graph -- without this registration the kr_l write stays
// baked at the first ubatch's kv_head, so later ubatches never write their recent index
// keys (those cells read uninitialized -> block-max-pool/top-k drops the genuinely
// attended recent block -> degraded/NaN sparse-FA decode). Register it like K/V so
// update_cache_copies() patches view_offs = kv_head * step each reuse.
// step = one index-key row = head_size * F16 = kr_cache->nb[1].
if ((size_t) il < lctx.dsa_cache_copies.size()) {
lctx.dsa_cache_copies[il].cpy = kr_cpy;
lctx.dsa_cache_copies[il].step = kr_cache->nb[1];
}
ggml_build_forward_expand(gf, kr_cpy);
}
// ---- read back the full cached key set: {head_size, n_kv} ----
ggml_tensor * cached_k = ggml_view_2d(ctx0, kr_cache, head_size, n_kv,
ggml_row_size(kr_cache->type, head_size), 0);
cb(cached_k, "dsa_cached_k", il);
// ---- indexer weights : {n_ihead, n_tokens} ----
ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, layer.indexer_proj, cur);
indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(head_size * n_ihead)));
cb(indexer_weights, "dsa_indexer_weights", il);
// ---- scores ----
// indexer_q : {head_size, n_ihead, n_tokens} -> {head_size, n_tokens, n_ihead}
// cached_k : {head_size, n_kv} -> {head_size, n_kv, 1}; broadcasts over q's n_ihead dim.
ggml_tensor * indexer_k_b = ggml_reshape_3d(ctx0, cached_k, head_size, n_kv, 1);
ggml_tensor * indexer_score = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0);
if (indexer_score->type != GGML_TYPE_F32) {
indexer_score = ggml_cast(ctx0, indexer_score, GGML_TYPE_F32);
cb(indexer_score, "indexer_score_f32", il);
}
for (int head = 0; head < n_ihead; ++head) {
int il_cb = 1000*(il + 1) + head;
// [1, n_tokens]
auto w = ggml_cont(ctx0, ggml_view_2d(ctx0, indexer_weights, 1, indexer_weights->ne[1], indexer_weights->nb[1], indexer_weights->nb[0]*head));
cb(w, "iweights", il_cb);
// [head_size, n_tokens]
auto q = ggml_view_2d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[2], indexer_q->nb[2], indexer_q->nb[1]*head);
// [n_kv, n_tokens]
auto kq = ggml_mul_mat(ctx0, indexer_k_b, q);
cb(kq, "ikq", il_cb);
// [n_kv, n_tokens]
kq = ggml_relu(ctx0, kq);
cb(kq, "ikq_relu", il_cb);
// [n_kv, n_tokens]
auto score = ggml_mul(ctx0, kq, w);
cb(score, "score", il_cb);
indexer_score = ggml_add_inplace(ctx0, indexer_score, score);
cb(indexer_score, "indexer_score", il_cb);
ggml_build_forward_expand(gf, indexer_score);
}
//// {n_kv(keys), n_tokens(q), n_ihead} (k's head dim broadcasts over n_ihead)
//ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k_b, indexer_q);
//cb(indexer_kq, "dsa_indexer_kq", il);
//// -> {n_ihead, n_tokens(q), n_kv(keys)} for per-head weighting
//indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
//ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq);
//// weights {n_ihead, n_tokens} -> {n_ihead, n_tokens, 1} broadcast over keys
//indexer_weights = ggml_reshape_3d(ctx0, indexer_weights, n_ihead, n_tokens, 1);
//indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
//// sum over heads -> {1, n_tokens(q), n_kv(keys)}
//indexer_score = ggml_sum_rows(ctx0, indexer_score);
//// -> {n_kv(keys), n_tokens(q), 1}
//indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
//cb(indexer_score, "dsa_indexer_score", il);
//// add base causal mask over the n_kv keys: first n_tokens query columns of KQ_mask {n_kv, n_tokens_pad}.
//ggml_tensor * causal = ggml_view_2d(ctx0, KQ_mask, n_kv, n_tokens, KQ_mask->nb[1], 0);
//// Under -fa 1 the dense KQ_mask is F16; CPU ggml_add only supports F32+F16 when src0 is F16,
//// not F32(score)+F16(mask) (it aborts). Cast the causal mask view to F32 so the add is valid on
//// CPU. (CUDA add accepts mixed types, so this only bit the CPU build.)
//if (causal->type != GGML_TYPE_F32) causal = ggml_cast(ctx0, ggml_cont(ctx0, causal), GGML_TYPE_F32);
//indexer_score = ggml_add(ctx0, indexer_score, causal);
//cb(indexer_score, "dsa_indexer_score_masked", il);
// Attention-sink force-inclusion: add a finite positive boost to each query's OWN SEQUENCE's
// first n_sink present tokens so the sink token(s) always survive the top-k selection. Masking
// the sink collapses most transformers; a heavily-quantized (IQ2) indexer does not reliably rank
// it high on its own. The boost is finite, so it cannot un-mask future/causal -inf positions
// (-inf + boost = -inf).
//
// MULTI-SEQUENCE: the boost MUST be per-(key,query), not a global per-key vector. With several
// sequences packed into one ubatch (seq 0 at cache cells [0,n0), seq 1 at [n0,n1), ...), a global
// "key index < n_sink" boost only protects sequence 0's sink; sequence 1's sink (at cell n0, not
// cell 0) is left unprotected and gets dropped from top-k once the mask bites, collapsing it.
// We therefore use a per-graph input tensor inp_dsa_sink {n_kv, n_tokens} (filled on the CPU from
// kv_self.cells like the KQ_mask, in llama_set_inputs): inp_dsa_sink[j,i] = 1e20 iff key cell i
// belongs to query j's sequence AND its pos is within [min present pos of that seq, +n_sink).
//
// SERVING: the anchor is each sequence's FIRST PRESENT pos, not absolute pos < n_sink. After
// multi-turn seq_rm drops a sequence's early tokens its earliest survivor has pos >= n_sink; an
// absolute test would then protect nothing and let the (now-)sink be masked out. For a fresh
// sequence starting at pos 0, min(pos)==0 so the boosted set is exactly the old "cell pos <
// n_sink" set with the same 1e20 magnitude — n_seq==1 from pos 0 stays byte-identical.
// DSA_SINK: DEBUG-ONLY env knob (no CLI surface). Default = 1 (protect each sequence's first
// present token from being masked out of top-k). Must stay in sync with the two fill sites.
if (lctx.inp_dsa_sink) {
indexer_score = ggml_add(ctx0, indexer_score, lctx.inp_dsa_sink);
cb(indexer_score, "dsa_indexer_score_sink", il);
ggml_build_forward_expand(gf, indexer_score);
}
// FULL descending argsort of the per-query scores over the n_kv axis: {n_kv, n_tokens} (I32).
// We return the full ranking (not just the top-k view): the sparse-mask builder writes a value
// into EVERY key slot keyed by its rank, which avoids relying on ggml_set_rows preserving an
// uninitialized base for partially-written destinations (a CUDA in-place quirk that corrupted
// decode when n_kv > top_k).
//ggml_tensor * sorted = ggml_cont(ctx0, ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC));
//ggml_tensor * sorted = ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC);
ggml_tensor * sorted;
if (cparams.flash_attn) {
int64_t n_top_k = (int64_t) hparams.indexer_top_k;
if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k;
if (n_top_k > indexer_score->ne[0]) n_top_k = indexer_score->ne[0];
sorted = ggml_top_k(ctx0, indexer_score, n_top_k);
sorted = ggml_cont(ctx0, sorted);
} else {
sorted = ggml_argsort(ctx0, indexer_score, GGML_SORT_ORDER_DESC);
}
cb(sorted, "dsa_sorted", il);
return sorted;
}
// Build an additive sparse causal mask {n_kv, n_tok} (F32): 0 for the top-k highest-scoring keys
// per query, a large negative value for the rest, then add the base causal KQ_mask so future/
// padding keys stay masked. ggml_soft_max_ext only requires mask->ne[1] >= q n_tokens, and
// n_tok == n_tokens here, so no padding is needed.
//
// `sorted` is the FULL descending argsort of the indexer scores: {n_kv, n_tok} (I32), where
// sorted[rank, j] = key index with the rank-th highest score for query j. We scatter a rank-based
// penalty into EVERY key slot: pen(rank) = 0 if rank < n_top_k else -BIG. Because every key slot
// is written exactly once (sorted is a per-column permutation), the result does NOT depend on the
// scatter destination's initial contents — sidestepping the ggml in-place set_rows quirk where a
// partially-written CUDA destination keeps uninitialized (garbage) rows.
ggml_tensor * llm_build_context::build_deepseek2_dsa_sparse_mask(
ggml_tensor * sorted,
ggml_tensor * KQ_mask) {
const int64_t n_kv_local = KQ_mask->ne[0];
const int64_t n_tok = sorted->ne[1];
int64_t n_top_k = (int64_t) hparams.indexer_top_k;
// Tuning knob: --dsa-top-k (cparams.dsa_top_k) lets us vary the kept-key count to characterize
// selection quality. <0 means use the model's configured top_k. With the model's configured
// top_k (2048) on heavily-quantized (IQ2_M) weights the indexer currently under-ranks some
// critical keys; a near-n_kv value stays coherent.
if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k;
if (n_top_k > n_kv_local) n_top_k = n_kv_local;
// Penalty magnitude for non-top-k keys. On the soft_max (-fa 0) path this F32 -BIG is added to
// the score and softmaxed -> effectively -inf, while staying finite avoids -inf*0 = NaN. On the
// FA (-fa 1) path this mask is cast to F16 (build_deepseek2_dsa_fa_mask): 1e30 saturates to the
// F16 max (~6.5e4), which is still a large-enough negative bias to zero the key in the FA softmax
// (the dense FA mask uses -INFINITY/F16 -inf there; our finite-but-huge value is equivalent in
// effect and cannot produce NaN). So -BIG masks the key on BOTH paths.
const float BIG = 1e30f;
// rank-based penalty vector: pen[rank] = 0 for rank < n_top_k, else -BIG. {n_kv}
// sel = step(n_top_k - 0.5 - rank) = 1 for rank <= n_top_k-1, else 0
ggml_tensor * rank = ggml_arange(ctx0, 0.0f, (float) n_kv_local, 1.0f); // {n_kv} F32
ggml_tensor * sel = ggml_step(ctx0, ggml_scale_bias(ctx0, rank, -1.0f, (float) n_top_k - 0.5f));
ggml_tensor * pen = ggml_scale_bias(ctx0, sel, BIG, -BIG); // 0 or -BIG
// shape penalty to {1, n_kv, n_tok} (broadcast the per-rank value across all query columns)
pen = ggml_reshape_3d(ctx0, pen, 1, n_kv_local, 1);
ggml_tensor * pen_b = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv_local, n_tok);
pen_b = ggml_repeat(ctx0, pen, pen_b); // {1, n_kv, n_tok}
// destination base {1, n_kv, n_tok} (contents irrelevant — fully overwritten by set_rows)
ggml_tensor * base = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_kv_local, n_tok);
base = ggml_fill(ctx0, base, -BIG);
// indices: {n_kv, n_tok, 1}. scatter pen_b[:, rank, j] into base[:, sorted[rank,j], j].
ggml_tensor * idx = ggml_reshape_3d(ctx0, sorted, n_kv_local, n_tok, 1);
ggml_tensor * scattered = ggml_set_rows(ctx0, base, pen_b, idx);
// {n_kv, n_tok}
ggml_tensor * sparse = ggml_reshape_2d(ctx0, ggml_cont(ctx0, scattered), n_kv_local, n_tok);
// add base causal mask (first n_tok query columns) so future/padding keys stay masked
ggml_tensor * causal = ggml_view_2d(ctx0, KQ_mask, n_kv_local, n_tok, KQ_mask->nb[1], 0);
// see note in build_deepseek2_dsa_indexer: cast F16 (-fa 1) mask to F32 for the CPU add.
if (causal->type != GGML_TYPE_F32) causal = ggml_cast(ctx0, ggml_cont(ctx0, causal), GGML_TYPE_F32);
sparse = ggml_add(ctx0, sparse, causal);
cb(sparse, "dsa_sparse_mask", -1);
return sparse;
}
static ggml_tensor * build_deepseek2_dsa_fa_mask(const llama_context & lctx, ggml_context * ctx0, ggml_tensor * KQ_mask, ggml_tensor * sorted) {
GGML_ASSERT(KQ_mask && KQ_mask->type == GGML_TYPE_F16);
GGML_ASSERT(sorted && sorted->type == GGML_TYPE_I32);
GGML_ASSERT(KQ_mask->ne[1] >= sorted->ne[1]);
int n_top_k = (int64_t) lctx.model.hparams.indexer_top_k;
if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k;
int n_kv_local = KQ_mask->ne[0];
if (n_top_k >= n_kv_local) {
return KQ_mask;
}
GGML_ASSERT(sorted->ne[1] == lctx.inp_mask_inf->ne[1]);
auto top_k = ggml_view_2d(ctx0, sorted, n_top_k, sorted->ne[1], sorted->nb[1], 0);
auto mask32 = ggml_blend(ctx0, lctx.inp_mask_inf, top_k, 0.0f);
if (KQ_mask->ne[1] == mask32->ne[1]) {
auto mask16 = ggml_add(ctx0, KQ_mask, mask32);
return mask16;
}
auto kq1 = ggml_view_2d(ctx0, KQ_mask, KQ_mask->ne[0], mask32->ne[1], KQ_mask->nb[1], 0);
auto kq2 = ggml_view_2d(ctx0, KQ_mask, KQ_mask->ne[0], KQ_mask->ne[1] - mask32->ne[1], KQ_mask->nb[1], mask32->ne[1]*KQ_mask->nb[1]);
kq1 = ggml_add(ctx0, kq1, mask32);
auto mask16 = ggml_concat(ctx0, kq1, kq2, 1);
return mask16;
}
// Adapt the (F32, unpadded {n_kv, n_tokens}) sparse mask for ggml_flash_attn_ext, which on this fork
// requires the mask to be F16, contiguous, and padded in ne[1] to GGML_PAD(n_queries, GGML_KQ_MASK_PAD)
// (build_inp_KQ_mask creates the dense -fa 1 mask exactly that way). We:
// 1) cast the sparse mask to F16,
// 2) concat the dense FA mask's padding rows [n_tok, n_pad) (already F16, causal -inf for the
// non-existent padded queries) onto the bottom so ne[1] matches the dense mask,
// 3) ggml_cont so the result is contiguous (the FA assert requires it).
// The padded rows feed only the discarded outputs of padded query slots, so reusing the dense mask's
// padding region is both correct and the cheapest way to get the exact dense shape.
ggml_tensor * llm_build_context::build_deepseek2_dsa_fa_mask(
ggml_tensor * sparse,
ggml_tensor * KQ_mask) {
const int64_t n_kv_local = KQ_mask->ne[0];
const int64_t n_tok = sparse->ne[1];
const int64_t n_pad = KQ_mask->ne[1]; // GGML_PAD(n_tokens, GGML_KQ_MASK_PAD)
GGML_ASSERT(KQ_mask->type == GGML_TYPE_F16 && "FA dense KQ_mask expected F16 on -fa 1");
ggml_tensor * fa_mask;
if (n_pad > n_tok) {
// dense padding rows: KQ_mask columns [n_tok, n_pad) -> {n_kv, n_pad - n_tok} (F16 view)
ggml_tensor * pad = ggml_view_2d(ctx0, KQ_mask, n_kv_local, n_pad - n_tok,
KQ_mask->nb[1], KQ_mask->nb[1] * n_tok);
// CPU ggml_concat only supports F16 along dim 0 (concat_any); the dim-1 row concat must be
// done in F32 (concat_f32 handles all dims), then cast the padded result to F16. On CUDA the
// F16 dim-1 concat is supported, so this path only needed adapting for the CPU build.
ggml_tensor * pad_f32 = ggml_cast(ctx0, ggml_cont(ctx0, pad), GGML_TYPE_F32);
ggml_tensor * fa_f32 = ggml_concat(ctx0, sparse, pad_f32, 1); // {n_kv, n_pad} F32
fa_mask = ggml_cast(ctx0, fa_f32, GGML_TYPE_F16); // {n_kv, n_pad} F16
} else {
fa_mask = ggml_cast(ctx0, sparse, GGML_TYPE_F16);
}
fa_mask = ggml_cont(ctx0, fa_mask);
cb(fa_mask, "dsa_fa_mask", -1);
return fa_mask;
}
// Layer-mode attention path (non-TP). Mirrors build_deepseek2_tp_attention's interface.
ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
ggml_cgraph * gf, int il,
@@ -343,6 +707,16 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
cur = llm_build_norm(ctx0, inpL, hparams, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, cb, il);
cb(cur, "attn_norm", il);
// DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). Built below from the q_lora latent
// and used to construct a sparse causal mask. Defaults to the dense KQ_mask.
// - sparse_mask : F32 additive sparse mask for the soft_max (-fa 0) path.
// - sparse_mask_fa : F16, padded variant for the ggml_flash_attn_ext (-fa 1) path.
// Both default to the dense KQ_mask so non-DSA / disabled builds are unchanged.
ggml_tensor * sparse_mask = KQ_mask;
ggml_tensor * sparse_mask_fa = KQ_mask;
//ggml_tensor * top_k = nullptr;
//(void) top_k; // captured for potential reuse/debug; only the masks are consumed downstream
// self_attention
{
ggml_tensor * q = nullptr;
@@ -391,6 +765,38 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
q = llm_build_norm(ctx0, q, hparams, model.layers[il].attn_q_a_norm, NULL, LLM_NORM_RMS, cb, il);
cb(q, "q", il);
// DSA lightning indexer (cache-backed): score the q_lora latent against the persistent
// indexer-key cache over the full n_kv, then build a sparse top-k causal mask. Correct
// for prefill AND decode (single sequence). Gate: --dsa opt-in (off by default) +
// GLM_DSA arch + indexer tensors + cache. When off, the model runs the dense MLA path,
// byte-identical to a build without this feature.
if (lctx.cparams.dsa && model.arch == LLM_ARCH_GLM_DSA && model.layers[il].indexer_attn_q_b
&& kv_self.kr_l.size() > (size_t) il && kv_self.kr_l[il]) {
// GLM-5.2 IndexShare: "full" layers compute their own lightning-indexer top-k;
// "shared" layers reuse the previous full layer's top-k (transformers reference:
// shared layer indexer=None, topk_indices=prev_topk_indices). The full/shared map is
// hparams.indexer_is_full (GGUF metadata or derived config rule). At a given step all
// layers share the same n_kv/n_tokens, so a full layer's argsort is valid to reuse.
ggml_tensor * sorted;
if (hparams.indexer_is_full[il] || dsa_last_full_sorted == nullptr) {
ggml_tensor * qr = q; // q_lora latent (after attn_q_a_norm, before wq_b)
sorted = build_deepseek2_dsa_indexer(gf, il, qr, cur, KQ_mask, inp_pos);
dsa_last_full_sorted = sorted;
} else {
sorted = dsa_last_full_sorted;
}
if (lctx.cparams.flash_attn) {
sparse_mask_fa = ::build_deepseek2_dsa_fa_mask(lctx, ctx0, KQ_mask, sorted);
} else {
sparse_mask = build_deepseek2_dsa_sparse_mask(sorted, KQ_mask);
}
//// For the FA path the mask must be F16 + padded; build it from the F32 sparse mask.
//if (lctx.cparams.flash_attn) {
// sparse_mask_fa = build_deepseek2_dsa_fa_mask(sparse_mask, KQ_mask);
//}
//top_k = sorted;
}
q = ggml_mul_mat(ctx0, model.layers[il].wq_b, q);
cb(q, "q", il);
} else {
@@ -566,7 +972,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
auto q_iter = ggml_view_3d(ctx0, q, q->ne[0], q->ne[1], n_max_head,
q->nb[1], q->nb[2], q->nb[2]*n_max_head*iter);
kqv = ggml_flash_attn_ext(ctx0, q_iter, k, v, KQ_mask, kq_scale, hparams.f_max_alibi_bias, 0.f);
kqv = ggml_flash_attn_ext(ctx0, q_iter, k, v, sparse_mask_fa, kq_scale, hparams.f_max_alibi_bias, 0.f);
if (use_f32_attn_precision || q->ne[1] <= 8) {
ggml_flash_attn_ext_set_prec(kqv, GGML_PREC_F32);
}
@@ -607,7 +1013,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
ggml_row_size(kv_self.k_l[il]->type, n_embd_head_qk_rope));
cb(kv_cache_lora, "kv_cache_lora", il);
kqv_compressed = ggml_flash_attn_ext(ctx0, q, kv_cache, kv_cache_lora, KQ_mask, kq_scale, hparams.f_max_alibi_bias, 0.f);
kqv_compressed = ggml_flash_attn_ext(ctx0, q, kv_cache, kv_cache_lora, sparse_mask_fa, kq_scale, hparams.f_max_alibi_bias, 0.f);
cb(kqv_compressed, "kqv_compressed", il);
if (use_f32_attn_precision) {
@@ -647,7 +1053,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
cb(kq, "kq_perm", il);
}
kq = ggml_soft_max_ext(ctx0, kq, KQ_mask, kq_scale, hparams.f_max_alibi_bias);
kq = ggml_soft_max_ext(ctx0, kq, sparse_mask, kq_scale, hparams.f_max_alibi_bias);
cb(kq, "kq_soft_max_ext", il);
if (!pp_opt) {
@@ -673,7 +1079,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
int this_ne12 = i_head + n_per_step <= q->ne[2] ? n_per_step : q->ne[2] - i_head;
ggml_tensor * q_i = ggml_view_3d(ctx0, q, q->ne[0], q->ne[1], this_ne12, q->nb[1], q->nb[2], q->nb[2]*i_head);
ggml_tensor * kq_i = ggml_mul_mat(ctx0, kv_cache, q_i);
kq_i = ggml_soft_max_ext(ctx0, kq_i, KQ_mask, kq_scale, hparams.f_max_alibi_bias);
kq_i = ggml_soft_max_ext(ctx0, kq_i, sparse_mask, kq_scale, hparams.f_max_alibi_bias);
ggml_tensor * kqv_i = ggml_mul_mat(ctx0, kv_cache_trans, kq_i);
if (i_head == 0) {
kqv_compressed = kqv_i;
@@ -762,6 +1168,7 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention(
}
ggml_cgraph * llm_build_context::build_deepseek2() {
dsa_last_full_sorted = nullptr; // GLM-5.2 IndexShare: reset shared-layer top-k reuse state (before any layer, incl. the MTP early-return path)
const bool tp_mode = (model.split_mode == LLAMA_SPLIT_MODE_GRAPH ||
model.split_mode == LLAMA_SPLIT_MODE_ATTN);
#ifdef GGML_USE_VULKAN
@@ -794,6 +1201,19 @@ ggml_cgraph * llm_build_context::build_deepseek2() {
// KQ_mask (mask for 1 head, it will be broadcasted to all heads)
struct ggml_tensor * KQ_mask = build_inp_KQ_mask();
if (lctx.cparams.dsa && model.arch == LLM_ARCH_GLM_DSA) {
static const int n_sink = []{ const char * e = getenv("DSA_SINK"); return e ? atoi(e) : 1; }();
if (n_sink > 0 && n_sink < (int) n_kv) {
lctx.inp_dsa_sink = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens);
cb(lctx.inp_dsa_sink, "dsa_sink", -1);
ggml_set_input(lctx.inp_dsa_sink);
}
auto minus_inf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, KQ_mask->ne[0], n_tokens);
minus_inf = ggml_fill_inplace(ctx0, minus_inf, -INFINITY);
ggml_build_forward_expand(gf, minus_inf);
lctx.inp_mask_inf = minus_inf;
}
// whether to use n_tokens as the matrix dimension during multiplication or n_head
// n_tokens is higher during prompt processing, this allows to optimize for this case
bool pp_opt = n_tokens >= 128 && lctx.cparams.mla_attn > 1;
src/llama-build-context.cpp
+116
View File
@@ -116,6 +116,7 @@ void llm_build_context::init() {
lctx.inp_pos_bucket = nullptr;
lctx.inp_embd_enc = nullptr;
lctx.inp_KQ_mask_cross = nullptr;
lctx.inp_dsa_sink = nullptr;
lctx.dflash.inputs.target_features = nullptr;
lctx.dflash.inputs.pos_ctx = nullptr;
lctx.dflash.inputs.kq_mask = nullptr;
@@ -194,6 +195,103 @@ ggml_cgraph * llm_build_context::build_k_shift() {
ggml_build_forward_expand(gf, tmp);
}
// DSA lightning-indexer key cache (GLM-5.2 / DeepSeek-V3.2): the persistent indexer keys in
// kv_self.kr_l[il] are RoPE-position-encoded at write time (build_deepseek2_dsa_indexer rotates
// the first rope_dim dims with the cell's absolute pos), so a context-shift that re-RoPEs the
// main K cache MUST apply the identical delta-rotation to the indexer keys, or query/key
// rotations desync and the indexer top-k silently degrades. We mirror the main K-shift here.
//
// Subtlety: the cached key is H * concat(RoPE(k_pe,pos), k_nope), where H is the symmetric
// orthonormal Walsh-Hadamard rotation applied AFTER RoPE (head-wide, mixing pe+nope dims). So we
// cannot RoPE the cached key directly. We un-rotate by H (H==H^T, H*H==I), RoPE-delta the pe
// sub-block, then re-rotate by H. This is exact because (a) for this model ext_factor==0 and
// attn_factor==1 (no YaRN: GLM-DSA carries no rope-scaling metadata), so NEOX RoPE is a pure
// rotation and composable: RoPE(x,pos+delta)==RoPE(RoPE(x,pos),delta); and (b) when the Hadamard
// is disabled (DSA_HADAMARD_DISABLE) the un/re-rotate by H is simply absent (kr_had==false).
//
// NOTE (engine gate): build_k_shift only runs after get_can_shift() passes, which currently
// returns FALSE for ALL MLA models (llama.cpp llama_kv_cache_update_internal -> get_can_shift:
// is_mla_model() includes GLM_DSA). So on GLM-5.2 a RoPE context-shift is refused by the engine
// (decode returns 1 == "failed to eval") BEFORE any kr_l graph is built, identically with or
// without the indexer. This block is therefore correct-and-dormant: it never executes on the
// current MLA path, but keeps the indexer keys consistent the instant MLA K-shift is ever enabled.
// It is NOT a way for the indexer to desync — the shift simply does not happen for MLA.
if (model.arch == LLM_ARCH_GLM_DSA && hparams.indexer_head_size > 0 && !kv_self.kr_l.empty()) {
const int64_t head_size = hparams.indexer_head_size;
const int64_t rope_dim = n_rot; // n_embd_head_qk_rope (indexer pe dims)
const int64_t nope_dim = head_size - rope_dim;
// Hadamard size used by the forward indexer: largest power-of-2 divisor of head_size that
// spans the full head row (else the forward path skips it). Rebuild the same matrix here so
// we can un/re-rotate the cached key. Must match build_deepseek2_dsa_indexer exactly.
bool do_hadamard = false;
static const bool dsa_had_disable = getenv("DSA_HADAMARD_DISABLE") != nullptr;
if (lctx.cparams.dsa_indexer_hadamard && !dsa_had_disable) {
GGML_ASSERT((head_size & ~(head_size-1)) == head_size);
do_hadamard = true;
}
int64_t nrot = 1;
while ((nrot * 2) <= head_size && head_size % (nrot * 2) == 0) nrot *= 2;
for (int il = 0; il < n_layer; ++il) {
if ((size_t) il >= kv_self.kr_l.size() || kv_self.kr_l[il] == nullptr) {
continue;
}
ggml_tensor * kr = kv_self.kr_l[il]; // {head_size, kv_size} F16
// The indexer forward RoPE (build_deepseek2_dsa_indexer) passes rope_factors==nullptr;
// mirror that exactly here so the delta-rotation matches the encoding.
struct ggml_tensor * rope_factors = nullptr;
// {head_size, kv_size} -> work in F32 for the rotation, write back F16.
ggml_tensor * kr_f32 = ggml_cast(ctx0, kr, GGML_TYPE_F32);
for (auto * backend : lctx.backends) {
if (ggml_backend_supports_buft(backend, lctx.model.buft_layer[il].buft)) {
ggml_backend_sched_set_tensor_backend(lctx.sched, kr_f32, backend);
break;
}
}
cb(kr_f32, "kr_f32", il);
// un-Hadamard: H * kr == concat(RoPE(k_pe,pos), k_nope) (H symmetric/orthonormal).
if (do_hadamard) {
kr_f32 = ggml_hadamard(ctx0, kr_f32, head_size);
cb(kr_f32, "kr_unhad", il);
}
// view the pe sub-block {rope_dim, 1, kv_size} and the nope sub-block {nope_dim,1,kv_size}.
ggml_tensor * kr_pe = ggml_view_3d(ctx0, kr_f32, rope_dim, 1, n_ctx,
ggml_row_size(kr_f32->type, head_size),
ggml_row_size(kr_f32->type, head_size), 0);
ggml_tensor * kr_nope = ggml_view_3d(ctx0, kr_f32, nope_dim, 1, n_ctx,
ggml_row_size(kr_f32->type, head_size),
ggml_row_size(kr_f32->type, head_size),
ggml_row_size(kr_f32->type, rope_dim));
// RoPE the pe block by the per-cell delta (inp_K_shift holds cells[i].delta) into a NEW
// tensor (non-in-place; ggml_cont copies the view first to avoid aliasing the source).
// NEOX, identical params to the indexer forward RoPE.
ggml_tensor * kr_pe_rot = ggml_rope_ext(ctx0, ggml_cont(ctx0, kr_pe),
lctx.inp_K_shift, rope_factors, n_rot, LLAMA_ROPE_TYPE_NEOX, n_ctx_orig,
freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
cb(kr_pe_rot, "kr_pe_shifted", il);
// reassemble {head_size, 1, kv_size} = concat(rotated pe, untouched nope) along dim0.
ggml_tensor * kr_cat = ggml_concat(ctx0, kr_pe_rot, ggml_cont(ctx0, kr_nope), 0);
kr_cat = ggml_reshape_2d(ctx0, kr_cat, head_size, n_ctx);
cb(kr_cat, "kr_cat_shifted", il);
// re-Hadamard: H * concat(RoPE(k_pe,pos+delta), k_nope) == the new cached key.
ggml_tensor * kr_new = kr_cat;
if (do_hadamard) {
kr_new = ggml_hadamard(ctx0, kr_cat, head_size);
cb(kr_new, "kr_rehad", il);
}
// write back into the F16 cache.
ggml_tensor * kr_back = ggml_cpy(ctx0, kr_new, kr);
cb(kr_back, "kr_shifted", il);
ggml_build_forward_expand(gf, kr_back);
}
}
return gf;
}
@@ -302,6 +400,24 @@ ggml_cgraph * llm_build_context::build_defrag(const std::vector<uint32_t> & ids)
if (view_v_src && view_v_dst) {
ggml_build_forward_expand(gf, ggml_cpy(ctx0, view_v_src, view_v_dst));
}
// DSA lightning-indexer key cache: move the indexer keys alongside k_l. Each cell's
// indexer key (kr_l row, one per cache cell) must follow its cell, or after defrag the
// kr_l rows no longer match the cells the indexer scores by cell index. The indexer key
// is RoPE-encoded at its (unchanged) absolute pos, and defrag does NOT change pos, so a
// plain row move is correct (no re-RoPE needed; only seq_add/K-shift change pos).
if ((size_t) il < kv_self.kr_l.size() && kv_self.kr_l[il] != nullptr) {
const int64_t head_size = hparams.indexer_head_size;
ggml_tensor * view_kr_src = ggml_view_2d(ctx0, kv_self.kr_l[il],
head_size, nm,
ggml_row_size(kv_self.kr_l[il]->type, head_size),
ggml_row_size(kv_self.kr_l[il]->type, head_size*i));
ggml_tensor * view_kr_dst = ggml_view_2d(ctx0, kv_self.kr_l[il],
head_size, nm,
ggml_row_size(kv_self.kr_l[il]->type, head_size),
ggml_row_size(kv_self.kr_l[il]->type, head_size*id));
ggml_build_forward_expand(gf, ggml_cpy(ctx0, view_kr_src, view_kr_dst));
}
}
i += nm - 1;
src/llama-build-context.h
+25
View File
@@ -97,6 +97,10 @@ struct llm_build_context {
struct ggml_context * ctx0 = nullptr;
// GLM-5.2 IndexShare: the most-recent "full" indexer layer's top-k selection (argsort indices).
// "shared" layers reuse this instead of computing their own. Reset to nullptr at each graph build.
struct ggml_tensor * dsa_last_full_sorted = nullptr;
// TODO: consider making the entire interface noexcept
llm_build_context(
llama_context & lctx,
@@ -287,6 +291,27 @@ struct llm_build_context {
bool is_lite,
bool pp_opt);
// DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). Cache-backed (persistent indexer-key cache).
// Returns the FULL descending argsort of the per-query indexer scores [n_kv, n_tokens] (I32).
ggml_tensor * build_deepseek2_dsa_indexer(
ggml_cgraph * gf,
int il,
ggml_tensor * qr, // q_lora latent [q_lora_rank, n_tokens] (after attn_q_a_norm)
ggml_tensor * cur, // attn_norm output [n_embd, n_tokens]
ggml_tensor * KQ_mask, // F32 causal mask [n_kv, n_tokens_pad]
ggml_tensor * inp_pos);
// Build the additive sparse causal mask from the full score ranking + the base causal KQ_mask.
ggml_tensor * build_deepseek2_dsa_sparse_mask(
ggml_tensor * sorted, // [n_kv, n_tokens] (I32) full descending argsort of scores
ggml_tensor * KQ_mask); // F32 causal mask [n_kv, n_tokens_pad]
// Adapt the (F32, unpadded) sparse mask to the shape/dtype ggml_flash_attn_ext requires on this
// fork: F16, contiguous, ne[1] padded to GGML_KQ_MASK_PAD (== the dense KQ_mask's padded shape).
ggml_tensor * build_deepseek2_dsa_fa_mask(
ggml_tensor * sparse, // [n_kv, n_tokens] (F32) additive sparse causal mask
ggml_tensor * KQ_mask); // FA dense mask [n_kv, n_tokens_pad] (F16 when -fa 1)
ggml_cgraph * build_glm4_moe();
ggml_cgraph * build_bitnet();
src/llama-context.h
+16
View File
@@ -61,6 +61,12 @@ struct llama_kv_cache {
std::vector<struct ggml_tensor *> v_l;
std::vector<struct ggml_tensor *> 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.
std::vector<struct ggml_tensor *> kr_l;
// When true, the delta_net graph builder will enable per-step SSM state saves
bool save_per_step_ssm = false;
@@ -378,6 +384,8 @@ 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_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;
ggml_backend_t ggml_backend_by_name(const char * name);
@@ -393,6 +401,14 @@ struct llama_context {
size_t step = 0;
};
std::vector<CacheCopy> cache_copies;
// GLM-DSA lightning indexer: the indexer-key cache (kr_l) write is a separate ggml_cpy that
// the K/V cache_copies fixup does NOT cover. Under graph reuse (FA pads KV to 256, so n_kv
// stays constant across consecutive decode ubatches and the graph IS reused) its view_offs
// would stay baked at the first ubatch's kv_head, scattering this ubatch's indexer keys to a
// stale slot. Later ubatches never populate their own recent index-key cells (those cells read
// 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<CacheCopy> dsa_cache_copies;
bool update_cache_copies();
src/llama-cparams.h
+3
View File
@@ -41,6 +41,9 @@ struct llama_cparams {
bool graph_reuse;
bool k_cache_hadamard;
bool v_cache_hadamard;
bool dsa_indexer_hadamard = true; // apply Walsh-Hadamard rotation to DSA indexer q/k (precision)
bool dsa = false; // enable GLM DSA sparse attention (off by default; opt-in via --dsa)
int dsa_top_k = -1; // DSA top-k override (<0 => use the model's configured indexer_top_k)
bool split_mode_graph_scheduling;
//bool split_mode_f16;
bool scheduler_async;
src/llama-hparams.cpp
+19
View File
@@ -1582,6 +1582,14 @@ void llm_load_hparams(
{
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
// GLM-DSA lightning-indexer k_norm is a (non-RMS) LayerNorm built via LLM_NORM,
// which uses hparams.f_norm_eps in ggml_norm(). The GGUF only carries the RMS eps,
// so f_norm_eps stays 0 and CPU ggml_norm aborts (GGML_ASSERT(eps > 0)). On CUDA
// the kernel does not assert (eps=0 is numerically tolerable), which is why the
// CPU attention path was never exercised. Mirror the RMS eps so the indexer
// LayerNorm gets a valid epsilon on all backends. (HF hardcodes 1e-6 for this
// LayerNorm, but 1e-6 vs 1e-5 is within 4-chunk PPL noise here, so keep the mirror.)
if (hparams.f_norm_eps <= 0.0f) hparams.f_norm_eps = hparams.f_norm_rms_eps;
ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false);
// MoE parameters
@@ -1605,6 +1613,17 @@ void llm_load_hparams(
ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
// GLM-5.2 IndexShare: per-layer full/shared indexer map. "full" layers compute their own
// top-k; "shared" layers reuse the previous full layer's selection (transformers
// modeling_glm_moe_dsa.py: shared layer indexer=None, topk_indices=prev_topk_indices).
// Derived from GLM-5.2's config indexer_types rule (full iff il<=1 or il%4==2), verified
// to reproduce the config's full set {0,1,2,6,10,...} exactly. Existing GGUFs carry no
// per-layer metadata, so the derivation is the source of truth; a future metadata key can
// override this for GLM-DSA variants with a different pattern.
for (uint32_t il = 0; il < hparams.n_layer; ++il) {
hparams.indexer_is_full[il] = (il <= 1) || (il % 4 == 2);
}
// Expert gating function (GLM-4.5 uses sigmoid)
ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);
if (hparams.expert_gating_func == LLM_EXPERT_GATING_FUNC_TYPE_NONE) {
src/llama-hparams.h
+4
View File
@@ -126,6 +126,10 @@ struct llama_hparams {
uint32_t indexer_n_head = 0;
uint32_t indexer_head_size = 0;
uint32_t indexer_top_k = 0;
// GLM-5.2 IndexShare: per-layer full/shared indexer map. "full" layers compute their own lightning-
// indexer top-k; "shared" layers reuse the previous full layer's top-k. Populated from GGUF
// indexer_types metadata if present, else derived from the GLM-5.2 config rule at load time.
std::array<bool, LLAMA_MAX_LAYERS> indexer_is_full = {};
// qwen3vl deepstack
uint32_t n_deepstack_layers = 0;
src/llama-load-tensors.cpp
+12 -4
View File
@@ -2814,16 +2814,24 @@ bool create_tensors_helper::create_glm_dsa_tensors(const LLM_TN & tn) {
layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, flags);
// DSA indexer
// GLM-5.2 only ships the DSA indexer on a subset of layers; the rest omit it.
// The DSA indexer runtime is not implemented (graph is plain MLA), so these
// tensors are loaded-but-unused. Mark them optional so layers without an
// indexer load as nullptr (ported from ggml-org/llama.cpp#24770).
// DSA indexer tensors. GLM-5.2's reference marks "shared" layers indexer=None, but the GGUFs
// in circulation (e.g. unsloth) ship indexer weights on every layer, and the DSA runtime IS
// implemented (see build_deepseek2_dsa_indexer + IndexShare reuse). These are TENSOR_NOT_REQUIRED
// so a faithful conversion that omits shared-layer indexer weights still loads: full layers must
// have them (they compute the top-k), shared layers reuse the previous full layer's top-k. A
// "full" layer (per hparams.indexer_is_full) that loads a null indexer will warn below.
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_k_norm_b = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "bias", 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);
// IndexShare: a "full" layer must carry indexer weights (it computes the top-k that shared layers
// reuse). If a conversion omits them on a full layer, DSA silently falls back to dense there.
if (i < (int) hparams.n_layer && hparams.indexer_is_full[i] && !layer.indexer_attn_q_b) {
LLAMA_LOG_WARN("%s: GLM-DSA layer %d is 'full' per indexer_types but has no indexer weights; DSA will run dense on it\n", __func__, i);
}
layer.ffn_norm = create_tensor(norm_ctx, tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
if (i < (int) hparams.n_layer_dense_lead) {
src/llama.cpp
+118 -3
View File
@@ -679,6 +679,28 @@ 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;
}
@@ -695,6 +717,9 @@ 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_cache_copies.resize(hparams.n_layer);
llama_all_contexts().push_back(this);
}
@@ -950,6 +975,13 @@ 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) {
cache.kr_l.resize(n_layer, nullptr);
}
std::vector<size_t> mem_split(model.splits.size(), 0);
const uint32_t qnext_state_slots = llama_qwen3next_state_slots(cparams, kv_size);
@@ -995,6 +1027,13 @@ static bool llama_kv_cache_init(
ggml_tensor * kv = ggml_new_tensor_2d(ctx, primary_kv_type, kv_lora_rank + n_embd_head_qk_rope, kv_size);
ggml_format_name(kv, "cache_k_l%d", i);
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 && !is_mtp_tail_layer) {
ggml_tensor * kr = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hparams.indexer_head_size, kv_size);
ggml_format_name(kr, "cache_kr_l%d", i);
cache.kr_l[i] = kr;
}
if (!cparams.flash_attn && cparams.mla_attn == 1) {
ggml_tensor * kvt = ggml_new_tensor_1d(ctx, cache.type_v, kv_lora_rank*kv_size);
ggml_format_name(kvt, "cache_v_l%d", i);
@@ -4284,6 +4323,58 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) {
const auto & cparams = lctx.cparams;
const auto & kv_self = lctx.kv_self;
if (lctx.inp_dsa_sink) {
// Per-sequence attention-sink boost for the DSA lightning indexer top-k selection.
// inp_dsa_sink {n_kv, n_tokens}: 1e20 iff key cell i is one of query j's sequence's FIRST
// n_sink present tokens, else 0. Filled from kv_self.cells like the KQ_mask so it is correct
// for multi-sequence ubatches (each sequence's OWN sink is force-included).
//
// SERVING FIX: anchor the sink on each sequence's FIRST PRESENT position, not absolute
// pos < n_sink. After multi-turn seq_rm drops a sequence's early tokens, its earliest
// surviving token has pos >= n_sink; an absolute "pos < n_sink" test would then protect
// NOTHING for that sequence and let the (now-)sink token be masked out of top-k, collapsing
// it. Anchoring on per-sequence min(pos) keeps the sink protection following the sequence's
// actual first present cell. For a fresh sequence starting at pos 0, min(pos)==0 so the
// boosted set is identical to the old behaviour (n_seq==1 byte-identical).
GGML_ASSERT(ggml_backend_buffer_is_host(lctx.inp_dsa_sink->buffer));
static const int n_sink = []{ const char * e = getenv("DSA_SINK"); return e ? atoi(e) : 1; }();
const int64_t n_kv = lctx.inp_dsa_sink->ne[0];
const int64_t n_tok_idx = lctx.inp_dsa_sink->ne[1];
float * data = (float *) lctx.inp_dsa_sink->data;
std::memset(data, 0, ggml_nbytes(lctx.inp_dsa_sink));
// Per-sequence first present pos: min over present cells of that seq, restricted to the
// n_kv key span the indexer scores. Cache across query rows in this ubatch.
// ASSUMPTION (latent today): min(pos) over the whole [0,n_kv) span equals the sequence's
// genuine first-present token only for a NON-SLIDING cache. GLM_DSA has no SWA, so the
// oldest cell of a seq is its true sink; if SWA were ever added, the window could evict the
// real sink and min(pos) would anchor on the wrong (window-floor) cell — revisit then.
std::unordered_map<llama_seq_id, llama_pos> seq_min_pos;
for (int64_t i = 0; i < n_kv; ++i) {
const auto & cell = kv_self.cells[i];
if (cell.pos < 0) continue;
for (const llama_seq_id sid : cell.seq_id) {
auto it = seq_min_pos.find(sid);
if (it == seq_min_pos.end() || cell.pos < it->second) {
seq_min_pos[sid] = cell.pos;
}
}
}
for (int64_t j = 0; j < n_tok_idx && j < (int64_t) batch.n_tokens; ++j) {
const llama_seq_id seq_id = batch.seq_id[j][0];
auto it = seq_min_pos.find(seq_id);
if (it == seq_min_pos.end()) continue; // no present cell for this seq in the key span
const llama_pos first_pos = it->second;
for (int64_t i = 0; i < n_kv; ++i) {
const auto & cell = kv_self.cells[i];
if (cell.pos >= first_pos && cell.pos < first_pos + n_sink && cell.has_seq_id(seq_id)) {
data[j*n_kv + i] = 1e20f;
}
}
}
}
if (batch.token && lctx.inp_tokens) {
#if IK_PRINT_TIMING == 2
auto tim1 = ggml_time_us();
@@ -5932,7 +6023,12 @@ static void llama_kv_cache_defrag_internal(struct llama_context & lctx) {
// - x2 for keys and values
//const uint32_t max_moves = model.max_nodes()/(6*n_layer);
// TODO: tmp fix https://github.com/ggerganov/llama.cpp/issues/6685#issuecomment-2057579516
const uint32_t max_moves = (lctx.model.max_nodes(1) - 2*n_layer)/(6*n_layer);
// DSA: build_defrag additionally moves the indexer-key cache (kr_l), +3 tensors/layer/move
// (src view, dst view, copy), so budget 9*n_layer per move when the indexer cache is present.
const bool has_dsa_indexer_defrag =
lctx.model.arch == LLM_ARCH_GLM_DSA && !kv_self.kr_l.empty();
const uint32_t tensors_per_move = has_dsa_indexer_defrag ? 9 : 6;
const uint32_t max_moves = (lctx.model.max_nodes(1) - 2*n_layer)/(tensors_per_move*n_layer);
// determine which KV cells to move where
//
@@ -6541,6 +6637,8 @@ struct llama_context_params llama_context_default_params() {
/*.fused_mmad =*/ true,
/*.rope_cache =*/ false,
/*.graph_reuse =*/ true,
/*.dsa =*/ false,
/*.dsa_top_k =*/ -1,
/*.min_experts =*/ -1,
/*.thtesh_experts =*/ 0.0f,
/*.only_active_experts =*/ false,
@@ -6958,6 +7056,14 @@ struct llama_context * llama_init_from_model(
cparams.fused_mmad = params.fused_mmad;
cparams.rope_cache = params.rope_cache;
cparams.graph_reuse = params.graph_reuse;
cparams.dsa = params.dsa;
cparams.dsa_top_k = params.dsa_top_k;
// The DSA lightning indexer is built only in the layer-mode (non-TP) attention path. Under
// -sm graph / -sm attn the model runs the tensor-parallel attention path, which has no indexer,
// so --dsa would silently run dense MLA. Warn instead of degrading silently.
if (cparams.dsa && (model->split_mode == LLAMA_SPLIT_MODE_GRAPH || model->split_mode == LLAMA_SPLIT_MODE_ATTN)) {
LLAMA_LOG_WARN("%s: --dsa is not active under -sm graph/attn (tensor-parallel attention has no indexer); running dense MLA\n", __func__);
}
cparams.k_cache_hadamard = params.k_cache_hadamard;
cparams.v_cache_hadamard = params.v_cache_hadamard;
// Folding H into wv_b/wk_b_pp permanently mutates the model; a later context
@@ -7316,12 +7422,18 @@ struct llama_context * llama_init_from_model(
{
size_t memory_size_k = 0;
size_t memory_size_v = 0;
size_t memory_size_k_indexer = 0;
for (auto & k : ctx->kv_self.k_l) {
if (k) {
memory_size_k += ggml_nbytes(k);
}
}
for (auto & k : ctx->kv_self.kr_l) {
if (k) {
memory_size_k_indexer += ggml_nbytes(k);
}
}
for (auto & v : ctx->kv_self.v_l) {
if (v) {
@@ -7329,8 +7441,8 @@ struct llama_context * llama_init_from_model(
}
}
if (memory_size_k + memory_size_v > 0) {
if (cparams.mla_attn != 0 && !cparams.flash_attn) {
if (memory_size_k + memory_size_v) {
if (cparams.mla_attn != 0 && !cparams.flash_attn) {
LLAMA_LOG_INFO("%s: KV self size = %7.2f MiB, c^KV (%s): %7.2f MiB, kv^T (%s): %7.2f MiB\n", __func__,
(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f),
ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),
@@ -7346,6 +7458,9 @@ struct llama_context * llama_init_from_model(
ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));
}
}
if (memory_size_k_indexer > 0) {
LLAMA_LOG_INFO("%s: KV self indexer size = %7.2f MiB\n", __func__, memory_size_k_indexer/1024./1024.);
}
}
// graph outputs buffer