Skip to content

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

Merged
ikawrakow merged 21 commits into
ikawrakow:mainfrom
mb8565:glm-dsa-upstream
Jul 2, 2026
Merged

deepseek2 : GLM-DSA sparse attention (lightning indexer), --dsa off by default#2045
ikawrakow merged 21 commits into
ikawrakow:mainfrom
mb8565:glm-dsa-upstream

Conversation

@mb8565

@mb8565 mb8565 commented Jun 27, 2026

Copy link
Copy Markdown
Contributor

What this does

This implements the GLM-5.2 / DeepSeek-V3.2 "lightning indexer" sparse-attention compute. #2017 (merged) makes the GLM-DSA indexer tensors load, but nothing consumes them yet, so GLM-5.2 currently runs the plain dense MLA graph with the indexer weights sitting unused. This adds the missing compute.

The indexer scores each query against the (Hadamard-rotated) indexer keys, keeps the top-k highest-scoring keys per query, and masks the rest out of attention. It is built inside the deepseek2 graph (build_deepseek2_dsa_indexer / build_deepseek2_dsa_sparse_mask / build_deepseek2_dsa_fa_mask in src/graphs/build_deepseek2.cpp), with a persistent per-layer indexer-key cache (kr_l, F16, [head_size, kv_size]) so a decoded token scores against all past indexer keys, not just the current batch. Both attention paths are handled: the soft_max path (-fa 0) consumes an F32 additive sparse mask, and the flash-attention path (-fa 1) consumes an F16 padded variant of the same mask.

The flag, off by default

Per the direction in #2040, the feature is a command-line argument, not an environment variable, and it is off by default:

  • --dsa enables it (GLM-DSA arch only).
  • --dsa-top-k N overrides the kept-key count; <0 uses the model's configured indexer_top_k.

This follows the maintainer's stated criteria in #2040 ("it should be added under the DSA arch ... turn it on via command line argument ... off by default for now"). The earlier env-var control (DSA_INDEXER_DISABLE) is gone. Two debug-only env knobs remain (DSA_HADAMARD_DISABLE, DSA_SINK) for bisection during development; neither is a primary control and both have safe defaults.

No regression

The change is arch-gated to LLM_ARCH_GLM_DSA and the --dsa flag is off by default, so for every other model and for GLM-DSA without --dsa the code takes the existing dense path unchanged. The DSA branch in build_deepseek2_layer_attention is reached only when cparams.dsa && model.arch == LLM_ARCH_GLM_DSA and the indexer tensors and cache are present.

Proof: with --dsa off, the branch is byte-equal in perplexity to clean upstream main across a zoo, including the deepseek2 model that exercises the exact file this PR touches. Hardware: Intel Xeon Platinum 8260, 1x Tesla P100-PCIE-16GB (sm_60), CUDA build. Baseline = upstream main at the branch merge-base (b84902d). Both binaries: -ngl 99 -fa 1 -c 512 -b 512 --chunks 8 --seed 1234, wikitext-2 test.

Model Arch PPL main PPL branch (--dsa off) Match
Qwen3.5-4B Q4_K_M qwen3 (dense/GQA) 10.0809 10.0809 byte-equal
Gemma-4-12B-qat Q4_K_XL gemma4 (GQA) 310.3393 310.3393 byte-equal
DeepSeek-V2-Lite-REAP48 Q4_K_M deepseek2 (MoE) 31.9774 31.9774 byte-equal

(The Gemma QAT chat model on raw wikitext is off-distribution, hence the high PPL; the point is the two binaries agree to the digit.)

Throughput is unchanged. llama-bench, same hardware, -ngl 99 -fa 1 -p 512 -n 128, interleaved runs to cancel GPU clock drift:

Model PP t/s main PP t/s branch TG t/s main TG t/s branch
DeepSeek-V2-Lite-REAP48 Q4_K_M 541.4 545.8 40.8 (39.6-42.1) 41.0 (40.4-42.3)
Qwen3.5-4B Q4_K_M 750.0 748.6 56.77 56.79

TG for the deepseek2 model is the mean of three interleaved 5-rep runs with the spread in parentheses; the branch-vs-main difference is inside the run-to-run variance.

Validation of the DSA path itself (--dsa on)

On GLM-5.2-UD-IQ2_M (arch glm-dsa, indexer top_k 2048), Xeon 8260 + 3x P100, --cpu-moe MoE offload, -mla 3 -fa 1, wikitext-2, 4 chunks at n_ctx 2560:

Run PPL
--dsa off (dense MLA) 2.4151
--dsa on, default top_k 2048 2.4697
--dsa on, --dsa-top-k 1024 3.5107

DSA-on is coherent and close to dense at the model's configured top_k, and PPL climbs as top_k tightens, which is what a sparse selector should do. Other validations I ran during development:

  • Selection-equivalence: the indexer top-k selection matches the HF reference implementation with diff 0 on a fixed prompt.
  • No-op-exact: at a context where n_kv <= top_k the sparse mask is a mathematical no-op, and DSA-on PPL is byte-identical to dense, so the mask machinery introduces no drift of its own.
  • Flash-attention decode (-fa 1) over a long prompt where the mask actively bites stays coherent with correct deep-context recall.
  • Graph-reuse fix: the persistent indexer-key cache (kr_l) write is now registered for the same kv_head fixup that update_cache_copies() applies to the K and V writes when a compute graph is reused. Without it, under flash-attention (which pads the cache so consecutive decode ubatches share n_kv and reuse the graph) later ubatches never wrote their own recent index keys. It was latent for GLM-5.2 at top_k 2048 but would bite at any tighter top_k or longer context.

Limitation: DSA and -sm graph

The 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 (build_deepseek2_tp_attention), which has no indexer, so --dsa there would silently run dense MLA. Rather than degrade silently, the model now emits a LLAMA_LOG_WARN at context creation when --dsa is set together with graph or attn split mode, and runs dense MLA. This matches the concern you raised in #2040 ("if adding DSA ... interferes with graph parallel ... a new arch can be added"). Wiring the indexer into the tensor-parallel path is deferred; it needs a working multi-GPU P2P test platform, which this rig does not have (peer DMA is corrupt on the Sky Lake-E root complex here).

test-backend-ops

./tests/test-backend-ops is not wired into this fork's CMake build (no target), so it cannot be run here. The ground it would cover (op correctness on the touched path) is covered by the byte-equal no-regression PPL table above and by the selection-equivalence and no-op-exact checks: with --dsa off the graph is the existing dense graph op-for-op, and with --dsa on the indexer ops are validated against the HF reference and against the dense no-op floor.

Co-authored-by: Claude Opus 4.8 (1M context) noreply@anthropic.com

@usrlocalben

usrlocalben commented Jun 28, 2026

Copy link
Copy Markdown
Contributor

Context length vs Quant gives surprising results re: VRAM buffers, I'm unable to start it with any useful context length:

sokann IQ2_KT
ctx = 100K -> +138GB VRAM, OOM
ctx = 500K -> +686GB VRAM, OOM

SixVolts Q3_K_M
ctx = 100K -> +227GB VRAM, OOM
ctx = 500K -> +1,131GB VRAM, OOM

<snip>
llama_init_from_model: n_ctx         = 100096
llama_init_from_model: n_batch       = 4096
llama_init_from_model: n_ubatch      = 4096
llama_init_from_model: flash_attn    = 1
llama_init_from_model: mla_attn      = 3
llama_init_from_model: attn_max_b    = 512
llama_init_from_model: fused_moe     = 1
llama_init_from_model: grouped er    = 0
llama_init_from_model: fused_up_gate = 1
llama_init_from_model: fused_mmad    = 1
llama_init_from_model: rope_cache    = 0
llama_init_from_model: graph_reuse   = 1
llama_init_from_model: k_cache_hadam = 1
llama_init_from_model: v_cache_hadam = 0
llama_init_from_model: split_mode_graph_scheduling = 0
llama_init_from_model: reduce_type   = f16
llama_init_from_model: sched_async   = 0
llama_init_from_model: ser           = -1, 0
llama_init_from_model: freq_base     = 8000000.0
llama_init_from_model: freq_scale    = 1
llama_kv_cache_init:      CUDA0 KV buffer size =  5128.48 MiB
llama_init_from_model: KV self size  = 4615.25 MiB, c^KV (q8_0): 4615.25 MiB, kv^T: not used
llama_init_from_model:        CPU  output buffer size =     0.61 MiB
ggml_backend_cuda_buffer_type_alloc_buffer: allocating 137755.08 MiB on device 0: cudaMalloc failed: out of memory
ggml_gallocr_reserve_n: failed to allocate CUDA0 buffer of size 144446668800
llama_init_from_model: failed to allocate compute buffers
<snip>

invocation:

N_NUMA=8
# M=/model/GLM-5.2/sokann/IQ2_KT/GLM-5.2-GGUF-2.244bpw.gguf
M=/model/GLM-5.2/SixVolts/GLM-5.2-Q3_K_M-00001-of-00008.gguf
GGML_CUDA_NO_PINNED=1 ./build/bin/llama-server \
  --host 0.0.0.0 --port 4972 --webui llamacpp \
  --numa distribute \
  -t $[ $N_NUMA * 16 ] \
  -b 4096 -ub 4096 -amb 512 \
  --cache-ram 200000 \
  -ngl 999 -cmoe \
  -mla 3 \
  --dsa \
  -muge \
  --jinja --parallel-tool-calls \
  --chat-template-kwargs '{"enable_thinking": true, "reasoning_effort": "high"}' \
  --chat-template-file /model/GLM-5.2/sokann/IQ2_KT/chat_template.jinja \
  -ctk q8_0 -khad \
  -c 500000 \
  --alias GLM-5.2 \
  --spec-type mtp:n_max=4,p_min=0.5 \
  --metrics \
  -m "$M"

@mb8565

mb8565 commented Jun 28, 2026

Copy link
Copy Markdown
Contributor Author

**Thanks for the detailed report and the buffer numbers. This is a real issue on our end, not a config problem with your setup. The compute buffer scales badly with context because the indexer currently builds full per-layer intermediates sized by context length times ubatch, so it grows with both your -c and your -ub. We're working on a fix to bound that, and we're adding long-context coverage to our tests, since our validation so far only exercised short context.

As an interim workaround, lowering -ub should cut the compute buffer roughly in proportion. At -ub 4096 you're paying about 8x what -ub 512 would cost, so a smaller ubatch should get you to a usable context while we land the fix. We'll follow up here once it's up.**

@mb8565
mb8565 force-pushed the glm-dsa-upstream branch from 371a7de to 03391cb Compare June 28, 2026 17:00
Comment thread src/graphs/build_deepseek2.cpp Outdated
nrot *= 2;
}
if (nrot == head_size) { // only apply when the rotation spans a full head row
if (!lctx.inp_dsa_hadamard) {

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In ik_llama.cpp we have a dedicated op for Hadamard transformations:

    GGML_API struct ggml_tensor * ggml_hadamard(
            struct ggml_context * ctx, 
            struct ggml_tensor  * a, 
            int                   n);  

Hence, you don't need to complicate things ala llama.cpp.

// 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

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, if n_kv <= n_top_k, wouldn't we want to skip all the sparse mask computation (it will end up being exactly the same as the attention mask)? In that case all we want to add to the indexer K cache.

Comment thread src/llama.cpp Outdated
const auto & cparams = lctx.cparams;
const auto & kv_self = lctx.kv_self;

if (lctx.inp_dsa_hadamard) {

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove this (see comment about ggml_hadamard above)

@ikawrakow

Copy link
Copy Markdown
Owner

@usrlocalben

Without a dedicated indexer op, the indexer compute graph materializes a K*Q matrix, which is of size n_ubatch * n_ctx * n_indexer_head * sizeof(float), so pretty giant for the context lengths and u-batch sizes you have used. The KQ tensor then gets element-wise multiplied with with another matrix, so that may cause the presence of a second giant matrix of that size (if the scheduler cannot figure out that this multiplication can be done in-place). My guess is that the PR misses calls to ggml_build_forward_expand to help the scheduler figure out when it can reuse a previous compute buffer, and this is why we end up requiring such huge amounts of VRAM (where to sprinkle calls to ggml_build_forward_expnd while building the compute graph is a bit of a black magic rather than a precise science).

If you have followed the DeepSeekV32/DeepSeekV4 efforts in llama.cpp, they have been struggling with this as well.

@mb8565
I guess, I need to think some more how to proceed in ik_llama.cpp. I.e, try to merge the PR even iff not really useful in practice, and then improve from there, or try to get a more useful version with a dedicated indexer op before merging. Apart from the huge amount of VRAM required, having the top_k op run on the CPU is not really very useful in practice for CUDA or hybrid CUDA/CPU inference.

@mb8565

mb8565 commented Jun 29, 2026

Copy link
Copy Markdown
Contributor Author

Thanks, this is helpful. On the compute buffer: mainline ran into the same thing and landed on a dedicated fused op. fairydreaming's #21149 added a full DSA path (CPU op plus CUDA kernel) and the standalone CPU op is being upstreamed in #24231, where the indexer compute buffer drops from 168368 MiB to 5808 MiB. It needs an op because the actual fix is dropping the n_indexer_head dimension from the output, which a graph rewrite cannot do; he first tried optimizing the materialized path and called it a dead end. So your dedicated-indexer-op instinct matches where upstream ended up.

Since --dsa is off by default, our plan for this PR is to make it usable now and track the op port as the follow-up: keep the query-chunking that bounds the score buffer (upstream's merged generic path has no such bound), drop the extra cont(permute) in our head-reduce so the score materializes once like the upstream reference, add the n_kv <= n_top_k dense short-circuit you pointed out, and adopt the native ggml_hadamard op. The mainline CUDA indexer kernel uses tensor-core (wmma) instructions, so a P100/sm_60 build would need a scalar fallback or to run the op on CPU. Would merge-now-then-improve (off by default) work for you, or would you rather hold for the fused op?

@ikawrakow

Copy link
Copy Markdown
Owner

One way to reduce the compute buffer size is to use the same trick as we already have for the DeepSeek arch with mla = 3. In that case we need to multiply the entire K-cache with the wkv_b tensor, which results in a tensor of size head_size * n_kv * n_head * sizeof(float). ik_llama.cpp has the -amb command line argument, which allows us to restrict the size of this intermediate result by performing the computation iteratively, computing some maximum number of heads per iteration. See

auto kv_f32_size = model.layers[il].wkv_b->ne[1] * kv_cache_nope->ne[1] * sizeof(float) / (1024*1024);

Implementing it in that way

  • Does not require a new op
  • Reduces maximum compute buffer size for large batches n_head times (if one head per iteration is computed)
  • Will skip the iteration for token generation (where looping over heads and associated kernel launch is costly) as K*Q is sufficiently small in that case.

Do you want to try? If not, I can do it after the PR is merged.

Please remove the llama.cpp-style Hadamard implementation as per comment, and we can see from there.

@ikawrakow

Copy link
Copy Markdown
Owner

Btw, the PPL results are concerning. By default perplexity gets calculated for a context of 512, which is less than the top_k of 2048, so DSA ON or OFF should give identical results as in both cases the exact same KV cache should be used. Same applies to --dsa 1 --dsa-top-k 1024.

@saood06

saood06 commented Jun 29, 2026

Copy link
Copy Markdown
Collaborator

Btw, the PPL results are concerning. By default perplexity gets calculated for a context of 512, which is less than the top_k of 2048, so DSA ON or OFF should give identical results as in both cases the exact same KV cache should be used. Same applies to --dsa 1 --dsa-top-k 1024.

He says that's at context 2560:

wikitext-2, 4 chunks at n_ctx 2560

The earlier discussion did say it was identical at context 512.

At c512 (n_kv well under top_k) the indexer is a pure no-op:
indexer ON equals dense byte-identical across all chunks (2.1957 ± 0.12031, chunks
2.2770 / 2.8741 / 2.3956 / 2.1957, four chunks at -mla 3 -fa 1

Also as an aside you say:

One way to reduce the compute buffer size is to use the same trick as we already have for the DeepSeek arch

Are both methods to reduce compute buffer mutually exclusive? If so do you think one would be strictly better than the other? If not could they work together to bring even lower compute buffers?

@ikawrakow

Copy link
Copy Markdown
Owner

Are both methods to reduce compute buffer mutually exclusive? If so do you think one would be strictly better than the other? If not could they work together to bring even lower compute buffers?

A dedicated indexer op is theoretically strictly better than a loop over attention heads where the indexer result is constructed out of existing ops. I say theoretically, because on CUDA one will inevitably run into issues with support for older CUDA architectures when implementing a dedicated op that fuses matrix multiplications with indexer score + top_k computation. Also, in practice one may get better performance by using a loop over attention heads plus GEMM kernels that have been optimized over a long period of time (vs newly written fused GEMM + indexer score).

@ikawrakow

Copy link
Copy Markdown
Owner

Here a perplexity test run for a context of 4096 on a 13x3090 system using the same loading parameters (--fit --fit-margin 6144)

No DSA

perplexity: calculating perplexity over 70 chunks, n_ctx=4096, batch_size=4096, n_seq=1
perplexity: 28.66 seconds per pass - ETA 33.42 minutes
[1]1.1599,[2]1.6922,[3]1.8054,[4]2.0518,[5]2.3822,[6]2.4999,[7]2.6761,[8]2.9101,[9]2.8473,[10]2.6441,
[11]2.5331,[12]2.4381,[13]2.3988,[14]2.5278,[15]2.5514,[16]2.5487,[17]2.5751,[18]2.5610,[19]2.5355,[20]2.5649

With DSA

perplexity: calculating perplexity over 70 chunks, n_ctx=4096, batch_size=4096, n_seq=1
perplexity: 39.86 seconds per pass - ETA 46.50 minutes
[1]1.3924,[2]2.0037,[3]2.1177,[4]2.3641,[5]2.7374,[6]2.8920,[7]3.0897,[8]3.3439,[9]3.3217,[10]3.1300,
[11]3.0406,[12]2.9918,[13]2.9961,[14]3.1378,[15]3.1488,[16]3.1371,[17]3.1721,[18]3.1428,[19]3.0973,[20]3.1168

I.e., it is almost 40% slower, uses 5.7 GiB extra VRAM per GPU, and produces a 20% higher PPL. If I use for the no-DSA run --fit --fit-margin 2048 --gpu-fit-margin 0,6144 to better utilize the VRAM (because we don't need 5.7 GiB extra VRAM per GPU for the indexer), I get

perplexity: calculating perplexity over 70 chunks, n_ctx=4096, batch_size=4096, n_seq=1
perplexity: 25.09 seconds per pass - ETA 29.27 minutes
[1]1.1599,[2]1.6922,[3]1.8054,[4]2.0518,[5]2.3822,[6]2.4999,[7]2.6761,[8]2.9101,[9]2.8473,[10]2.6441,
[11]2.5331,[12]2.4381,[13]2.3988,[14]2.5278,[15]2.5514,...

so DSA is 1.6X slower than that.

I suspect that at sufficiently long context DSA will become competitive with no-DSA performance wise, but with the current VRAM usage I cannot go anywhere near such a long context.

@saood06

saood06 commented Jun 29, 2026

Copy link
Copy Markdown
Collaborator

Here a perplexity test run for a context of 4096 on a 13x3090 system using the same loading parameters (--fit --fit-margin 6144)

I'm not sure PPL is ideal for seeing the results of how much response quality changes using MLA on a DSA native model. Using MLA clearly works (otherwise no one would run V3.2/Speciale/GLM-5[.1/.2]). But the difference does show up at high context, high complexity situations. See this:

image

(source: fairydreaming)

I suspect that at sufficiently long context DSA will become competitive with no-DSA performance wise, but with the current VRAM usage I cannot go anywhere near such a long context.

I know that Deepseek showed that for them the cutoff was ~8K for decode and ~16K for prefill

image

But also I heard people talking saying that vLLM and SGlang implementations were worse and thus couldn't replicate the above graphs (but this was a while ago, so maybe they have optimized their implementation since then).

Only bringing this up because I think that a sufficiently optimized implementation would be an obvious default because of the improved quality at high context, and the performance should be better except for very low context situations (and in general people tend to use smaller models for those low context use cases).

@ikawrakow

Copy link
Copy Markdown
Owner

See PR #2058 that shows how one can calculate the indexer score via a loop over attention heads. This not only reduces required memory, but also improves performance.

I did not check TG performance, but my guess is that when the batch size is less than some threshold it would be better to use the original implementation in this PR.

@mb8565

mb8565 commented Jun 29, 2026

Copy link
Copy Markdown
Contributor Author

Thanks, the per-head loop looks like exactly the memory reduction we were after, and a cleaner way to get it than what we were attempting.

We did build ik/dsa_loop and run it CPU-only (-ngl 0). It aborts on the indexer key-norm, GGML_ASSERT(eps > 0.0f). The glm-dsa hparam loader only loads f_norm_rms_eps, so f_norm_eps stays 0, which only bites on CPU (the same path runs fine on CUDA). We have a fix for that, plus a couple of other CPU-side gaps we hit getting the DSA path running -ngl 0. We can fold them in whichever way is cleanest for you: rebase #2045 onto your loop change, or send them as a small follow-up.

@mb8565

mb8565 commented Jun 29, 2026

Copy link
Copy Markdown
Contributor Author

Confirmed and fixed. We built ik/dsa_loop, put our CPU-only fixes on top, and validated it on both backends. The four divergences were the indexer key-norm eps=0 (above); concat for non-F32 types only supporting dim 0 on the CPU; ggml_set_rows with an F32 destination (no from_float trait, a NULL deref); and the DSA sparse-mask add of the F16 KQ_mask to the F32 mask. The last three look like general CPU-kernel gaps that only the DSA path reaches today.

Validation (GLM-5.2 IQ2_M, wikitext): no-op-exact holds (at n_ctx 512, below the 2048 top_k, DSA-on equals DSA-off, PPL 2.1870 both); CPU-only 4-chunk PPL at 4096 is 3.18 with DSA vs 2.70 without, a +18% that lines up with the +21% you measured on the GPU; a single-chunk 3x P100 hybrid run lands at 2.26 against the CPU 2.30, within the expected cross-backend spread. It is one small commit, 3 files, +56/-15, on your loop.

Happy to land it however suits you: rebase #2045 onto your loop, or a small follow-up against #2058.

@usrlocalben

Copy link
Copy Markdown
Contributor

@mb8565 does it give the /32 improvement in buffer size? @ikawrakow 's #2058 doesn't for me, not clear why. I also still don't understand why the quant (model) used changes the ram requirement either.

Meanwhile, mainline merged dsv4 support and its similarly enormous buffers. It looks like it will be resolved with a custom op.

mb8565 pushed a commit to mb8565/ik_llama.cpp that referenced this pull request Jun 30, 2026
PR ikawrakow#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>
mb8565 added a commit to mb8565/ik_llama.cpp that referenced this pull request Jun 30, 2026
GLM-DSA: fix CPU-only crashes in the sparse-attention path (ikawrakow#2045 was GPU-validated only)
@mb8565

mb8565 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

The CPU-only fixes are done and validated on the head-loop branch (your #2058 plus four fixes): CPU-only -fa 1 PPL 3.18 at 4K and coherent decode. The fixes do not touch the GPU path, and the 3x P100 hybrid is unchanged at 2.26. mgkwill independently reproduced the same four from a clean CPU-only build and opened a PR against our branch; his FA-mask concat is cleaner than ours (he does the row concat in F32 rather than touching the shared ggml_concat op), so we adopted it and merged his PR with credit.

For tracking: one of the four, the set_rows F32 fix, already landed on main via #2038, so only three still need to go in: the indexer-norm eps, the F16-mask to F32 casts, and the F32 FA-mask concat.

Merging mgkwill's PR also put the CPU fixes onto the #2045 branch, which now conflicts with main, and the conflict is exactly that already-merged #2038 set_rows region. So the cleanest path depends on how you would like it structured: rebase #2045 onto the head-loop and carry the three remaining fixes, or keep #2045 as the indexer and take the CPU fixes as a small follow-up. Either is easy on our side, and your call there also resolves the conflict.

@mb8565

mb8565 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

On the /32: no, our current update to the PR does not change the buffer. We are getting the CPU-only support fully in order before moving forward with more; the head-loop and its memory behavior are ikawrakow's (#2058), and our fixes do not touch them.

As for why the loop comes out below /32: it only bounds the per-head score, from {n_kv, n_tokens, n_ihead} down to {n_kv, n_tokens} (n_ihead is 32 here, which is where the /32 ceiling comes from). The other indexer buffers are still {n_kv, n_tokens} and the loop does not touch them, so the total drop is smaller, which is why #2058 lands around the 1.75x you saw rather than /32. Our P100 numbers line up with that: at 16K the compute buffer grows from 1697 MiB at ub 512 to 6376 at ub 2048, so it grows with the token count. We have not measured the exact per-term split at your 500K, but the custom op you pointed to in #24231 avoids materializing the per-head score at all, and that is the real fix.

On the quant: that compute buffer is F32/F16 activation tensors sized by n_kv and n_tokens, so it follows context and batch rather than the weight quant. If the RAM you are seeing moves with quant, that is most likely the weight footprint rather than the indexer buffer.

Happy to keep at the GLM DSA work with everyone here.

mb8565 and others added 6 commits June 30, 2026 18:29
…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>
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>
…-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>
… 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>
…eq>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>
…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>
mb8565 and others added 9 commits June 30, 2026 18:29
…off by default)

Implements ikawrakow's direction from discussion ikawrakow#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>
…ns 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>
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>
PR ikawrakow#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>
@mb8565
mb8565 force-pushed the glm-dsa-upstream branch from 93140f9 to 1184ddd Compare July 1, 2026 00:31
mb8565 pushed a commit to mb8565/ik_llama.cpp that referenced this pull request Jul 1, 2026
PR ikawrakow#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>
@mb8565

mb8565 commented Jul 1, 2026

Copy link
Copy Markdown
Contributor Author

Rebased onto main. The conflict was the set_rows F32 fix duplicating #2038, resolved by taking main's version; the three remaining CPU-only fixes and the indexer are unchanged. Confirmed it still builds cleanly.

mb8565 added a commit to mb8565/ik_llama.cpp that referenced this pull request Jul 1, 2026
The GLM-DSA sparse-mask destination `base` was created with ggml_new_tensor_3d,
i.e. an op-NONE leaf. ggml-alloc allocates all graph leaves eagerly before the
node loop runs, so one {1,n_kv,n_tok} base leaf per DSA layer is live at once at
the compute-buffer high-water mark, and the buffer grows as (n DSA layers) x
n_kv x n_ubatch x 4 -> it OOMs at long context (usrlocalben's report on ikawrakow#2045).
Here that is 78 layers: n_layer 79 minus the one MTP/next-n layer, which has no
indexer (guard build_deepseek2.cpp: dsa && GLM_DSA && indexer_attn_q_b).

Build `base` as a FILL node off the pen_b node instead. Nodes are allocated
lazily and returned to the free list after their last consumer, so base is
reused across layers and the per-layer pile disappears. Output is byte-identical:
ggml_fill writes -BIG everywhere and set_rows overwrites every key slot exactly
once (sorted is a per-column permutation over n_kv), so base's initial contents
were always irrelevant; ggml_fill dup_tensors and FILL is not can-inplace, so
the fill output cannot alias pen_b (still read by set_rows). No new op.

Measured CPU-only on this branch (GLM-5.2-IQ2_M, -ngl 0 --dsa -fa 1 -c 32768
-ub 512), controlled A/B differing only in this line: CPU compute buffer
12617 -> 7629 MiB (delta 4989 = 78 x 64 MiB), and 4-chunk 4K PPL byte-identical
(3.2102 == 3.2102). Independent of the indexer-score approach, so it stacks with
the head-loop in ikawrakow#2058; on that branch a GGML_ALLOCATOR_DEBUG dump confirms the
78 per-layer 64 MiB base leaves drop to 0 in the peak live-set (10569 -> 5580).

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

mb8565 commented Jul 1, 2026

Copy link
Copy Markdown
Contributor Author

On the long-context DSA memory: we think we found where it comes from, and a small fix, but we haven't finished testing it, so treat this as preliminary. Sharing it now in case it's useful to test against.

On our side (CPU-only, GGML_ALLOCATOR_DEBUG, GLM-5.2-IQ2_M) the compute-buffer peak at 32K is filled by the sparse-mask destination base in build_deepseek2_dsa_sparse_mask. It's created with ggml_new_tensor_3d, so it's an op-NONE leaf. As we read ggml-alloc, leaves are allocated eagerly before the node loop, so one {1, n_kv, n_tokens} base leaf per DSA layer is live at once at the high-water mark, and the buffer grows as (num DSA layers) x n_kv x n_ubatch x 4. In this model that's 78 such leaves (every layer with an indexer); the allocator peak line shows exactly 78 of them at 64 MiB each. This is a separate object from the KQ materialization you flagged (which #2058 handles): an op-NONE leaf the allocator pins once per layer.

The one-line change we're testing is to build base as a node instead of a leaf, reusing the existing ggml_fill on the pen_b node rather than filling a fresh tensor:

ggml_tensor * base = ggml_fill(ctx0, pen_b, -BIG);   // was: new_tensor_3d + fill

As a node it gets the normal lazy alloc + free-list reuse and stops piling up per layer. It should be output-identical: the fill still writes -BIG everywhere and set_rows overwrites every key slot exactly once (sorted is a per-column permutation over n_kv), so base's initial contents were always irrelevant. As far as we can tell the change is also alloc-safe (ggml_fill dup_tensors, and FILL isn't can-inplace, so the fill output can't land on pen_b's buffer). No new op.

What we've measured so far: on the current #2045 head, a controlled A/B differing only in that line drops the CPU compute buffer 12617 -> 7629 MiB (CPU, -fa 1, 32K, ub 512). Output is identical by construction (the fill/set_rows argument above), and 4K PPL is unchanged at 3.2102. That ~4989 MiB is the per-layer base pile (78 layers x 64 MiB); an allocator-debug dump on the head-loop branch shows those 78 base leaves going to 0 in the peak live-set. It's the same pile regardless of the indexer-score approach, so it also stacks on your head-loop (10569 -> 5580 there). Across 16K-131K the DSA-vs-dense overhead becomes O(1) in layer count, growing only with n_kv. What we haven't checked yet: -fa 0, long-context decode, the GPU backend, and multi-sequence / context-shift numerical identity. That's why we're calling it preliminary and would value another set of eyes.

One note so it doesn't cost you a detour: we initially suspected the per-head indexer scores weren't being reused, but forcing them inplace changed nothing and the dump shows they're reused fine (they're nodes, not leaves), so it really is just the base leaf. We're steering clear of a dedicated fused op given your point about older CUDA archs (our P100 case). We've put it on #2045 so people can test against it.

@usrlocalben

Copy link
Copy Markdown
Contributor

2cdb513 is an improvement in buffer size, but there's probably no need to get into those details since the output is bad:

tell me a story about a laptop and a cat:

<snip>
She bought him a heated cat bed. He ignored it. She put a box next to her desk. He sat in it for two minutes, then returned to the
 laptop.

The truth was, it wasn't about warmth. It was about *her*. The laptop had her attention, and Mochi wanted it for himself.

One rainy Tuesday, Sarah closed the laptop, pulled Mochi into her lap, and just sat there, listening to the rain and his purring.
<snip>

That seems ok. However a 10K tok prompt "Summarize this EULA" although not complete gibberish is incoherent at best.

% cat ~/prompt/nintendo-eula.md | llm -m glm
If any3, 16 and  Section 16(j) - 1
  Section 16(3) j): 1. Section  20. Bell
EULA mentions

Section 16(j) of the EULA and Section 16(j)E; Section 16(d) - j - 1.
B.3; and 16; 021; Section 0)
* E. 13 and 5
* 3. Section </think>1: EULA. 5.13) - 0
*1
<snip>

@mb8565

mb8565 commented Jul 1, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the clear repro. The base-fill (2cdb513) is allocation-only and PPL-identical with and without it in our CPU runs (chunk-for-chunk, on -fa 1 and -fa 0), so it shouldn't be the cause. The incoherence looks like the DSA sparse path once the prompt passes the indexer top_k, which would fit: short prompts stay dense and read fine, while a long one like yours crosses into the pruned path. We're isolating it on CPU now and will follow up.

@ikawrakow ikawrakow mentioned this pull request Jul 1, 2026
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 ikawrakow#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>
@mb8565
mb8565 force-pushed the glm-dsa-upstream branch from 2cdb513 to 9b84916 Compare July 2, 2026 03:17
@mb8565

mb8565 commented Jul 2, 2026

Copy link
Copy Markdown
Contributor Author

Updated this PR. We rebased it onto @ikawrakow's #2063 branch, so it carries his head-loop, ggml_hadamard, ggml_blend, argsort, and FA-mask work, with the IndexShare fix on top.

IndexShare: GLM-5.2's indexer_types marks 21 "full" indexer layers that compute their own top-k and 57 "shared" layers that reuse the previous full layer's top-k. This port (and #2063) were computing an independent top-k on every layer. On an unsloth IQ2_M GGUF, 4K, top_k 2048, CPU, fixing that takes DSA-on perplexity from 3.19 to 2.71 against a dense 2.70. Root cause and the full write-up are on #2063 (issuecomment-4861866823), where @usrlocalben independently identified IndexShare.

One caveat: with -fa 0 (no flash attention) the DSA path produces garbage perplexity. It is not from the IndexShare fix (it reproduces with the shared-layer reuse off), and we root-caused it to an in-place write into the shared KQ_mask on the -fa 0 soft_max path (-fa 1 is unaffected). We have a candidate fix in testing.

Comment thread ggml/src/ggml-cuda/blend.cu Outdated

@ikawrakow ikawrakow left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems to work and performance is quite a bit better than originally.

I'm not worried about the no-FA path not working. If you cannot find a fix, we can add a guard disabling DSA + no-FA.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants