Skip to content

GLM-DSA: fix -fa 0 garbage perplexity on the batch>8 indexer path (private seed, do not corrupt KQ_mask) - #2069

Merged
ikawrakow merged 1 commit into
ikawrakow:mainfrom
mb8565:glm-dsa-fa0-fix
Jul 2, 2026
Merged

ikawrakow merged 1 commit into
ikawrakow:mainfrom
mb8565:glm-dsa-fa0-fix

Conversation

@mb8565

@mb8565 mb8565 commented Jul 2, 2026

Copy link
Copy Markdown
Contributor

Follow-up to the GLM-DSA work (#2045, #2066, #2067). With -fa 0 (no flash attention) and --dsa, the DSA path still produces garbage perplexity on the batch >8 (prefill) indexer path. -fa 1 --dsa is fine. #2067 already fixed the small-batch (<=8, TG/decode) path as a side effect of switching it to a single non-inplace ggml_add. This PR handles the remaining batch >8 path.

Cause: build_deepseek2_dsa_indexer seeds the indexer score accumulator with a view of KQ_mask. On -fa 0, KQ_mask is the raw F32 input tensor. The small-batch path (#2067) accumulates with a non-inplace ggml_add, so it does not touch KQ_mask. The batch >8 path still accumulates the heads with ggml_add_inplace, which writes into the shared KQ_mask buffer and corrupts the causal mask that build_deepseek2_dsa_sparse_mask and the later softmax layers read back. On -fa 1 the mask is F16, so the indexer's cast to F32 already gives a private buffer.

Fix: in the batch >8 branch, take a private copy of the F32 seed before the in-place accumulation, matching what the small-batch path and -fa 1 already do:

} else {
    // -fa 0: the seed still aliases KQ_mask; copy it so the in-place accumulation below does not
    // corrupt the shared causal mask. -fa 1's F16 mask was already cast to a private F32 buffer above.
    if (KQ_mask->type == GGML_TYPE_F32) {
        indexer_score = ggml_cont(ctx0, indexer_score);
    }
    for (int head = 0; head < n_ihead; ++head) {
        ...
    }
}

The small-batch / TG path is left untouched, so it keeps the #2067 performance work, and the copy is gated to -fa 0 (raw F32 mask) so -fa 1 prefill pays nothing. Cost is one n_kv x n_tokens F32 copy per full layer on the batch >8 -fa 0 path only.

Results (unsloth IQ2_M GGUF, top_k 2048, CPU unless noted; on current main dbe2ecb / #2067):

  • -fa 0 --dsa 4K PPL (batch 512, the batch >8 path): thousands (broken) -> 2.7134 with the fix (dense 2.6972, -fa 1 --dsa 2.7111)
  • -fa 0 --dsa 4K PPL on 3x P100 (CUDA, --cpu-moe): 2.6980, sane and in the dense / -fa 1 band (not garbage), so the fix works on CUDA (pre-GLM-DSA: minor optimization #2066 tree, same copy fix)
  • Dense-equivalence oracle (DSA top_k >= n_kv must == dense, chunk-for-chunk): bit-exact on both -fa 1 and -fa 0-with-fix -> the attention machinery is correct, the DSA delta is pure key selection (pre-GLM-DSA: minor optimization #2066 tree, same copy fix)
  • Coherence: -fa 0 clean at 9K (CPU), coherent at 24K (GPU, 21272-token prompt); -fa 1 unaffected throughout (pre-GLM-DSA: minor optimization #2066 tree, same copy fix)
  • Note: -fa 0 --dsa at very long context is memory-heavy (per-layer sparse mask); 24K is the max on 48GB VRAM

Performance: decode is unaffected. The #2067 small-batch fast path is batch-gated, not FA-gated, so -fa 0 decode also takes it and never reaches the patched batch >8 loop. The fix adds one n_kv x n_tokens F32 copy of the causal-mask seed per full layer on the -fa 0 prefill path, small next to the per-head matmul loop it guards, and it leaves the TG compute-buffer path that #2067 optimized untouched. If you would rather fold the copy into the accumulation, making head 0's add non-inplace and keeping the rest in-place is an equivalent, slightly cheaper alternative to the ggml_cont.

The batch >8 indexer path in build_deepseek2_dsa_indexer accumulates the
per-head scores with ggml_add_inplace into an accumulator that is seeded
from a view of KQ_mask. On -fa 0, KQ_mask is the raw F32 input tensor, so
the in-place writes land in the shared KQ_mask buffer and corrupt the causal
mask that build_deepseek2_dsa_sparse_mask and the later softmax layers read
back, which gives garbage perplexity.

-fa 1 is unaffected (its F16 mask is cast to a private F32 buffer), and the
small-batch path added in ikawrakow#2067 is unaffected (it uses a non-inplace add).
Take a private copy of the seed in the batch >8 -fa 0 path (raw F32 mask)
before the accumulation, matching what those two paths already do.

4K -fa 0 --dsa PPL goes from thousands to 2.7134 (dense 2.6972, -fa 1 --dsa
2.7111). -fa 1 and non-DSA builds are byte-identical.

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

mb8565 commented Jul 2, 2026

Copy link
Copy Markdown
Contributor Author

Noticed #2068 is in flight on the same file. This one only touches the batch >8 indexer accumulation, which #2068 leaves alone, so they are independent and it should rebase cleanly on top. Happy to rebase once #2068 lands.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants