Eval bug: First decode after prefill runs 2-45x slower: HIP CUDA-graph capture warmup (Qwen3.8-Flash-Next / qwen4exp QSA, gfx1201)

[Draco-Lunaris]

Name and Version

Summary

On the turboquant fork (PR #340 QSA code), the first decode after any significant prefill within the same task runs dramatically slower than subsequent decodes — e.g. 5 t/s vs 36 t/s with MTP at 126K context. The root cause is HIP CUDA-graph capture warmup: the first graph compute with a new batch shape pays capture cost, the second pays validation, and only the third+ reuses the captured graph. A 1-token warmup decode at server init fixes the common case (small prefills); large-prefill first decodes remain slow because the graph cache key includes n_kv-dependent QSA tensor shapes.

Environment

• Model: Qwen3.8-Flash-Next 177B UD-IQ4_XS (93.7GB, 3 parts) + dzannotti MTP draft (Q4_K_M)
• Arch: qwen4exp (QSA sparse attention, 12 indexer layers, compress_ratio 4, top_k 2048 blocks)
• GPUs: 2x AMD Radeon R9700 (gfx1201), ROCm 7.14, --tensor-split 10,8
• Server flags: --spec-type draft-mtp --spec-draft-n-max 2 --spec-draft-ngl 1 --n-cpu-moe 10 --ctx-size 262144 --cache-type-k q8_0 --cache-type-v turbo4 --cache-ram 12288 --moe-cache off --load-mode mmap --tensor-read-lazy on --parallel 1 --jinja
• Fork: TheTom/llama-cpp-turboquant (with PR qwen4exp: upstream catch-up, sparse-fa, and NextN/MTP support #340 QSA fixes)

Observed behavior

Scenario	Decode speed
1K prefill + decode, same task (before fix)	~12 t/s first decode
1K prefill + decode, same task (with warmup patch)	40 t/s
126K prefill + decode, same task	5-6.4 t/s first decode
2nd task on warm slot @126K	35-42 t/s
Without MTP: first decode vs second	12.15 t/s (82ms/tok) vs 24.30 t/s (41ms/tok)

The effect hits /v1/chat/completions and /v1/completions identically. MTP amplifies it (5 vs 36 t/s) because draft acceptance drops on the first decode (0.52 vs 0.76).

Root cause (instrumented)

Timing instrumentation around llama-context.cpp graph build/compute (ggml_backend_sched_graph_compute_async) shows, for a fresh 126K prefill followed by decodes on the same slot:

prefill graph computes:  5096ms → 1420/1225/1165ms → 124/93/89ms
first decode compute:     339ms   (n_tokens=2)
second decode compute:    172ms
third+ decode compute:    7.5ms   (45x faster than first)

Graph build is always <1ms — it's the graph compute that's slow on first calls with a new batch shape. This is the classic CUDA-graph capture pattern: first call captures, second validates, third+ reuses.

Two contributing bugs in the fork:

1. graphs reused = 0 counter is broken for QSA models. llm_graph_result::can_reuse (qwen4exp.cpp:739) compares n_kv-dependent tensor shapes (cell_blk->ne[0] == n_kv etc.) that change every token, so it always returns false. The CPU-side graph is rebuilt every token — normally harmless, but it means nothing hides the GPU-side capture cost.

2. The GPU-level CUDA graph cache (ggml-cuda.cu:2917, ggml_cuda_graph_get_key) hashes node count + first/last node shapes, and those shapes include n_kv-dependent QSA tensors. A graph captured at n_kv=0 (warmup) therefore cannot be reused at n_kv=126K — which is why the warmup fix below works for small prefills but not large ones.

CUDA graphs are NOT disabled by MoE here: for TQ weights needs_sync=(ne[2]>MMVQ_MAX_BATCH_SIZE || !contiguous) and for IQ4_XS needs_sync=(!quantized || ne[2]>mmvq_mmid_max) — both false for decode (ne[2]=1).

The fix: 1-token warmup decode at server init

In server_context::init() (tools/server/server-context.cpp), before the chat-template block:

// WARMUP DECODE: do a 1-token decode to capture the CUDA graph for
// the decode batch shape (n_tokens=1). Without this, the first decode
// after a large prefill runs ~45x slower (339ms vs 7.5ms per compute)
// because the CUDA graph for the decode shape hasn't been captured yet.
{
    SRV_INF("%s", "running warmup decode to capture CUDA graph\n");
    llama_batch warmup_batch = llama_batch_init(1, 0, 1);
    llama_token warmup_tok = 0;
    common_batch_add(warmup_batch, warmup_tok, 0, {0}, false);
    const int warmup_ret = llama_decode(ctx_tgt, warmup_batch);
    if (warmup_ret != 0) {
        SRV_WRN("%s: warmup decode failed (ret=%d), first decode may be slow\n", __func__, warmup_ret);
    } else {
        SRV_INF("%s", "warmup decode done, CUDA graph captured\n");
    }
    llama_batch_free(warmup_batch);
    llama_memory_seq_rm(llama_get_memory(ctx_tgt), 0, 0, 1);
}

Result: small-prefill first decodes go from ~12 t/s to 40 t/s (3.3x). This is a clean, general fix — it captures the (n_tokens=1, n_kv=0) decode graph shape at startup so the first user decode doesn't pay capture cost.

Note: a server-side priming decode inserted after each prefill (DONE_PROMPT→GENERATING transition) was also tried and crashed — the QSA model's set_input needs proper batch config (seq IDs, indexer cache setup) that a raw llama_batch_get_one doesn't provide. The init-time warmup above works because it goes through the full batch path.

Remaining issue (needs deeper fix)

First decode after a large prefill (e.g. 126K, conversation restore) is still ~6.4 t/s. The graph cache key includes n_kv-dependent QSA tensor shapes, so the n_kv=0 warmup graph can't be reused at n_kv=126K. Possible directions:

• Pad/round n_kv in the graph shapes so the key becomes n_kv-independent (upstream fork work)
• Capture warmup graphs at several n_kv checkpoints (e.g. 0, 16K, 64K, 128K, 256K)
• Make can_reuse (qwen4exp.cpp:739) n_kv-tolerant so CPU-side reuse also works

Theories tested and disproven (so nobody re-chases them)

1. Chat-template markers (<|im_start|>/<|im_end|> causing QSA top-k ties): raw completions with a fresh 126K prefill show the same first-decode slowdown; markers at small context are fast (32-38 t/s). Not the cause.
2. QSA top-k sort bottleneck: enabled hipCUB (upstream PR CUDA: enable the CUB path on HIP via hipCUB ggml-org/llama.cpp#26592, radix argsort path compiled and verified in binary) — no improvement. The sort is O(n) radix and not the cost.
3. QSA scattered-gather from top-k ties: applied upstream PR qwen4exp : gather-based sparse attention for QSA decode ggml-org/llama.cpp#28213 (gather-based sparse attention) on top of PR qwen4exp: upstream catch-up, sparse-fa, and NextN/MTP support #340 — no improvement, and it crashed (ggml_abort) during prefill; incompatible with the fork's QSA code. Reverted.

Also verified NOT the cause: hipMalloc (0.5ms/call), GPU clocks (2.3-3.1 GHz, 100-390W during slow decode), CPU usage (~130-160%, same fast/slow), KV-cache size differences (n_kv identical between fast/slow cases).

Benchmark data (with warmup patch, MTP on)

Test	Result
1K prefill + 200 decode	23-40 t/s (wall/eval varies with slot warmth)
126K prefill + 200 decode, same task	6.37 t/s eval, acceptance 0.79
2nd request @126K warm slot	28-35 t/s wall, 31-39 t/s eval
3 consecutive warm requests	29.5 / 28.8 / 36.1 t/s
Chat endpoint warm	25 t/s (acceptance 0.45-0.49 — template lowers acceptance)
Prefill @126K	240 t/s (4.15 ms/tok)

Operating systems

Linux

GGML backends

CUDA

Hardware

AMD Ryzen 9900x 64gb DDR5, Dual r9700 pro GPU's 64gb total vram.

Models

Unsolth Qwen3.8-Flash-Next 177B UD-IQ4_XS

Problem description & steps to reproduce

Reproduction

1. Start server with the flags above on gfx1201 (any QSA model should show it; magnitude scales with context)
2. Send a fresh large prefill (e.g. 126K tokens) with max_tokens: 200, non-stream
3. Observe eval time in server logs: first task ~200ms/tok; resend the same request (warm slot) → ~25-28ms/tok

Instrumentation patch for debugging (thresholded at 50ms, production-safe): wrap model.build_graph(gparams) and ggml_backend_sched_graph_compute_async(sched.get(), gf) in llama-context.cpp with ggml_time_us() timers and print n_tokens.

First Bad Commit

No response

Relevant log output

Logs