GLM-DSA: improve TG performance even more - #2068
Conversation
|
OK, with I can go to a context of 55k tokens, but then get OOM in device 2. But here the sweep bench result in this case
|
Nice, the performance drop for TG looks really flat. Any chance you have the patience to run the sweep bench comparison with pure CPU to see how flat it is and where the performance cutoff is without any GPUs. |
|
Running CPU only is sluggish - <4 t/s TG and <30 t/s PP (for short context), so getting up to 64k tokens will take quite some time. For CPU-only I'm getting somewhat more usable performance with MTP - about 5.5 t/s for a query with not particularly high acceptance rate. But we don't have something like I'm also interested to see what happens with full GPU offload. In these benchmarks TG performance is strongly dominated by computing the MoE FFN for 37 layers on the CPU, so DSA performance improvement for long context is limited by that. But I'm curious to know how this implementation compares to what z.AI have reported (i.e., do we get 6X better TG compared to MLA at very long context), so if someone has the hardware to run with full GPU offload, please let us know. |
I was going to suggest something but then realized that because of the fact that the indexer weights exist for every block means that unless someone splits their GGUF in a way that intentionally makes it possible to just download the indexer weights you can't download them separately (and as far as I can find no one has done that yet), and bolt them onto an existing GGUF. |
|
OK, here the sweep-bench CPU-only on a Ryzen-3995WX using GLM-5.2-Q4_K_M from Unsloth: MLA
DSA
So, nearly 1.5X for TG ar 64k tokens. As stated earlier, PP needs more work. Btw, existing GGUFs do contain the indexer tensors, at least those from Unsloth do. That's where I took the GLM-5.2 model that I'm using for testing. |
9f2fe80 to
e942a20
Compare
Thank you! The curve did flatten significantly, and the crossover point for performance stays around 8K. This plus MTP (assuming my workload gets good enough acceptance rates) will make V3.2 look a lot more attractive than the older versions of Deepseek (where I started at 3.5 t/s at 0 context and dropped fast). Performance with DSA on V3.2 might not be as nice given no IndexShare/IndexCache on that model.
Yes for GLM-5/5.1/5.2 that seems to be the case but not for Speciale which I thought the box you use had a pure GPU quant for. |
|
I haven't tried any of this with DeepSeek Speciale. Has anyone else tried? |
Just in case anyone wants links of premade GGUFs that have the indexer for DSA, here is Speciale, V3.2, and V3.2-Exp. They come in Q8_0 and Q4_K_M, but no imatrix was used for either quant. As reference, "DeepSeek-V3.2-Speciale variant is designed exclusively for deep reasoning tasks and does not support the tool-calling functionality", V3.2 is the "normal" variant, V3.2-Exp is an earlier checkpoint with weaker benchmarks than V3.2. |
|
@ikawrakow I've been testing this PR w/MTP, no crash - what conditions cause crash? |
Running CPU-only. Possibly the crash is not here but in the next PR. As I wasn't sure, I decided to push the change to the base of the currently opened performance improvement PR. |
|
Btw, I saw you discussing elsewhere that even with DSA the model becomes incoherent after 100k tokens. On the current branches the matrix multiplication between the indexer cache and the query is done using From past experience, some models may also need adjustment of the FA offset for sufficiently long context. You could try adding |
My plan for today was to work on a clean, minimal reproduction, prompt, etc. for an Issue but on a re-test earlier my 100K test gave answer-correct. 🤷 I'll keep these two ideas in mind as I continue. |
|
edit: disregard. |
| } | ||
|
|
||
| last_sparse_mask = last_sparse_mask_fa = nullptr; | ||
| last_sparse_mask = last_sparse_mask_fa = dsa_last_full_sorted = dsa_tg_fast_mask = nullptr; |
I tried V3.2 (Q8_0 which is the same shape as Speciale) DSA:
MLA:
I don't know if anyone else has tried but if anyone does try with the quant above (or something similar to it) the patch below can be used to make them load. git diffdiff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index bf07387f..9d94b7b2 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -81,7 +81,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_SEED_OSS, "seed_oss" },
{ LLM_ARCH_STEP35, "step35" },
{ LLM_ARCH_LAGUNA, "laguna" },
- { LLM_ARCH_GLM_DSA, "glm-dsa" },
+ { LLM_ARCH_GLM_DSA, "deepseek32" },
{ LLM_ARCH_MISTRAL4, "mistral4" },
{ LLM_ARCH_GEMMA4, "gemma4" },
{ LLM_ARCH_GEMMA4_MTP, "gemma4_mtp" },
diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp
index 3ece0f96..9989fa9d 100644
--- a/src/llama-load-tensors.cpp
+++ b/src/llama-load-tensors.cpp
@@ -2942,8 +2942,8 @@ bool create_tensors_helper::create_glm_dsa_tensors(const LLM_TN & tn) {
}
// Incompatible mainline model. Let's see if we can still load it
- layer.wk_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, 0);
- layer.wv_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, 0);
+ layer.wk_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, flags);
+ layer.wv_b = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, flags);
layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, flags);
// DSA indexer
I'm actually working on refining my own quants of this while testing to see if I like this model. Once I get something I like I plan to put it up on huggingface. It will be made with ik_llama.cpp in mind so I'll use the architecture name from here, and include the MTP tensors which should eliminate the need for the above patch. My current mix:
I know it can be faster given I did extensive testing on ~4.5BPW Deepseek models with this new machine and certain quants got 5.35 TG at N_KV 0 with MLA ( and hopefully even faster once I add MTP to the model if it is beneficial with my workloads and system). I'm going to pare this down closer to IQ4_K. Also I did find mirror to be the best performing in my tests holding it's advantage even at high context values. |
3 participants
This PR follows in the footsteps of #2066 and #2067 and further improves GLM TG performance when using DSA. TG performance is now better than no DSA for context length greater than 8k tokens or so.
For TG the trick to take advantage of sparse attention is relatively simple: we just need to take out from the MLA cache the
n_top_kkeys selected by the indexer, and then do FA with just those. Hence, past 2k tokens (n_top_k = 2048for GLM-5), FA is constant time per generated token. TG performance still slightly degrades with context length as the cost of the indexer increases linearly with context, but overall the performance decrease is very slow (see table below).Below are some sweep-bench results with this PR compared to no-DSA. GLM-5.2-Q4_K_M, 13x3090 (limited to 200W each), Ryzen-3995WX CPU. Command line is
which results in routed experts in 37 layers being left in RAM.
No DSA
DSA (this PR)
The DSA run aborts with OOM on device 12 at a context of 43008 tokens, so I guess I need a larger margin there as well (so one less MoE layer offloaded). But we see that at a context of 40k tokens DSA TG is ~15% better than no DSA, with the gap likely increasing even further at longer context. PP is of course still worse by 10-15%, but improving this will be a separate effort.