Cap n_outputs_max on MTP draft contexts

[Pento95]

Summary

When using an MTP draft head, the draft context reserves a sampling/compute buffer scaling with n_batch * n_vocab — on Gemma 4 (vocab 262144) ~0.5–2 GB depending on batch size. n_outputs_max is simply never set on the draft context, while the built-in MTP path does set it.

Root cause

mtp drafting added (40084bb) introduced two MTP paths:

• built-in MTP context (gpttype_adapter.cpp:884) correctly sets mtp_ctx_params.n_outputs_max = 1;
• separate draft-head path (speculative_decoding_setup()) builds draft_ctx_params from llama_context_default_params() and copies many fields, but never sets n_outputs_max.

It therefore stays at the default (0), which llama-context.cpp resolves to n_batch:

// llama-context.cpp
cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model)
                        ? cparams.n_batch : params.n_outputs_max;

The draft's backend sampling buffers are sized by n_outputs_max:

// llama-context.cpp
backend_float_count = 2 * n_vocab * n_outputs_max;   // logits + probs

so the draft context reserves for n_batch output positions even though it only ever produces a handful.

draft-mtp generates tokens autoregressively — one output per sequence per llama_decode, looped up to --draftamount (see common_speculative_impl_draft_mtp::draft() in common/speculative.cpp). The "multi-token" comes from the loop / chained heads, not from a wide forward, so the draft never needs more than n_seq outputs — independent of --draftamount.

Fix

Two spots, same omission. Mirror the built-in path and cap outputs at the sequence count (n_seq_max, not a hard-coded 1, to stay correct for n_seq > 1):

1. speculative_decoding_setup() — the real draft context. This is what reserves the VRAM at runtime.
2. estimate_draft_autofit_tax_mb() — the autofit pre-flight estimate (added in "another fix for drafting"). It builds an equivalent draft_ctx_params to predict the drafter's VRAM cost for layer offloading; without the cap it over-estimates the MTP draft tax by the same ~0.5–2 GB, so autofit would offload fewer main-model layers than it could. Same one-liner keeps the estimate consistent with the real context.

draft_ctx_params.n_outputs_max = base_ctx_params.n_seq_max;

Measurements (Gemma 4 31B, vocab 262144, concedo_experimental)

Draft-context graph compute buffer (sched_reserve), and n_outputs_max of the draft context, before → after (same branch, single variable):

Config	n_outputs_max	draft compute buffer
batch 1024, q8 KV	1024 → 1	1049.00 → 498.04 MiB (−551 MiB)
batch 512, f16 KV	512 → 1	524.50 → 118.88 MiB (−406 MiB)

Total VRAM drops by more than the compute buffer alone (the separate logits/probs/candidates sampling buffers also scale with n_outputs_max): ~0.7 GB at batch 512 / ~2 GB at batch 1024 on this model. Batch 1024 + f16 KV, which previously OOM'd at load and silently fell back to "Speculative Decoding will not be used!", now loads.

Throughput and draft acceptance are unchanged (EvalScope perf, parallel 1: same tok/s and MTP acceptance before and after).

Effect scales with vocabulary size, so large-vocab models (Gemma 262k) are near worst case; the built-in MTP path (and integrated-MTP models such as Qwen via that path) were already unaffected.

Related upstream: ggml-org#24340.

AI disclosure