GLM-DSA: much better PP long context performance (CUDA)

[ikawrakow]

This PR implements a much more efficient sparse attention computation for GLM-DSA on CUDA with two caveats

• mla = 1 is requuired
• f16 K-cache is required
• flash attention (FA) is required

No new op or command line options are added. Instead, when the CUDA FA is invoked and the computation can be done in the better way (i.e., there is a top_k tensor, the above requirements are met, and the context >= 4*n_top_k), the newly added function ggml_cuda_dsa_attn_ext instead of the standard FA is used.

The original idea was to integrate the top_k selection directly into the FA kernel. But, after getting confused multiple times and not arriving at a working version, I ended up putting together an implementation that iterates over the batch tokens in chunks (with the only reason for chunking being to restrict the size of the intermediate compute buffers to a reasonable maximum - about 220 MiB for GLM-DSA (nearly) independent of batch size).

At a context of 100k tokens I observe a ~2.7X better PP compared to the not yet merged PR #2103.

Below are sweep-bench results for GLM-5.2-Q4_K_M on a 13x3090+Ryzen-3995WX system (GPUs are limited to 200W). u-batch size is 2048, 26 layers are offloaded to the GPUs. The dip in the black curve is strange, possibly a power management issue.

[Skelectric]

Hi, I ran some benchmarks on GLM-5.2_UD-Q4-K-XL:

This is on a server with 2x Epyc 9355 and an RTX Pro 6000.

Thank you!

[usrlocalben]

at the risk of going off-topic, @Skelectric could you share your ik config? My system is very similar (2x 9B14, 1x R6KP) --decode speed is similar for my q4/q4/q6 quant but your prefill is significantly better. I see 315t/s down to 227t/s @ 100K. maybe you have the 600W version?

mine, for this PR:

MEMORY_CONCURRENCY_IN_NUMA_NODES=8
N_BATCH=8192

# this quant is Q6 attn, BF16 index, _exps Q4/Q4/Q6 gate/up/down
COMPUTE_POWER_IN_THREADS_PER_NUMA=6
M=/model/GLM-5.2/usrlocalben/HQ4_K/GLM-5.2-HQ4_K.gguf

N_THREADS=$[ $MEMORY_CONCURRENCY_IN_NUMA_NODES * $COMPUTE_POWER_IN_THREADS_PER_NUMA ]
# N_THREADS=144

GGML_CUDA_NO_PINNED=1 \
exec ./ik_llama.cpp/build/bin/llama-server \
  --host 0.0.0.0 --port 4972 --webui llamacpp \
  --numa distribute \
  -t $N_THREADS \
  --cache-ram 200000 \
  -b $N_BATCH -ub $N_BATCH -amb 1024 \
  -mla 1 -dsa -fidx \
  -ngl 999 -ot exps=CPU \
  --jinja --parallel-tool-calls \
  --chat-template-kwargs '{"enable_thinking": true, "reasoning_effort": "high"}' \
  -ctk f16 \
  -c $[ 2**18 ] \
  --alias GLM-5.2 \
  --metrics \
  -m "$M"

[Skelectric]

@usrlocalben Its a 600W model power limited to 400W.

    cmd: >
      /etc/ik-llama-cpp/llama-server-dev
      --model /mnt/nvme/models/GLM-5.2/UD-Q4_K_XL/GLM-5.2-UD-Q4_K_XL-00001-of-00011.gguf
      --host 0.0.0.0
      --port ''${PORT}
      --jinja
      --chat-template-file /mnt/nvme/models/GLM-5.2/chat_template.jinja
      --dsa
      -mla 1
      -amb 512
      -ctk f16
      -ngl 999
      -ot exps=CPU
      -fidx
      --ctx-size 256000
      --batch-size 8192
      --ubatch-size 8192
      --threads 64
      --webui llamacpp
      --log-file ${logDir}/glm-5.2-dev
    aliases: ["glm-5.2-dev"]```