speculative: --spec-draft-window (MTP drafting at every depth)

[professorpalmer]

One of the five pieces of #285, split per bri-prism's request on #317. Two commits on prism (eaecb50c7): the change, and the fixes from the review. Applies without conflicts.

What it does

--spec-draft-window N: the draft (MTP) context keeps only the last N rows of each sequence. Before each batch, the server drops the rows of each sequence in the batch that are older than the window, measured from that sequence's first position in the batch; sequences that are not in the batch keep their rows. The draft context is sized for the window (common_speculative_init), so its cells are reused and a draft pass costs the same at any depth. This retires the --spec-draft-depth-max 24576 cutoff from #221 (the flag stays, default 0). The MTP draft loop now honours the per-draft cap that the server sets for each slot.

The tail draft size (--spec-draft-n-max-tail) is no longer in this PR: it needs the tiered-KV line from #319, so it will come as a follow-up stacked on #319. This branch now builds on its own (CPU-only build of llama-server checked); a CUDA serving check with two slots follows in the thread.

Receipts (RTX 4070, Bonsai 2 27B PTQ1_0 with the MTP head, q8_0 K/V)

• Drafting at every depth instead of stopping at 24k: 32k 54.5 -> 103.6 tok/s, 64k 48.0 -> 90.1, acceptance unchanged. These numbers were taken together with cuda+server: q8_0 K/V at the full 262k window on 12 GB (tiered VMM KV cache), MTP drafting at every depth #285's FA GQA-decode commit (now going into cuda: use the MMA flash attention kernel for GQA above 4 with quantized K/V on Ada #307 with the ggml_cuda_fattn_mma_kv_native_supported guard); the window alone removes the depth cutoff, the speed at depth depends on the attention kernel the decode step uses.
• Draft size is a quality-neutral knob: greedy continuations stay valid greedy decodes; the identity notes on drafting from cuda+server: q8_0 K/V at the full 262k window on 12 GB (tiered VMM KV cache), MTP drafting at every depth #285 apply (draft-on is not always byte-identical to draft-off at the rounding level, with the stock kernels too, because the FA KV split follows the kernel instance and the padded KV length).

Serving recipe and receipts: https://github.com/professorpalmer/bonsai-ada-surgery/blob/main/docs/Q8_FULL_CONTEXT.md

[bri-prism]

I found three issues that need fixing before this is ready: the standalone server does not compile, window eviction crosses sequence boundaries, and sleep/wake loses the normal draft limit. Details and focused reproductions are attached inline.

The server translation unit fails its generated CPU syntax check; the base source passes the same check. Runtime counterexamples use the extracted changed code with actual KV-cell metadata or the actual draft-limit method, under AddressSanitizer and UndefinedBehaviorSanitizer. These are focused reproductions, not an end-to-end MTP serving run, because the current head cannot compile.

[bri-prism]

[P1] Remove the undeclared tiered-KV dependency from the standalone PR

common_params has no n_kv_vram_cells field on this branch. The generated server compiler command fails here with no member named 'n_kv_vram_cells' in 'common_params', including when neither new flag is used. The base source passes the same syntax check.

The field belongs to the separate tiered-KV change, so the window currently cannot work independently as described. Please make this branch compile without that prerequisite, or explicitly stack/include the dependency. The tail threshold should remain zero until a real tiered-KV API is available.

[bri-prism]

[P1] Trim only batched sequences, using a separate cutoff for each

The cutoff is computed from the minimum position of the whole batch, then applied to every slot, including absent slots. With window 100 and only slot 0 batched at position 5000, cutoff 4900 deletes all 500 cached rows of slot 1 at positions 0-499. If both slots are batched at positions 5000 and 500, cutoff 400 instead leaves 1000 old cells in the deep slot where only 100 should remain.

Both cases reproduced using the exact eviction block and real KV-cell metadata/removal code. This can discard another slot's draft prefix or defeat the smaller draft context's bounded-cache assumption; continued growth can exhaust that context. Please compute the cutoff per sequence ID present in the batch and only trim that sequence.

[bri-prism]

[P2] Preserve the requested normal draft limit across sleep/wake

The initial load saves the normal limit, then widens params_base.speculative.draft.n_max for allocation. On wake, handle_sleeping_state(false) calls load_model(params_base), so line 990 recaptures the already-widened value as the normal limit.

With normal limit 2 and tail limit 4, the extracted initialization code plus the actual get_n_draft_max() method returns 2 before sleep and 4 after reload at position 100, below a 32768 tail threshold. Please keep the requested limits separate from allocation capacity and avoid recapturing the normal limit from mutated parameters on resume.

[professorpalmer]

@bri-prism, thank you. All three findings are correct. Fixes are in 93fcb95:

[P1] Standalone build. Correct: n_kv_vram_cells comes from #319. I removed --spec-draft-n-max-tail from this PR, with its per-slot fields and the n_max widening at load. The window does not need the tail. The tail comes back later as a follow-up stacked on #319. Check: a CPU-only build of llama-server on this branch now compiles and links with no errors, and --help lists --spec-draft-window only. Title and description are updated.

[P1] Eviction across sequences. Correct. The cutoff now comes from each sequence's own first position in the batch, and only sequences that are in the batch are trimmed:

std::map<llama_seq_id, llama_pos> seq_pos_min;   // per sequence in batch_view (all seq_id entries of each token)
...
for (const auto & [sid, pos_min] : seq_pos_min) {
    const llama_pos hi = pos_min - params_base.speculative.draft.n_window;
    if (hi > 0) {
        llama_memory_seq_rm(llama_get_memory(ctx_dft), sid, 0, hi);
    }
}

Check with the extracted block (plain arrays in place of the batch and the memory call), your two cases plus two more:

Batch (window 100)	Rows removed
only seq 0 at 5000	seq 0 [0, 4900); seq 1 keeps all its rows
seq 0 at 5000, seq 1 at 500	seq 0 [0, 4900), seq 1 [0, 400)
seq 0 at 5000, 5001, 5002 (a verify batch)	seq 0 [0, 4900)
seq 1 at 50	nothing

[P2] Sleep/wake. Correct. The cause was the widening of params_base.speculative.draft.n_max at load, which load_model(params_base) then read back as the normal limit on wake. That code is gone with the tail. When the tail returns on #319, it will keep the requested normal and tail limits in their own members, set once from the user's parameters, and size the drafter from them without writing params_base. I will add your sleep/wake case as a check there.

Still to come here: a CUDA serving run on the RTX 4070 with two slots and MTP. One slot will be deep and one shallow, with the window set below the deep slot's position. I will report the acceptance of the shallow slot, and greedy output against one-slot runs. The GPU is busy for about two hours, so that comes later today.

[professorpalmer]

CUDA serving check of 93fcb95, as promised. RTX 4070, Bonsai 2 27B PTQ1_0 with the MTP head, this branch built for sm_89, --spec-type draft-mtp --spec-draft-n-max 2 --spec-draft-window 1024, q8_0 K/V, greedy, thinking off (bench/pr320_two_slot_check.py in our repo):

1. Two slots (-np 2). Slot 1 answers a short prompt (B1). Then slot 0 runs alone with an 8.9k-token prompt, far past the window (A), so every one of its batches trims its own draft rows. Then slot 1 continues its conversation (B2).
2. One slot: B1 and B2 only, for comparison.

	Two slots	One slot
B1: prompt / answer tokens, acceptance, text hash	185 / 200, 0.452, b968dbb6f70d	185 / 200, 0.452, b968dbb6f70d
A (slot 0, alone, deep)	8,869 / 115, 0.585	-
B2: prompt tokens evaluated / answer, acceptance, text hash	26 / 200, 0.495, c3ae1c5a9f45	26 / 200, 0.495, c3ae1c5a9f45

Slot 1's draft rows survive slot 0's deep batches: B2 has the same text, the same draft acceptance and the same prompt reuse as without the other slot. With the old cutoff (the minimum position of the whole batch, applied to every slot), A's batches would have removed all of slot 1's rows.