[sanmai]

Overview

Just like in qwen4exp (#30965), attention layers are QSA so we should not need to allocate kq_mask.

Again, this saves 2 GB GPU/host memory at ub 4096 and 4 GB at ub 8192 that we could put towards the expert cache or a larger context.

Additional information

The same trick from #30965 but for GLM 5 and DSA arches.

• With 512k context and ub 8192 it was impossible to use GLM-5.3-Flash UD-IQ1_M on 5090
• Now I'm looking at 651 t/s prefill, 12.6 t/s decode
• For DSA we are looking for similar VRAM savings.

As I also included a similar change to DSA as per #30815 (comment) I only tested it only with a slice of a full model; logits were idential.

Benchmarks

Compute buffers, --moe-cache-mib 4000:

ubatch	build	CUDA0 MiB	CUDA_Host MiB
4096	base	6938	2819
4096	PR	4896	771
8192	base	13874	5636
8192	PR	9784	1540

All benchmarks below are on 5090@400W with 128GB DDR4.

Lots of benchmarks

32GB, 262k context

Performance benchmarks:

ubatch	build	--moe-cache-mib	prefill t/s	decode t/s
4096	base	15600	443	10.3
4096	PR	17650	475	11.4
8192	base	7950	642	7.6
8192	PR	12050	673	9.4

These are with q5_1 KV cache and a full context and --no-mmproj -cmoe --load-mode mmap.

512k context

build	ubatch	--moe-cache-mib	prefill t/s	decode t/s	VRAM peak MiB
base	4096	8100	433	7.9	31892
PR	4096	12200	443	9.5	31906
PR	8192	0	651	12.6	28952
PR	8192	2800	TBD	TBD	TBD

5090 limited to 16 GB

On a simulated 16 GB card ( torch.empty(16 << 30, ...)) with 65k context I'm looking at:

ubatch	build	--moe-cache-mib	prefill t/s	decode t/s	MoE cache hit rate
4096	base	3900	414	5.5	16%
4096	PR	4400	410	6.4	29%
2048	base	5400	269	6.6	30%
2048	PR	5650	271	6.9	32%

Without the MoE cache:

build	context	prefill t/s	decode t/s
base	64k	666	12.8
PR	96k	649	12.3

+32k context is not a lot, but it could be a deal breaker for some.

24 GB limited

build	ubatch	--moe-cache-mib	prefill t/s	decode t/s	VRAM peak MiB
base	4096	8100	433	7.9	31892
PR	4096	12200	443	9.5	31906
PR	8192	0	651	12.6	28952

Trying it out

#!/bin/sh

hf download unsloth/GLM-5.3-Flash-GGUF --include "*UD-IQ1_M*"

exec llama-server -hf unsloth/GLM-5.3-Flash-GGUF:UD-IQ1_M \
    --no-mmproj -fa on \
    -c 786432 -np 2 \
    -ctk q5_1 -ctv q5_1 \
    -b 8192 -ub 8192 \
    -cmoe --load-mode mmap "$@"

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, paired with Claude

[sanmai]

My machine is too weak to run the real thing so I tested it with a slice of DeepSeek-V3.2

[ServeurpersoCom]

My machine is too weak to run the real thing so I tested it with a slice of DeepSeek-V3.2

I ran it on the real GLM-5.3-Flash UD-IQ1_S on a rootless container with an RTX PRO 6000, full offload, FA on, q8_0 cache, against the merge base.

Memory: both the compute buffer and the pinned host buffer shrink by the size of the mask, so about 512 MiB each at 128k context and 2 GiB each at 256k. 128k with -ub 8192 now runs, master runs out of memory there.

Output: this model isn't deterministic run to run on CUDA, even on master, so I applied your #31084 to both the PR and master. With it, both give the exact same perplexity and KLD per chunk over 7 runs. The non-fused path and the toy DSA models (deepseek32, glm-dsa, hy_v4) match master too.

Speed: prefill is about 2% faster at -ub 8192, decode doesn't move. So it's a bit faster and saves a good amount of memory.

LGTM! And #31084 makes this model deterministic here, so it's worth pushing on its own.

[ggerganov]

Here we are indirectly using the kq_mask shape. So even though the kq_mask remains unallocated, the kpool logic is still "coupled" to the shape of the KQ mask. We should do follow-up refactors to fix that because it is a bit difficult to follow and also it causes problems such as: #29510 (comment)

[ggerganov]

The refactoring that I mentioned is proposed by the @am17an in 58650ed for qwen4exp - I think we can extract it in a dedicated PR and apply it to GLM5 after merging the current PR here.

[am17an]

#31243

[ggerganov]

@sanmai You should be able to merge this when ready.

[sanmai]

@ggerganov Appreciate that! Going to give it another look before merging.