Commit 3fb47677e for llama.cpp
commit 3fb47677eaf0ebb569a1f94c756dcf58b8584cc0
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Wed Oct 14 13:58:05 2026 +0300
common : add --spec-draft-n-ubatch (#30459)
* common : add --spec-draft-n-ubatch
Allow configuring the physical batch size of the draft context. The
drafters feed it at most the drafted tokens per sequence and chunk larger
batches by its own n_ubatch, so it does not need the ubatch of the target.
A smaller value shrinks the draft compute buffer and leaves more VRAM for
the target KV cache (the draft is fitted together with the target).
The default 0 keeps the current behavior of following --n-ubatch.
Assisted-by: pi:llama.cpp/Qwen3.8-Flash-Next
* cont : remove comment
[no ci]
diff --git a/common/arg.cpp b/common/arg.cpp
index 6bc4a84e2..29f6df32f 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -4173,6 +4173,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
params.speculative.draft.n_min = value;
}
).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN"));
+ add_opt(common_arg(
+ {"--spec-draft-n-ubatch"}, "N",
+ string_format("physical maximum batch size of the draft context, use a smaller value to spare VRAM (default: %d - same as --n-ubatch)", params.speculative.draft.n_ubatch),
+ [](common_params & params, int value) {
+ if (value < 0) {
+ throw std::invalid_argument("invalid value");
+ }
+ params.speculative.draft.n_ubatch = value;
+ }
+ ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_UBATCH"));
add_opt(common_arg(
{"--spec-synth-len"}, "L",
"target mean synthetic acceptance length, including the target token (benchmarking only)",
diff --git a/common/common.h b/common/common.h
index b6683fe67..2048d6a79 100644
--- a/common/common.h
+++ b/common/common.h
@@ -343,6 +343,8 @@ struct common_params_speculative_draft {
int32_t n_gpu_layers = -1; // number of layers to store in VRAM for the draft model (-1 - use default)
+ int32_t n_ubatch = 0; // physical batch size of the draft context (0 - use the target's n_ubatch)
+
ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V
diff --git a/common/speculative.cpp b/common/speculative.cpp
index c5043a4a8..50647dbb2 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -2565,6 +2565,10 @@ common_params common_base_params_to_speculative(const common_params & params) {
result.n_outputs_max = params.n_parallel;
result.n_outputs_max_per_seq = 1;
+ if (params_spec.n_ubatch > 0) {
+ result.n_ubatch = params_spec.n_ubatch;
+ }
+
// the MoE cache is only used by the target context
result.moe_cache_size = 0;
diff --git a/docs/speculative.md b/docs/speculative.md
index ffb1e34c7..a32842a25 100644
--- a/docs/speculative.md
+++ b/docs/speculative.md
@@ -247,6 +247,9 @@ Use exactly one of these options:
--spec-draft-n-min N
minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN)
+--spec-draft-n-ubatch N
+ physical maximum batch size of the draft context, use a smaller value to spare VRAM (default: 0 - same as --n-ubatch)
+ (env: LLAMA_ARG_SPEC_DRAFT_N_UBATCH)
--spec-draft-p-split, --draft-p-split P
speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT)
diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp
index e0907631a..543da2aed 100644
--- a/tests/test-arg-parser.cpp
+++ b/tests/test-arg-parser.cpp
@@ -95,10 +95,20 @@ static void test(void) {
common_params base;
base.n_parallel = 4;
base.n_outputs_max_per_seq = 8;
+ base.n_ubatch = 512;
- const auto draft = common_base_params_to_speculative(base);
- assert(draft.n_outputs_max == 4);
- assert(draft.n_outputs_max_per_seq == 1);
+ {
+ const auto draft = common_base_params_to_speculative(base);
+ assert(draft.n_outputs_max == 4);
+ assert(draft.n_outputs_max_per_seq == 1);
+ assert(draft.n_ubatch == 512); // inherits the target by default
+ }
+
+ base.speculative.draft.n_ubatch = 8;
+ {
+ const auto draft = common_base_params_to_speculative(base);
+ assert(draft.n_ubatch == 8);
+ }
}
printf("test-arg-parser: make sure there is no duplicated arguments in any examples\n\n");
diff --git a/tools/cli/README.md b/tools/cli/README.md
index 284b8f308..163774119 100644
--- a/tools/cli/README.md
+++ b/tools/cli/README.md
@@ -200,6 +200,7 @@
| `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) |
| `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) |
| `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) |
+| `--spec-draft-n-ubatch N` | physical maximum batch size of the draft context, use a smaller value to spare VRAM (default: 0 - same as --n-ubatch)<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_UBATCH) |
| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)<br/>(env: LLAMA_ARG_SPEC_SYNTH_LEN) |
| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)<br/>(env: LLAMA_ARG_SPEC_SYNTH_RATES) |
| `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)<br/>(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) |
diff --git a/tools/server/README.md b/tools/server/README.md
index 48c40eda7..37c41f1bb 100644
--- a/tools/server/README.md
+++ b/tools/server/README.md
@@ -261,6 +261,7 @@ For the full list of features, please refer to [server's changelog](https://gith
| `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) |
| `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) |
| `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) |
+| `--spec-draft-n-ubatch N` | physical maximum batch size of the draft context, use a smaller value to spare VRAM (default: 0 - same as --n-ubatch)<br/>(env: LLAMA_ARG_SPEC_DRAFT_N_UBATCH) |
| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)<br/>(env: LLAMA_ARG_SPEC_SYNTH_LEN) |
| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)<br/>(env: LLAMA_ARG_SPEC_SYNTH_RATES) |
| `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)<br/>(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) |
diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py
index 22b523954..7a0354302 100644
--- a/tools/server/tests/unit/test_speculative.py
+++ b/tools/server/tests/unit/test_speculative.py
@@ -152,6 +152,32 @@ def test_synth_ignores_target_tokens():
assert res.body["stop_type"] == "limit"
+def test_draft_n_ubatch_smaller_than_target():
+ global server
+ request = {
+ "prompt": "I believe the meaning of life is",
+ "temperature": 0.0,
+ "top_k": 1,
+ "n_predict": 16,
+ }
+
+ # draft with the default ubatch size of the draft context
+ server.start()
+ res_default = server.make_request("POST", "/completion", data=request)
+ server.stop()
+ assert res_default.status_code == 200
+ assert res_default.body["timings"]["draft_n"] > 0
+
+ # a smaller draft ubatch splits the draft batches into more ubatches, but must not change the result
+ server.spec_draft_n_ubatch = 8
+ server.start()
+ res_small = server.make_request("POST", "/completion", data=request)
+ server.stop()
+ assert res_small.status_code == 200
+ assert res_small.body["timings"]["draft_n"] > 0
+ assert res_small.body["tokens"] == res_default.body["tokens"]
+
+
def test_slot_ctx_not_exceeded():
global server
server.n_ctx = 256
diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py
index 76e7b2bf4..b3207d387 100644
--- a/tools/server/tests/utils.py
+++ b/tools/server/tests/utils.py
@@ -101,6 +101,7 @@ class ServerProcess:
spec_type: str | None = None
spec_draft_n_min: int | None = None
spec_draft_n_max: int | None = None
+ spec_draft_n_ubatch: int | None = None
spec_synth_len: float | None = None
spec_synth_rates: List[float] | None = None
no_ui: bool | None = None
@@ -252,6 +253,8 @@ class ServerProcess:
server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max])
if self.spec_draft_n_min:
server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min])
+ if self.spec_draft_n_ubatch:
+ server_args.extend(["--spec-draft-n-ubatch", self.spec_draft_n_ubatch])
if self.spec_synth_len is not None:
server_args.extend(["--spec-synth-len", self.spec_synth_len])
if self.spec_synth_rates is not None: