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
(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) | | `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) | | `--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-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_LEN) | | `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_RATES) | | `--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/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
(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) | | `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) | | `--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-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_LEN) | | `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_RATES) | | `--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/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: