[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.

[bri-prism]

This needs a small rebase: prism added a LAB_SPEC_DRAFT_N_MAX knob to the same MTP stop condition, so keep both checks with an OR. With that it builds and behaves well on an H200. With a 9B PQ2_0 plus MTP head at about 29k tokens of context, a 4096 window keeps output token-identical to plain decode, acceptance is unchanged, and the draft context drops from 40960 to 8448 cells. Speed is level there (about 362 t/s with or without the window), which I expect, since the gain should show up deeper and on smaller cards. The dp.n_max check is worth having regardless, because the MTP drafter is the only one that ignores the server's per-draft cap today. Two requests: add a comment or assert tying the depth_max-based sizing to the server not feeding the draft context past depth_max, and add a small test that runs with a short window (for example 64) so the trim path is covered.

[professorpalmer]

Thank you for the H200 run. Rebased on prism c96baf4ab, head 6f5709a96:

• LAB_SPEC_DRAFT_N_MAX: the MTP stop condition keeps both checks with an OR: min(n_max, LAB_SPEC_DRAFT_N_MAX) and the caller's per-draft cap dp.n_max.
• depth_max sizing: writing the assert showed that the claim was only true for one sequence. The server feeds a view to the draft context when any of its tokens is at or below depth-max. With several slots one view can hold a shallow and a deep sequence, so the deep one's rows also reach the draft context, which was sized for depth-max. Now the depth-max sizing applies only with one sequence (n_seq_max <= 1). The comment says why, and the server asserts that a fed view ends below depth_max + n_batch. The window sizing is per sequence, so it stays for any slot count.
• Short-window test: test_draft_window_keeps_output in tools/server/tests/unit/test_speculative.py (with a spec_draft_window field in utils.py). It runs the draft-simple pair with --spec-draft-window 64 and a prompt longer than 64 tokens, so every batch trims, and it checks that drafting happens and that the greedy tokens equal a run without the window. I could not run it here, because the harness downloads the test models.

I ran the same check with MTP on an RTX 4070 (Bonsai 2 27B with its MTP head, --spec-draft-n-max 2, q8_0 K/V, a 1,980-token prompt, 128 greedy tokens):

Run	Draft tokens / accepted	Draft context	Tokens
no window	128 / 62	8192	reference
--spec-draft-window 64	148 / 52	4608	identical
--spec-draft-depth-max 1024, one slot	0 (the prompt is past the cutoff)	5376	(no drafting, so not compared)

0 CUDA errors, and the new assert did not fire. A CPU-only build of this branch alone also compiles and links llama-server.