[BUG] MTP draft acceptance collapses at specific context sizes (KV slot boundary)

[ronnieops]

Summary

MTP (multi-token prediction) speculative decoding experiences dramatic draft acceptance rate collapse at specific context sizes in llama-server. A 256-token context difference can mean the difference between a 2.1x speedup and 1.04x (near-zero acceptance). The pattern repeats at ~2048-aligned boundaries, suggesting an internal KV cache buffer alignment issue in the MTP draft context.

Hardware

• GPU: NVIDIA GeForce RTX 5060 Ti 16GB (SM120 Blackwell, 16311 MiB usable)
• CPU: AMD Ryzen 5 7600

Software

• llama.cpp b9295 (95405ac): MTP works correctly at small contexts
• llama.cpp b9305 (63248fc): Includes fit-margin commit 83eebe9
• Build: CMAKE_BUILD_TYPE=Release, LLAMA_CUDA=ON, CMAKE_CUDA_ARCHITECTURES=120
• KV cache: --cache-type-k q8_0 --cache-type-v q8_0

Model

• Unsloth Qwen3.6-35B-A3B-UD-IQ2_M (GGUF, 11.9 GB, 35B params, 3B active per token)
• MTP weights included (NextN architecture)
• Server flags: -ngl 99 --spec-type draft-mtp --spec-draft-n-max 2

Reproduction

llama-server \
  --model Qwen3.6-35B-A3B-UD-IQ2_M.gguf \
  --ctx-size 17408 \
  --cache-type-k q8_0 --cache-type-v q8_0 \
  -ngl 99 \
  --spec-type draft-mtp --spec-draft-n-max 2 \
  --port 18888 --no-ui

Then send via HTTP API:

POST /v1/completions
{
  "prompt": "<test prompt>",
  "max_tokens": 256,
  "temperature": 0.0,
  "speculative": {"n_max": 2}
}

The collapse is triggered by --ctx-size value, NOT prompt length.

Results (b9305, q8_0 KV cache, MTP N=2)

Context	Baseline TG	MTP N=2 TG	Speedup	AR%	Status
4096	103.5	163.9	1.58x	55%	OK
8192	101.3	159.9	1.58x	55%	OK
10240	66.0	69.3	1.05x	7%	BROKEN
11776	76.6	76.9	1.00x	0.6%	BROKEN
12000	33.5	30.4	0.91x	0%	BROKEN
12032	76.5	79.9	1.04x	6%	BROKEN
12288	74.2	155.5	2.10x	78%	OK
12544	75.3	82.6	1.10x	13%	BROKEN
12800	97.1	123.6	1.27x	32%	Partial
13312	76.5	85.3	1.12x	15%	BROKEN
14336	73.2	75.3	1.03x	4%	BROKEN
17024	28.6	25.2	0.88x	0%	BROKEN
17280	69.2	79.3	1.15x	19%	Degraded
18432	90.5	152.1	1.68x	61%	OK

Same pattern at larger context with q4_0 KV (halved size):

Context	Baseline TG	MTP N=2 TG	Speedup	Status
12288	75.6	160.6	2.12x	OK
22016	25.2	23.7	0.94x	BROKEN
22336	85.9	146.6	1.74x	OK
22528	85.3	106.2	1.25x	Partial
24064	25.1	22.3	0.89x	BROKEN
24576	84.3	145.7	1.75x	OK

Key Observations

1. Narrowest transition: 256 tokens between 1.04x (ctx=12032) and 2.10x (ctx=12288). A single power-of-2 boundary separates regression from peak performance.

2. Pre-dates fit-margin (b9295): Same collapse at ctx=17024 on b9295. Fit-margin commit 83eebe9 only shifted boundary positions by ~288 tokens — didn't fix the root cause.

3. Both q8_0 and q4_0 KV affected: q8_0 collapses at ctx=24576 (0.87x) but q4_0 works (1.75x). Yet both show the same ~2048-aligned structural pattern — NOT a VRAM limitation.

4. AR is KV type independent: 12288 with q4_0 KV: AR 79% (2.12x). 12288 with q8_0 KV: AR 78% (2.10x). Identical acceptance rate proves the AR collapse is in the draft context interaction, not KV quantization.

5. No crash, no OOM: Server stays operational — draft simply stops accepting predictions (silent performance bug).

6. Early EOS at collapse points (b9295): At ctx=17024 with q8_0 KV, server emits early EOS (12 tokens vs 128 requested), corrupting throughput measurements.

Expected Behavior

MTP acceptance rate should be consistent across context sizes (55-78% AR for this model). Collapses of 100x (78% → 0.6%) over 256-token boundaries indicate a structural issue in how the MTP draft context shares the KV cache with the target context.

Hypothesis

The ~256-2048 token granularity suggests an internal buffer alignment issue in slot allocation for the MTP draft context vs target context. When the KV pool boundary aligns unfavorably at specific context sizes, the draft context's slot management interferes, causing near-zero token acceptance.

Suggested Investigation

• Check how LLAMA_CONTEXT_TYPE_MTP context is created relative to slot allocation in server-context.cpp and llama-context.cpp
• Look for power-of-2 or block-aligned boundary conditions in slot management for the draft vs target contexts
• Consider whether the combined target+draft KV cache exceeds slot capacity at specific context sizes

[am17an]

Follow the issue template

[ronnieops]

Update: b9318 Testing Results

Tested same model (Qwen3.6-35B-A3B IQ2_M) on b9318 (6c4cbdc, vanilla, no local patches).

Key Changes in b9318

• PR server: MTP layer kv-cache should respect draft type ctk #23646: MTP draft KV-cache now respects draft ctk/ctv
• PR server: fix checkpoints creation #22929: Checkpoint creation refactor + TYPE_RS rollback (subsumes our PART_BOUNDED workaround)

b9318 Results (q8_0 KV, N=2, gen=512)

Context	Speedup	AR%	Status
4096	1.59x	56%	OK
8192	1.66x	59%	OK
10240	1.13x	17%	BROKEN
12288	2.14x	80%	Sweet spot
16128	0.93x	--	BROKEN
16256	1.12x	16%	BROKEN
16384	1.14x	18%	BROKEN
16640	2.11x	79%	New sweet spot
17024	0.87x	--	BROKEN

Notable Change: 16640 Fixed

Most interesting: 16640 previously 1.82x on b9305, now 2.11x (79% AR) on b9318. 256-token offset from 16384 (2^14) boundary. PR #23646 (draft KV cache respecting ctk/ctv) likely involved.

Extended Context (q4_0 KV)

Context	Speedup	AR%	Status
12288	2.21x	82%	OK
24576	1.75x	64%	OK
24832	1.76x	65%	Max working
25088	0.83x	--	BROKEN
32768	0.86x	--	BROKEN

Observations

1. Collapse pattern persists -- b9318 did not eliminate ~2048-aligned collapse zones
2. 16640 unblocked -- suggests issue partially related to draft context KV cache type handling
3. q4_0 boundary at 24832 -- 256 tokens before 25088 (another 2^N + offset pattern)
4. No patch needed -- PART_BOUNDED workaround dropped; upstream TYPE_RS handling in PR server: fix checkpoints creation #22929 sufficient

KV slot boundary root cause (collapse at specific ctx-size values) remains open. b9318 improves boundary positions but does not eliminate them entirely.