From 4388ebaa2af27db5054a6a4a975e63b05bc0f0cd Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 14 Sep 2026 01:21:43 +0200 Subject: [PATCH 1/4] speculative : rename --spec-mtp-rs-planes to --spec-draft-rs-planes Keep --spec-mtp-rs-planes as an alias. Rename the capped-replay helpers that are not specific to MTP. Assisted-by: Claude Opus 5 --- common/arg.cpp | 6 +- common/common.h | 10 +- common/speculative.cpp | 34 +++---- common/speculative.h | 4 +- docs/speculative.md | 6 +- tests/test-arg-parser.cpp | 41 ++++---- tests/test-chat.cpp | 38 +++---- tools/cli/README.md | 2 +- tools/server/README.md | 2 +- tools/server/server-context.cpp | 124 +++++++++++------------ tools/server/server-speculative-replay.h | 42 ++++---- 11 files changed, 155 insertions(+), 154 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 3acfbab6e544..815fef66666d 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4208,12 +4208,12 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } ).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-mtp-rs-planes"}, "N", + {"--spec-draft-rs-planes", "--spec-mtp-rs-planes"}, "N", "total target recurrent-state planes for draft-mtp, including the current state (default: 0, allocate spec-draft-n-max + 1)", [](common_params & params, int value) { - params.speculative.mtp_rs_planes = value; + params.speculative.rs_planes = value; } - ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_MTP_RS_PLANES")); + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_RS_PLANES")); add_opt(common_arg( {"--spec-draft-ubatch-size", "--ubatch-size-draft", "-ubd"}, "N", "physical maximum batch size for the draft context (default: 0, inherit target ubatch); " diff --git a/common/common.h b/common/common.h index 8edd79f3d053..a36636996f51 100644 --- a/common/common.h +++ b/common/common.h @@ -385,14 +385,14 @@ struct common_params_speculative { common_params_speculative_ngram_cache ngram_cache; - int32_t mtp_rs_planes = 0; // total target recurrent planes (0 = draft.n_max + 1) + int32_t rs_planes = 0; // total target recurrent planes (0 = draft.n_max + 1) bool has_dft() const { return !draft.mparams.empty(); } - bool is_mtp_rs_capped() const { - return mtp_rs_planes > 0 && int64_t(mtp_rs_planes) < int64_t(draft.n_max) + 1; + bool is_rs_capped() const { + return rs_planes > 0 && int64_t(rs_planes) < int64_t(draft.n_max) + 1; } bool has_synth() const { @@ -405,8 +405,8 @@ struct common_params_speculative { }); const bool has_mtp = std::find(types.begin(), types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != types.end(); - if (has_mtp && mtp_rs_planes > 0) { - return uint32_t(mtp_rs_planes - 1); + if (has_mtp && rs_planes > 0) { + return uint32_t(rs_planes - 1); } return needs_rs_seq ? uint32_t(std::max(0, draft.n_max)) : 0u; diff --git a/common/speculative.cpp b/common/speculative.cpp index 49f8c26656b1..1e0465bd916f 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -175,8 +175,8 @@ struct common_speculative_impl { virtual bool get_state(llama_seq_id /*seq_id*/, std::vector & /*data*/) const { return false; } virtual void set_state(llama_seq_id /*seq_id*/, const std::vector & /*data*/) {} - virtual bool get_mtp_replay_state(llama_seq_id /*seq_id*/, std::vector & /*data*/) const { return false; } - virtual bool set_mtp_replay_state(llama_seq_id /*seq_id*/, const std::vector & /*data*/) { return false; } + virtual bool get_replay_state(llama_seq_id /*seq_id*/, std::vector & /*data*/) const { return false; } + virtual bool set_replay_state(llama_seq_id /*seq_id*/, const std::vector & /*data*/) { return false; } }; struct common_speculative_impl_draft_simple : public common_speculative_impl { @@ -1782,7 +1782,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes); } - bool get_mtp_replay_state(llama_seq_id seq_id, std::vector & data) const override { + bool get_replay_state(llama_seq_id seq_id, std::vector & data) const override { if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq || pending_h[seq_id].size() != (size_t) n_embd) { return false; @@ -1797,7 +1797,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { return true; } - bool set_mtp_replay_state(llama_seq_id seq_id, const std::vector & data) override { + bool set_replay_state(llama_seq_id seq_id, const std::vector & data) override { if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { return false; } @@ -2335,22 +2335,22 @@ void common_validate_speculative_params( params.draft.n_ubatch, target_ubatch_raw)); } - if (params.mtp_rs_planes == 0) { + if (params.rs_planes == 0) { return; } if (!has_mtp) { - throw std::invalid_argument("spec-mtp-rs-planes requires --spec-type draft-mtp"); + throw std::invalid_argument("spec-draft-rs-planes requires --spec-type draft-mtp"); } if (params.draft.n_max < 1) { - throw std::invalid_argument("spec-mtp-rs-planes requires spec-draft-n-max >= 1"); + throw std::invalid_argument("spec-draft-rs-planes requires spec-draft-n-max >= 1"); } const int64_t max_planes = int64_t(params.draft.n_max) + 1; - if (params.mtp_rs_planes < 2 || int64_t(params.mtp_rs_planes) > max_planes) { + if (params.rs_planes < 2 || int64_t(params.rs_planes) > max_planes) { throw std::invalid_argument(string_format( - "spec-mtp-rs-planes must be 0 or in [2, %" PRId64 "] for spec-draft-n-max=%d", + "spec-draft-rs-planes must be 0 or in [2, %" PRId64 "] for spec-draft-n-max=%d", max_planes, params.draft.n_max)); } @@ -2360,15 +2360,15 @@ void common_validate_speculative_params( type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; }); - if (params.is_mtp_rs_capped() && has_other_recurrent_mode) { + if (params.is_rs_capped() && has_other_recurrent_mode) { throw std::invalid_argument( - "spec-mtp-rs-planes cannot be combined with another speculative mode that requires recurrent rollback"); + "spec-draft-rs-planes cannot be combined with another speculative mode that requires recurrent rollback"); } - if (params.is_mtp_rs_capped() && target_ubatch_effective > 0 && + if (params.is_rs_capped() && target_ubatch_effective > 0 && int64_t(params.draft.n_max) + 1 > target_ubatch_effective) { throw std::invalid_argument(string_format( - "capped spec-mtp-rs-planes requires ubatch-size >= spec-draft-n-max + 1 (%" PRId64 ")", + "capped spec-draft-rs-planes requires ubatch-size >= spec-draft-n-max + 1 (%" PRId64 ")", int64_t(params.draft.n_max) + 1)); } } @@ -3081,14 +3081,14 @@ void common_speculative_set_state(common_speculative * spec, llama_seq_id seq_id } } -bool common_speculative_get_mtp_state(common_speculative * spec, llama_seq_id seq_id, std::vector & data) { +bool common_speculative_get_replay_state(common_speculative * spec, llama_seq_id seq_id, std::vector & data) { data.clear(); if (spec == nullptr) { return false; } for (auto & impl : spec->impls) { - if (impl->get_mtp_replay_state(seq_id, data)) { + if (impl->get_replay_state(seq_id, data)) { return true; } } @@ -3096,14 +3096,14 @@ bool common_speculative_get_mtp_state(common_speculative * spec, llama_seq_id se return false; } -bool common_speculative_set_mtp_state( +bool common_speculative_set_replay_state( common_speculative * spec, llama_seq_id seq_id, const std::vector & data) { if (spec == nullptr) { return data.empty(); } for (auto & impl : spec->impls) { - if (impl->set_mtp_replay_state(seq_id, data)) { + if (impl->set_replay_state(seq_id, data)) { return true; } } diff --git a/common/speculative.h b/common/speculative.h index ebce4cc5911d..8d66db038d46 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -94,8 +94,8 @@ void common_speculative_accept(common_speculative * spec, llama_seq_id, uint16_t bool common_speculative_get_state(common_speculative * spec, llama_seq_id seq_id, std::vector & data); void common_speculative_set_state(common_speculative * spec, llama_seq_id seq_id, const std::vector & data); -bool common_speculative_get_mtp_state(common_speculative * spec, llama_seq_id seq_id, std::vector & data); -bool common_speculative_set_mtp_state(common_speculative * spec, llama_seq_id seq_id, const std::vector & data); +bool common_speculative_get_replay_state(common_speculative * spec, llama_seq_id seq_id, std::vector & data); +bool common_speculative_set_replay_state(common_speculative * spec, llama_seq_id seq_id, const std::vector & data); // print statistics about the speculative decoding void common_speculative_print_stats(const common_speculative * spec); diff --git a/docs/speculative.md b/docs/speculative.md index 43a9a7e2db1e..2a1bd259acaa 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -247,9 +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-mtp-rs-planes N +--spec-draft-rs-planes, --spec-mtp-rs-planes N total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1) - (env: LLAMA_ARG_SPEC_MTP_RS_PLANES) + (env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) --spec-draft-p-split, --draft-p-split P speculative decoding split probability (default: 0.10) (env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) @@ -266,7 +266,7 @@ Use exactly one of these options: #### Capped MTP recurrent planes -`--spec-mtp-rs-planes` applies only to `draft-mtp`. The default value `0` allocates `--spec-draft-n-max + 1` target recurrent-state planes. An explicit value must be in `[2, --spec-draft-n-max + 1]`; a smaller value enables capped replay. Capped replay cannot be combined with Eagle3, DFlash, or DSpark because those modes also control recurrent rollback. +`--spec-draft-rs-planes` (formerly `--spec-mtp-rs-planes`, still accepted) applies only to `draft-mtp`. The default value `0` allocates `--spec-draft-n-max + 1` target recurrent-state planes. An explicit value must be in `[2, --spec-draft-n-max + 1]`; a smaller value enables capped replay. Capped replay cannot be combined with Eagle3, DFlash, or DSpark because those modes also control recurrent rollback. The effective target ubatch is `min(batch-size, ubatch-size)`, or `batch-size` when `ubatch-size` is zero. For capped replay it must be at least `--spec-draft-n-max + 1`. A nonzero draft ubatch override must equal the target ubatch. diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 57c32fe43c98..0b9559009491 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -257,20 +257,21 @@ static void test(void) { { { "--spec-type", "draft-mtp", "--ubatch-size", "512", "--spec-draft-ubatch-size", "128" }, false, -1, -1, -1 }, { { "--spec-type", "draft-mtp", "-b", "256", "-ub", "512", "-ubd", "512" }, true, -1, -1, -1 }, { { "--spec-type", "draft-mtp", "-b", "256", "-ub", "512", "-ubd", "256" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "256", "--spec-mtp-rs-planes", "2", "-b", "256", "-ub", "512", "-ubd", "512" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "256", "--spec-mtp-rs-planes", "2", "-b", "256", "-ub", "0" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "256", "--spec-draft-rs-planes", "2", "-b", "256", "-ub", "512", "-ubd", "512" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "256", "--spec-draft-rs-planes", "2", "-b", "256", "-ub", "0" }, false, -1, -1, -1 }, { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8" }, true, 0, 8, false }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-mtp-rs-planes", "0" }, true, 0, 8, false }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "0" }, true, 0, 8, false }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "2" }, true, 2, 1, true }, { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-mtp-rs-planes", "2" }, true, 2, 1, true }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--ubatch-size", "8", "--spec-mtp-rs-planes", "2" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--ubatch-size", "9", "--spec-mtp-rs-planes", "2" }, true, 2, 1, true }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--batch-size", "8", "--ubatch-size", "9", "--spec-mtp-rs-planes", "2" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-mtp-rs-planes", "9" }, true, 9, 8, false }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-mtp-rs-planes", "-1" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-mtp-rs-planes", "1" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-mtp-rs-planes", "10" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "8", "--spec-mtp-rs-planes", "4" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-mtp,draft-eagle3", "--spec-draft-n-max", "8", "--spec-mtp-rs-planes", "4" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--ubatch-size", "8", "--spec-draft-rs-planes", "2" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--ubatch-size", "9", "--spec-draft-rs-planes", "2" }, true, 2, 1, true }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--batch-size", "8", "--ubatch-size", "9", "--spec-draft-rs-planes", "2" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "9" }, true, 9, 8, false }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "-1" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "1" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "10" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "4" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp,draft-eagle3", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "4" }, false, -1, -1, -1 }, }; for (const auto & test_case : mtp_cases) { params = common_params(); @@ -278,9 +279,9 @@ static void test(void) { argv.insert(argv.end(), test_case.args.begin(), test_case.args.end()); assert(test_case.valid == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SERVER)); if (test_case.planes >= 0) { - assert(params.speculative.mtp_rs_planes == test_case.planes); + assert(params.speculative.rs_planes == test_case.planes); assert(params.speculative.need_n_rs_seq() == (uint32_t) test_case.n_rs_seq); - assert(params.speculative.is_mtp_rs_capped() == (bool) test_case.capped); + assert(params.speculative.is_rs_capped() == (bool) test_case.capped); } } @@ -602,7 +603,7 @@ static void test_mtp_draft_ubatch_validation() { params.draft.n_ubatch = 0; params.draft.n_max = 8; - params.mtp_rs_planes = 2; + params.rs_planes = 2; rejected = false; try { common_validate_speculative_params(params, 512, 512); @@ -615,14 +616,14 @@ static void test_mtp_draft_ubatch_validation() { common_validate_speculative_params(params, 512, 512); } -static void test_mtp_state_boundaries() { +static void test_replay_state_boundaries() { std::vector state; const std::vector malformed = { 0x01, 0x02, 0x03 }; - assert(!common_speculative_get_mtp_state(nullptr, 0, state)); + assert(!common_speculative_get_replay_state(nullptr, 0, state)); assert(state.empty()); - assert(common_speculative_set_mtp_state(nullptr, 0, {})); - assert(!common_speculative_set_mtp_state(nullptr, 0, malformed)); + assert(common_speculative_set_replay_state(nullptr, 0, {})); + assert(!common_speculative_set_replay_state(nullptr, 0, malformed)); } int main(void) { @@ -630,7 +631,7 @@ int main(void) { test(); test_draft_ubatch_override(); test_mtp_draft_ubatch_validation(); - test_mtp_state_boundaries(); + test_replay_state_boundaries(); } catch (std::exception & e) { fprintf(stderr, "test-arg-parser: exception: %s\n", e.what()); return 1; diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 638ab3d3b9f7..cba5b9acd764 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -160,20 +160,20 @@ static std::unique_ptr build_grammar(const std::string & grammar_ static void test_speculative_replay_state_transitions() { server_speculative_replay_state state; - assert_equals(false, state.mtp_gpu_snapshots_armed()); - assert_equals(false, state.mtp_gpu_replay_pending()); + assert_equals(false, state.gpu_snapshots_armed()); + assert_equals(false, state.gpu_replay_pending()); assert_equals(false, state.excludes_replayed_token_from_acceptance()); - state.arm_mtp_gpu_snapshots(); - state.arm_mtp_gpu_snapshots(); - assert_equals(true, state.mtp_gpu_snapshots_armed()); - state.discard_mtp_gpu_snapshot_arm(); - assert_equals(false, state.mtp_gpu_snapshots_armed()); + state.arm_gpu_snapshots(); + state.arm_gpu_snapshots(); + assert_equals(true, state.gpu_snapshots_armed()); + state.discard_gpu_snapshot_arm(); + assert_equals(false, state.gpu_snapshots_armed()); state.begin_checkpoint_replay(); state.begin_checkpoint_replay(); assert_equals(true, state.excludes_replayed_token_from_acceptance()); - state.discard_mtp_gpu_snapshot_arm(); + state.discard_gpu_snapshot_arm(); assert_equals(true, state.excludes_replayed_token_from_acceptance()); state.finish_verification(); assert_equals(false, state.excludes_replayed_token_from_acceptance()); @@ -185,15 +185,15 @@ static void test_speculative_replay_state_transitions() { common_sampler_ptr sampler(common_sampler_init(sampler_model.get(), sampler_params)); common_sampler * sampler_ptr = sampler.get(); - state.arm_mtp_gpu_snapshots(); - state.begin_mtp_gpu_replay(expected_tokens, std::move(sampler), 2); + state.arm_gpu_snapshots(); + state.begin_gpu_replay(expected_tokens, std::move(sampler), 2); const bool sampler_moved = sampler == nullptr; - const bool replay_pending = state.mtp_gpu_replay_pending(); - const uint32_t selected_token = state.mtp_gpu_replay_selected_token(); + const bool replay_pending = state.gpu_replay_pending(); + const uint32_t selected_token = state.gpu_replay_selected_token(); llama_tokens replayed_tokens = { 99 }; common_sampler_ptr replayed_sampler; - const uint32_t accepted = state.consume_mtp_gpu_replay(replayed_tokens, replayed_sampler); + const uint32_t accepted = state.consume_gpu_replay(replayed_tokens, replayed_sampler); const bool sampler_transferred = replayed_sampler.get() == sampler_ptr; state.reset(); @@ -203,12 +203,12 @@ static void test_speculative_replay_state_transitions() { assert_equals(true, replayed_tokens == expected_tokens); assert_equals(true, sampler_transferred); assert_equals(uint32_t(2), accepted); - assert_equals(false, state.mtp_gpu_replay_pending()); - assert_equals(false, state.mtp_gpu_snapshots_armed()); + assert_equals(false, state.gpu_replay_pending()); + assert_equals(false, state.gpu_snapshots_armed()); - state.arm_mtp_gpu_snapshots(); + state.arm_gpu_snapshots(); state.reset(); - assert_equals(false, state.mtp_gpu_snapshots_armed()); + assert_equals(false, state.gpu_snapshots_armed()); } static void test_speculative_replay_failure_slot_scope() { @@ -232,7 +232,7 @@ static void test_speculative_replay_failure_slot_scope() { for (const auto & test : cases) { std::array states; for (auto & state : states) { - state.arm_mtp_gpu_snapshots(); + state.arm_gpu_snapshots(); } for (size_t slot_id = 0; slot_id < states.size(); ++slot_id) { if (server_sparse_batch_slot_is_affected( @@ -241,7 +241,7 @@ static void test_speculative_replay_failure_slot_scope() { } } for (size_t slot_id = 0; slot_id < states.size(); ++slot_id) { - assert_equals(test.affected[slot_id], !states[slot_id].mtp_gpu_snapshots_armed()); + assert_equals(test.affected[slot_id], !states[slot_id].gpu_snapshots_armed()); } } } diff --git a/tools/cli/README.md b/tools/cli/README.md index a58ed6095fbd..9eaecb003a81 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -200,7 +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-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1); see [capped MTP recurrent planes](../../docs/speculative.md#capped-mtp-recurrent-planes)
(env: LLAMA_ARG_SPEC_MTP_RS_PLANES) | +| `--spec-draft-rs-planes, --spec-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1); see [capped MTP recurrent planes](../../docs/speculative.md#capped-mtp-recurrent-planes)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) | | `--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 7a1caf006dcb..5cb35f7863dc 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -261,7 +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-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1); see [capped MTP recurrent planes](../../docs/speculative.md#capped-mtp-recurrent-planes)
(env: LLAMA_ARG_SPEC_MTP_RS_PLANES) | +| `--spec-draft-rs-planes, --spec-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1); see [capped MTP recurrent planes](../../docs/speculative.md#capped-mtp-recurrent-planes)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) | | `--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/server-context.cpp b/tools/server/server-context.cpp index bc5fbf937fef..646fd554ba84 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -128,7 +128,7 @@ struct server_batch { // track if given slot can be batched with slots already in the batch server_slot * slot_batched = nullptr; server_slot * replay_slot = nullptr; - bool mtp_sparse_snapshots = false; + bool rs_sparse_snapshots = false; // in embd mode, we temporarily swap out the tokens arr and restore it on clear() bool has_embd = false; @@ -184,7 +184,7 @@ struct server_batch { common_batch_clear(batch); slot_batched = nullptr; replay_slot = nullptr; - mtp_sparse_snapshots = false; + rs_sparse_snapshots = false; alora_scale = -1.0f; alora_disabled_id = 0; batch_rendered = false; @@ -890,8 +890,8 @@ struct server_context_impl { llama_context * ctx_tgt = nullptr; server_batch batch; - llama_seq_id capped_mtp_next_verification_slot = 0; - bool capped_mtp_prompt_turn = false; + llama_seq_id capped_rs_next_verification_slot = 0; + bool capped_rs_prompt_turn = false; llama_model * model_dft = nullptr; llama_context * ctx_dft = nullptr; @@ -1039,7 +1039,7 @@ struct server_context_impl { const bool spec_mtp = std::find(params_base.speculative.types.begin(), params_base.speculative.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params_base.speculative.types.end(); - const bool capped_mtp = spec_mtp && params_base.speculative.is_mtp_rs_capped(); + const bool capped_rs = spec_mtp && params_base.speculative.is_rs_capped(); const bool has_spec = has_draft || spec_mtp; if (callback_state) { @@ -1131,17 +1131,17 @@ struct server_context_impl { return false; } - if (spec_mtp && params_base.speculative.mtp_rs_planes > 0 && - llama_n_rs_seq(ctx_tgt) != uint32_t(params_base.speculative.mtp_rs_planes - 1)) { + if (spec_mtp && params_base.speculative.rs_planes > 0 && + llama_n_rs_seq(ctx_tgt) != uint32_t(params_base.speculative.rs_planes - 1)) { SRV_ERR("requested %d MTP recurrent planes, but the target context allocated %u; the target model does not support this recurrent rollback configuration\n", - params_base.speculative.mtp_rs_planes, llama_n_rs_seq(ctx_tgt) + 1); + params_base.speculative.rs_planes, llama_n_rs_seq(ctx_tgt) + 1); return false; } - if (capped_mtp && !llama_recurrent_sparse_snapshots_supported(ctx_tgt)) { + if (capped_rs && !llama_recurrent_sparse_snapshots_supported(ctx_tgt)) { SRV_ERR("%s", "capped MTP recurrent planes require a model graph and recurrent-state backend with selected sparse-snapshot support\n"); return false; } - if (capped_mtp && params_base.speculative.draft.n_max + 1 > (int32_t) llama_n_ubatch(ctx_tgt)) { + if (capped_rs && params_base.speculative.draft.n_max + 1 > (int32_t) llama_n_ubatch(ctx_tgt)) { SRV_ERR("capped MTP replay requires an ubatch of at least %d tokens, but the target context has %u\n", params_base.speculative.draft.n_max + 1, llama_n_ubatch(ctx_tgt)); return false; @@ -1271,8 +1271,8 @@ struct server_context_impl { } slots.clear(); - capped_mtp_next_verification_slot = 0; - capped_mtp_prompt_turn = false; + capped_rs_next_verification_slot = 0; + capped_rs_prompt_turn = false; ctx_tgt_seq_rm_type = common_context_can_seq_rm(ctx_tgt); if (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_NO) { @@ -1285,10 +1285,10 @@ struct server_context_impl { if (spec_mtp) { const uint32_t total_planes = llama_n_rs_seq(ctx_tgt) + 1; - const uint32_t direct_rollback = capped_mtp ? total_planes - 2 : total_planes - 1; + const uint32_t direct_rollback = capped_rs ? total_planes - 2 : total_planes - 1; SRV_INF("MTP recurrent-plane policy: draft depth = %d, total planes = %u, direct rollback horizon = %u, full-shape GPU replay = %s\n", params_base.speculative.draft.n_max, total_planes, direct_rollback, - capped_mtp ? "enabled" : "disabled"); + capped_rs ? "enabled" : "disabled"); } // setup slots @@ -2814,7 +2814,7 @@ struct server_context_impl { } void iterate_post_decode(std::function callback) { - if (batch.replay_slot == nullptr && !batch.mtp_sparse_snapshots) { + if (batch.replay_slot == nullptr && !batch.rs_sparse_snapshots) { iterate(slots, std::move(callback)); return; } @@ -2835,12 +2835,12 @@ struct server_context_impl { void clear_failed_sparse_batch_state() { const int32_t replay_slot_id = batch.replay_slot != nullptr ? batch.replay_slot->id : -1; - if (replay_slot_id < 0 && !batch.mtp_sparse_snapshots) { + if (replay_slot_id < 0 && !batch.rs_sparse_snapshots) { return; } for (auto & slot : slots) { if (slot.is_processing() && server_sparse_batch_slot_is_affected( - replay_slot_id, batch.mtp_sparse_snapshots, batch.tokens, slot.id)) { + replay_slot_id, batch.rs_sparse_snapshots, batch.tokens, slot.id)) { slot.prompt_clear(); } } @@ -3078,9 +3078,9 @@ struct server_context_impl { std::vector generating; std::vector drafting; - if (params_base.speculative.is_mtp_rs_capped()) { + if (params_base.speculative.is_rs_capped()) { iterate(slots, [&](server_slot & slot) { - if (batch.replay_slot == nullptr && slot.state == SLOT_STATE_GENERATING && slot.spec_replay.mtp_gpu_replay_pending()) { + if (batch.replay_slot == nullptr && slot.state == SLOT_STATE_GENERATING && slot.spec_replay.gpu_replay_pending()) { batch.replay_slot = &slot; } }); @@ -3099,7 +3099,7 @@ struct server_context_impl { fail_speculative_replay(slot, "failed to rewind the MTP draft context for GPU replay"); return; } - if (!common_speculative_set_mtp_state(spec.get(), slot.id, ckpt.data_spec)) { + if (!common_speculative_set_replay_state(spec.get(), slot.id, ckpt.data_spec)) { fail_speculative_replay(slot, "failed to restore speculative state for GPU replay"); return; } @@ -3147,8 +3147,8 @@ struct server_context_impl { llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id), llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id)); slot.spec_ckpt.data_spec.clear(); - if (params_base.speculative.is_mtp_rs_capped()) { - if (!common_speculative_get_mtp_state(spec.get(), slot.id, slot.spec_ckpt.data_spec)) { + if (params_base.speculative.is_rs_capped()) { + if (!common_speculative_get_replay_state(spec.get(), slot.id, slot.spec_ckpt.data_spec)) { throw std::runtime_error("failed to capture capped MTP speculative state"); } } @@ -3202,12 +3202,12 @@ struct server_context_impl { } if (!draft.empty()) { - const bool capped_mtp = params_base.speculative.is_mtp_rs_capped(); - if (capped_mtp) { - slot.spec_replay.arm_mtp_gpu_snapshots(); + const bool capped_rs = params_base.speculative.is_rs_capped(); + if (capped_rs) { + slot.spec_replay.arm_gpu_snapshots(); } const bool use_ckpt_tgt = - !capped_mtp && (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || + !capped_rs && (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_tgt))); const bool use_ckpt_dft = @@ -3234,62 +3234,62 @@ struct server_context_impl { }); // update the batch with the sampled/drafted tokens - const bool capped_mtp_verification = params_base.speculative.is_mtp_rs_capped() && + const bool capped_rs_verification = params_base.speculative.is_rs_capped() && batch.replay_slot == nullptr && !generating.empty(); - if (capped_mtp_verification) { + if (capped_rs_verification) { const auto first = std::find_if(generating.begin(), generating.end(), [&](const server_slot * slot) { - return slot->id >= capped_mtp_next_verification_slot; + return slot->id >= capped_rs_next_verification_slot; }); if (first != generating.end()) { std::rotate(generating.begin(), first, generating.end()); } } - const bool capped_mtp_empty_draft_class = capped_mtp_verification && generating.front()->spec_draft.empty(); - const size_t capped_mtp_verification_size = capped_mtp_verification && !capped_mtp_empty_draft_class + const bool capped_rs_empty_draft_class = capped_rs_verification && generating.front()->spec_draft.empty(); + const size_t capped_rs_verification_size = capped_rs_verification && !capped_rs_empty_draft_class ? generating.front()->spec_draft.size() + 1 : 0; - const bool capped_mtp_prompt_ready = capped_mtp_verification && std::any_of(slots.begin(), slots.end(), [](const server_slot & slot) { + const bool capped_rs_prompt_ready = capped_rs_verification && std::any_of(slots.begin(), slots.end(), [](const server_slot & slot) { return slot.is_processing() && (slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT); }); - bool capped_mtp_prompt_only = false; - if (capped_mtp_verification_size > 1 && capped_mtp_prompt_ready) { - capped_mtp_prompt_only = capped_mtp_prompt_turn; - capped_mtp_prompt_turn = !capped_mtp_prompt_turn; + bool capped_rs_prompt_only = false; + if (capped_rs_verification_size > 1 && capped_rs_prompt_ready) { + capped_rs_prompt_only = capped_rs_prompt_turn; + capped_rs_prompt_turn = !capped_rs_prompt_turn; } - if (capped_mtp_prompt_only) { + if (capped_rs_prompt_only) { slot_batched = nullptr; } - const size_t capped_mtp_batch_limit = std::min(llama_n_batch(ctx_tgt), llama_n_ubatch(ctx_tgt)); - size_t capped_mtp_verification_tokens = 0; - server_slot * capped_mtp_scheduled_anchor = nullptr; - if (!capped_mtp_prompt_only) { + const size_t capped_rs_batch_limit = std::min(llama_n_batch(ctx_tgt), llama_n_ubatch(ctx_tgt)); + size_t capped_rs_verification_tokens = 0; + server_slot * capped_rs_scheduled_anchor = nullptr; + if (!capped_rs_prompt_only) { iterate(generating, [&](server_slot & slot) { - if (capped_mtp_verification) { - if (slot.spec_draft.empty() != capped_mtp_empty_draft_class) { + if (capped_rs_verification) { + if (slot.spec_draft.empty() != capped_rs_empty_draft_class) { return; } - if (!capped_mtp_empty_draft_class && slot.spec_draft.size() + 1 != capped_mtp_verification_size) { + if (!capped_rs_empty_draft_class && slot.spec_draft.size() + 1 != capped_rs_verification_size) { return; } - if (!capped_mtp_empty_draft_class && - capped_mtp_verification_tokens + capped_mtp_verification_size > capped_mtp_batch_limit) { + if (!capped_rs_empty_draft_class && + capped_rs_verification_tokens + capped_rs_verification_size > capped_rs_batch_limit) { return; } } slot.handle_last_sampled_token(batch); - if (capped_mtp_verification && capped_mtp_scheduled_anchor == nullptr) { - capped_mtp_scheduled_anchor = &slot; + if (capped_rs_verification && capped_rs_scheduled_anchor == nullptr) { + capped_rs_scheduled_anchor = &slot; } - if (!capped_mtp_empty_draft_class && capped_mtp_verification_size != 0) { - capped_mtp_verification_tokens += capped_mtp_verification_size; - batch.mtp_sparse_snapshots = true; + if (!capped_rs_empty_draft_class && capped_rs_verification_size != 0) { + capped_rs_verification_tokens += capped_rs_verification_size; + batch.rs_sparse_snapshots = true; } }); } - if (capped_mtp_scheduled_anchor != nullptr) { - capped_mtp_next_verification_slot = capped_mtp_scheduled_anchor->id + 1; + if (capped_rs_scheduled_anchor != nullptr) { + capped_rs_next_verification_slot = capped_rs_scheduled_anchor->id + 1; } // process in chunks of params.n_batch @@ -3300,7 +3300,7 @@ struct server_context_impl { auto & alora_disabled_id = batch.alora_disabled_id; // next, batch any pending prompts without exceeding n_batch - if (batch.replay_slot == nullptr && (capped_mtp_prompt_only || capped_mtp_verification_size == 0) && + if (batch.replay_slot == nullptr && (capped_rs_prompt_only || capped_rs_verification_size == 0) && (params_base.cont_batching || batch.size() == 0)) { bool add_ok = true; // false means the batch is full, skip remaining slots @@ -3878,12 +3878,12 @@ struct server_context_impl { for (int32_t i = 0; i < batch_view.n_tokens; ++i) { GGML_ASSERT(batch.tokens[i].id_slot == replay_slot->id); } - GGML_ASSERT(replay_slot->spec_replay.mtp_gpu_replay_selected_token() < (uint32_t) batch_view.n_tokens); + GGML_ASSERT(replay_slot->spec_replay.gpu_replay_selected_token() < (uint32_t) batch_view.n_tokens); } - const bool sparse_snapshots = replay_slot != nullptr || batch.mtp_sparse_snapshots; + const bool sparse_snapshots = replay_slot != nullptr || batch.rs_sparse_snapshots; const int32_t selected_token = replay_slot != nullptr - ? (int32_t) replay_slot->spec_replay.mtp_gpu_replay_selected_token() : -1; + ? (int32_t) replay_slot->spec_replay.gpu_replay_selected_token() : -1; if (sparse_snapshots) { GGML_ASSERT(off == 0); GGML_ASSERT(batch_view.n_tokens == batch.size()); @@ -4147,7 +4147,7 @@ struct server_context_impl { const size_t n_draft = slot.spec_draft.size(); GGML_ASSERT(n_draft > 0); - const bool gpu_snapshot_replay = slot.spec_replay.mtp_gpu_replay_pending(); + const bool gpu_snapshot_replay = slot.spec_replay.gpu_replay_pending(); // verify and try to accept the draft { @@ -4167,7 +4167,7 @@ struct server_context_impl { slot.spec_i_batch.clear(); if (gpu_snapshot_replay) { - const uint32_t replay_accepted = slot.spec_replay.consume_mtp_gpu_replay(accepted, slot.smpl); + const uint32_t replay_accepted = slot.spec_replay.consume_gpu_replay(accepted, slot.smpl); const auto & ckpt = slot.spec_ckpt; const llama_pos accepted_end = ckpt.pos_max + 2 + replay_accepted; @@ -4188,7 +4188,7 @@ struct server_context_impl { const uint32_t n_rollback = slot.spec_draft.size() + 1 - accepted.size(); const uint32_t direct_horizon = llama_n_rs_seq(ctx_tgt) > 0 ? llama_n_rs_seq(ctx_tgt) - 1 : 0; - const bool use_gpu_replay = slot.spec_replay.mtp_gpu_snapshots_armed() && n_rollback > direct_horizon; + const bool use_gpu_replay = slot.spec_replay.gpu_snapshots_armed() && n_rollback > direct_horizon; if (use_gpu_replay) { const uint32_t n_accepted = accepted.size() - 1; @@ -4199,7 +4199,7 @@ struct server_context_impl { n_accepted, slot.spec_draft.size()); } - slot.spec_replay.begin_mtp_gpu_replay( + slot.spec_replay.begin_gpu_replay( std::move(accepted), std::move(slot.smpl), n_accepted); slot.prompt.tokens.keep_first(ckpt.n_tokens); slot.smpl = std::move(smpl_save); @@ -4247,7 +4247,7 @@ struct server_context_impl { common_speculative_accept(spec.get(), slot.id, accepted.size() - 1); slot.spec_draft = std::move(accepted); - slot.spec_replay.discard_mtp_gpu_snapshot_arm(); + slot.spec_replay.discard_gpu_snapshot_arm(); } } diff --git a/tools/server/server-speculative-replay.h b/tools/server/server-speculative-replay.h index 30c5b1072182..88d4296bd094 100644 --- a/tools/server/server-speculative-replay.h +++ b/tools/server/server-speculative-replay.h @@ -27,26 +27,26 @@ bool server_sparse_batch_slot_is_affected( // Checkpoint and capped-MTP GPU replay use the same slot lifetime. // Only GPU replay owns accepted tokens and its sampler snapshot. struct server_speculative_replay_state { - bool mtp_gpu_snapshots_armed() const { - return phase == phase_type::MTP_GPU_SNAPSHOTS_ARMED; + bool gpu_snapshots_armed() const { + return phase == phase_type::GPU_SNAPSHOTS_ARMED; } - bool mtp_gpu_replay_pending() const { - return phase == phase_type::MTP_GPU_REPLAY_PENDING; + bool gpu_replay_pending() const { + return phase == phase_type::GPU_REPLAY_PENDING; } bool excludes_replayed_token_from_acceptance() const { return phase == phase_type::CHECKPOINT_REPLAY; } - uint32_t mtp_gpu_replay_selected_token() const { - GGML_ASSERT(mtp_gpu_replay_pending()); + uint32_t gpu_replay_selected_token() const { + GGML_ASSERT(gpu_replay_pending()); return n_accepted; } - void arm_mtp_gpu_snapshots() { - GGML_ASSERT(phase == phase_type::IDLE || phase == phase_type::MTP_GPU_SNAPSHOTS_ARMED); - phase = phase_type::MTP_GPU_SNAPSHOTS_ARMED; + void arm_gpu_snapshots() { + GGML_ASSERT(phase == phase_type::IDLE || phase == phase_type::GPU_SNAPSHOTS_ARMED); + phase = phase_type::GPU_SNAPSHOTS_ARMED; } void begin_checkpoint_replay() { @@ -54,16 +54,16 @@ struct server_speculative_replay_state { phase = phase_type::CHECKPOINT_REPLAY; } - void begin_mtp_gpu_replay(llama_tokens tokens, common_sampler_ptr sampler, uint32_t accepted) { - GGML_ASSERT(phase == phase_type::MTP_GPU_SNAPSHOTS_ARMED && sampler != nullptr); + void begin_gpu_replay(llama_tokens tokens, common_sampler_ptr sampler, uint32_t accepted) { + GGML_ASSERT(phase == phase_type::GPU_SNAPSHOTS_ARMED && sampler != nullptr); accepted_tokens = std::move(tokens); accepted_sampler = std::move(sampler); n_accepted = accepted; - phase = phase_type::MTP_GPU_REPLAY_PENDING; + phase = phase_type::GPU_REPLAY_PENDING; } - uint32_t consume_mtp_gpu_replay(llama_tokens & tokens, common_sampler_ptr & sampler) { - GGML_ASSERT(mtp_gpu_replay_pending() && accepted_sampler != nullptr); + uint32_t consume_gpu_replay(llama_tokens & tokens, common_sampler_ptr & sampler) { + GGML_ASSERT(gpu_replay_pending() && accepted_sampler != nullptr); tokens = std::move(accepted_tokens); sampler = std::move(accepted_sampler); const uint32_t accepted = n_accepted; @@ -72,18 +72,18 @@ struct server_speculative_replay_state { return accepted; } - void discard_mtp_gpu_snapshot_arm() { + void discard_gpu_snapshot_arm() { GGML_ASSERT(phase == phase_type::IDLE || - phase == phase_type::MTP_GPU_SNAPSHOTS_ARMED || + phase == phase_type::GPU_SNAPSHOTS_ARMED || phase == phase_type::CHECKPOINT_REPLAY); - if (phase == phase_type::MTP_GPU_SNAPSHOTS_ARMED) { + if (phase == phase_type::GPU_SNAPSHOTS_ARMED) { phase = phase_type::IDLE; } } void finish_verification() { - GGML_ASSERT(phase != phase_type::MTP_GPU_SNAPSHOTS_ARMED && - phase != phase_type::MTP_GPU_REPLAY_PENDING); + GGML_ASSERT(phase != phase_type::GPU_SNAPSHOTS_ARMED && + phase != phase_type::GPU_REPLAY_PENDING); reset(); } @@ -95,9 +95,9 @@ struct server_speculative_replay_state { private: enum class phase_type { IDLE, - MTP_GPU_SNAPSHOTS_ARMED, + GPU_SNAPSHOTS_ARMED, CHECKPOINT_REPLAY, - MTP_GPU_REPLAY_PENDING, + GPU_REPLAY_PENDING, }; void clear_payload() { From 5844ed26d4a9d65041025e6a69ce2272e05d5dff Mon Sep 17 00:00:00 2001 From: piggidragon Date: Mon, 14 Sep 2026 01:24:29 +0200 Subject: [PATCH 2/4] speculative : allow capped recurrent planes for DFlash DFlash keeps its draft state in the draft KV cache, which the capped replay already rewinds with seq_rm, so it has no replay state to save. The planes stay owned by a single speculative mode. Assisted-by: Claude Opus 5 --- common/arg.cpp | 2 +- common/common.h | 8 +++++--- common/speculative.cpp | 27 ++++++++++++++++++++++----- docs/speculative.md | 10 ++++++---- tests/test-arg-parser.cpp | 8 +++++++- tools/cli/README.md | 2 +- tools/server/README.md | 2 +- tools/server/server-context.cpp | 20 ++++++++++---------- 8 files changed, 53 insertions(+), 26 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 815fef66666d..5ed110aac3f7 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4209,7 +4209,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).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-rs-planes", "--spec-mtp-rs-planes"}, "N", - "total target recurrent-state planes for draft-mtp, including the current state (default: 0, allocate spec-draft-n-max + 1)", + "total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, allocate spec-draft-n-max + 1)", [](common_params & params, int value) { params.speculative.rs_planes = value; } diff --git a/common/common.h b/common/common.h index a36636996f51..3604ee9bed78 100644 --- a/common/common.h +++ b/common/common.h @@ -385,7 +385,7 @@ struct common_params_speculative { common_params_speculative_ngram_cache ngram_cache; - int32_t rs_planes = 0; // total target recurrent planes (0 = draft.n_max + 1) + int32_t rs_planes = 0; // total target recurrent planes for draft-mtp or draft-dflash (0 = draft.n_max + 1) bool has_dft() const { return !draft.mparams.empty(); @@ -404,8 +404,10 @@ struct common_params_speculative { return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; }); - const bool has_mtp = std::find(types.begin(), types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != types.end(); - if (has_mtp && rs_planes > 0) { + const bool can_cap = std::any_of(types.begin(), types.end(), [](auto t) { + return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH; + }); + if (can_cap && rs_planes > 0) { return uint32_t(rs_planes - 1); } diff --git a/common/speculative.cpp b/common/speculative.cpp index 1e0465bd916f..d1e952b0477a 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1328,6 +1328,19 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override { // noop } + + // the draft state is the ctx_dft KV cache, and a capped replay rewinds it with seq_rm + bool get_replay_state(llama_seq_id seq_id, std::vector & data) const override { + if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { + return false; + } + data.clear(); + return true; + } + + bool set_replay_state(llama_seq_id seq_id, const std::vector & data) override { + return seq_id >= 0 && seq_id < (llama_seq_id) n_seq && data.empty(); + } }; struct common_speculative_impl_draft_mtp : public common_speculative_impl { @@ -2339,8 +2352,10 @@ void common_validate_speculative_params( return; } - if (!has_mtp) { - throw std::invalid_argument("spec-draft-rs-planes requires --spec-type draft-mtp"); + const bool has_dflash = std::find( + params.types.begin(), params.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != params.types.end(); + if (!has_mtp && !has_dflash) { + throw std::invalid_argument("spec-draft-rs-planes requires --spec-type draft-mtp or draft-dflash"); } if (params.draft.n_max < 1) { @@ -2354,13 +2369,15 @@ void common_validate_speculative_params( max_planes, params.draft.n_max)); } - const bool has_other_recurrent_mode = std::any_of( + // the planes are shared by every mode that rolls back the recurrent state, so only one may use them + const auto n_recurrent_modes = std::count_if( params.types.begin(), params.types.end(), [](common_speculative_type type) { - return type == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || + return type == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || + type == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; }); - if (params.is_rs_capped() && has_other_recurrent_mode) { + if (params.is_rs_capped() && n_recurrent_modes > 1) { throw std::invalid_argument( "spec-draft-rs-planes cannot be combined with another speculative mode that requires recurrent rollback"); } diff --git a/docs/speculative.md b/docs/speculative.md index 2a1bd259acaa..f5668f465f64 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -248,7 +248,7 @@ Use exactly one of these options: minimum number of draft tokens to use for speculative decoding (default: 0) (env: LLAMA_ARG_SPEC_DRAFT_N_MIN) --spec-draft-rs-planes, --spec-mtp-rs-planes N - total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1) + total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, use spec-draft-n-max + 1) (env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) --spec-draft-p-split, --draft-p-split P speculative decoding split probability (default: 0.10) @@ -264,11 +264,13 @@ Use exactly one of these options: (use --list-devices to see available devices) ``` -#### Capped MTP recurrent planes +#### Capped recurrent planes -`--spec-draft-rs-planes` (formerly `--spec-mtp-rs-planes`, still accepted) applies only to `draft-mtp`. The default value `0` allocates `--spec-draft-n-max + 1` target recurrent-state planes. An explicit value must be in `[2, --spec-draft-n-max + 1]`; a smaller value enables capped replay. Capped replay cannot be combined with Eagle3, DFlash, or DSpark because those modes also control recurrent rollback. +`--spec-draft-rs-planes` (formerly `--spec-mtp-rs-planes`, still accepted) applies to `draft-mtp` and `draft-dflash`. The default value `0` allocates `--spec-draft-n-max + 1` target recurrent-state planes. An explicit value must be in `[2, --spec-draft-n-max + 1]`; a smaller value enables capped replay. The planes belong to one speculative mode, so capped replay cannot be combined with a second mode that rolls back the recurrent state (MTP, Eagle3, DFlash or DSpark), and DSpark and Eagle3 cannot use it. -The effective target ubatch is `min(batch-size, ubatch-size)`, or `batch-size` when `ubatch-size` is zero. For capped replay it must be at least `--spec-draft-n-max + 1`. A nonzero draft ubatch override must equal the target ubatch. +Each plane holds one copy of the recurrent state for every sequence, so the saving is `n_parallel * (n_max + 1 - planes)` copies. DFlash keeps its own state in the draft KV cache, which a replay rewinds together with the target, so no extra speculative state is saved for it. + +The effective target ubatch is `min(batch-size, ubatch-size)`, or `batch-size` when `ubatch-size` is zero. For capped replay it must be at least `--spec-draft-n-max + 1`. For `draft-mtp`, a nonzero draft ubatch override must equal the target ubatch. Every device selected for the recurrent graph operations must support sparse snapshots. The server rejects an unsupported device layout at startup; it does not move the operations to another device or fall back to CPU. diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 0b9559009491..8651de0a3617 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -270,7 +270,10 @@ static void test(void) { { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "-1" }, false, -1, -1, -1 }, { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "1" }, false, -1, -1, -1 }, { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "10" }, false, -1, -1, -1 }, - { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "4" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "4" }, true, 4, 3, true }, + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "8", "--ubatch-size", "8", "--spec-draft-rs-planes", "4" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-dspark", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "4" }, false, -1, -1, -1 }, + { { "--spec-type", "draft-mtp,draft-dflash", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "4" }, false, -1, -1, -1 }, { { "--spec-type", "draft-mtp,draft-eagle3", "--spec-draft-n-max", "8", "--spec-draft-rs-planes", "4" }, false, -1, -1, -1 }, }; for (const auto & test_case : mtp_cases) { @@ -604,6 +607,9 @@ static void test_mtp_draft_ubatch_validation() { params.draft.n_ubatch = 0; params.draft.n_max = 8; params.rs_planes = 2; + common_validate_speculative_params(params, 512, 512); + + params.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK }; rejected = false; try { common_validate_speculative_params(params, 512, 512); diff --git a/tools/cli/README.md b/tools/cli/README.md index 9eaecb003a81..2556717f6993 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -200,7 +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-rs-planes, --spec-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1); see [capped MTP recurrent planes](../../docs/speculative.md#capped-mtp-recurrent-planes)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) | +| `--spec-draft-rs-planes, --spec-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, use spec-draft-n-max + 1); see [capped recurrent planes](../../docs/speculative.md#capped-recurrent-planes)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) | | `--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 5cb35f7863dc..1926fdb7eefd 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -261,7 +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-rs-planes, --spec-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp, including the current state (default: 0, use spec-draft-n-max + 1); see [capped MTP recurrent planes](../../docs/speculative.md#capped-mtp-recurrent-planes)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) | +| `--spec-draft-rs-planes, --spec-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, use spec-draft-n-max + 1); see [capped recurrent planes](../../docs/speculative.md#capped-recurrent-planes)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) | | `--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/server-context.cpp b/tools/server/server-context.cpp index 646fd554ba84..9ae69f19f315 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1039,7 +1039,7 @@ struct server_context_impl { const bool spec_mtp = std::find(params_base.speculative.types.begin(), params_base.speculative.types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params_base.speculative.types.end(); - const bool capped_rs = spec_mtp && params_base.speculative.is_rs_capped(); + const bool capped_rs = params_base.speculative.is_rs_capped(); const bool has_spec = has_draft || spec_mtp; if (callback_state) { @@ -1131,18 +1131,18 @@ struct server_context_impl { return false; } - if (spec_mtp && params_base.speculative.rs_planes > 0 && + if (params_base.speculative.rs_planes > 0 && llama_n_rs_seq(ctx_tgt) != uint32_t(params_base.speculative.rs_planes - 1)) { - SRV_ERR("requested %d MTP recurrent planes, but the target context allocated %u; the target model does not support this recurrent rollback configuration\n", + SRV_ERR("requested %d recurrent planes, but the target context allocated %u; the target model does not support this recurrent rollback configuration\n", params_base.speculative.rs_planes, llama_n_rs_seq(ctx_tgt) + 1); return false; } if (capped_rs && !llama_recurrent_sparse_snapshots_supported(ctx_tgt)) { - SRV_ERR("%s", "capped MTP recurrent planes require a model graph and recurrent-state backend with selected sparse-snapshot support\n"); + SRV_ERR("%s", "capped recurrent planes require a model graph and recurrent-state backend with selected sparse-snapshot support\n"); return false; } if (capped_rs && params_base.speculative.draft.n_max + 1 > (int32_t) llama_n_ubatch(ctx_tgt)) { - SRV_ERR("capped MTP replay requires an ubatch of at least %d tokens, but the target context has %u\n", + SRV_ERR("capped recurrent replay requires an ubatch of at least %d tokens, but the target context has %u\n", params_base.speculative.draft.n_max + 1, llama_n_ubatch(ctx_tgt)); return false; } @@ -1283,10 +1283,10 @@ struct server_context_impl { SRV_TRC("%s", "speculative decoding will use checkpoints\n"); } - if (spec_mtp) { + if (has_spec && llama_n_rs_seq(ctx_tgt) > 0) { const uint32_t total_planes = llama_n_rs_seq(ctx_tgt) + 1; const uint32_t direct_rollback = capped_rs ? total_planes - 2 : total_planes - 1; - SRV_INF("MTP recurrent-plane policy: draft depth = %d, total planes = %u, direct rollback horizon = %u, full-shape GPU replay = %s\n", + SRV_INF("recurrent-plane policy: draft depth = %d, total planes = %u, direct rollback horizon = %u, full-shape GPU replay = %s\n", params_base.speculative.draft.n_max, total_planes, direct_rollback, capped_rs ? "enabled" : "disabled"); } @@ -3096,7 +3096,7 @@ struct server_context_impl { } if (ctx_dft != nullptr && !llama_memory_seq_rm( llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) { - fail_speculative_replay(slot, "failed to rewind the MTP draft context for GPU replay"); + fail_speculative_replay(slot, "failed to rewind the draft context for GPU replay"); return; } if (!common_speculative_set_replay_state(spec.get(), slot.id, ckpt.data_spec)) { @@ -3149,7 +3149,7 @@ struct server_context_impl { slot.spec_ckpt.data_spec.clear(); if (params_base.speculative.is_rs_capped()) { if (!common_speculative_get_replay_state(spec.get(), slot.id, slot.spec_ckpt.data_spec)) { - throw std::runtime_error("failed to capture capped MTP speculative state"); + throw std::runtime_error("failed to capture the speculative state for capped replay"); } } @@ -4177,7 +4177,7 @@ struct server_context_impl { } if (ctx_dft != nullptr && !llama_memory_seq_rm( llama_get_memory(ctx_dft), slot.id, accepted_end, -1)) { - fail_speculative_replay(slot, "failed to align the MTP draft context after GPU replay"); + fail_speculative_replay(slot, "failed to align the draft context after GPU replay"); return; } From e46fd67ed014ad7f9533fe18ad24ac7d087a577a Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 19 Sep 2026 03:17:01 +0200 Subject: [PATCH 3/4] ggml : write leading state snapshots in gated_delta_net The reserved-input mode only kept the states of the last tokens. Add leading_snapshots, which keeps the states after the first tokens in the slots below the input slot, so a caller can pick where the snapshots are. Assisted-by: Claude Opus 5 --- ggml/include/ggml.h | 6 +- ggml/src/ggml-cpu/ops.cpp | 8 ++- ggml/src/ggml-cuda/gated_delta_net.cu | 29 ++++++---- ggml/src/ggml-cuda/ggml-cuda.cu | 6 +- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 3 + .../vulkan-shaders/gated_delta_net.comp | 9 ++- ggml/src/ggml.c | 17 ++++-- tests/test-backend-ops.cpp | 57 +++++++++++-------- 8 files changed, 86 insertions(+), 49 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index e294420d9825..f7773e3d270f 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2651,7 +2651,8 @@ extern "C" { // By default, slot 0 is the final state and up to min(n_tokens, K) trailing states are written most-recent first. // For K > 1, trailing-only mode accepts (trailing_snapshots, selected_token, reserve_input) == ([0, K], -1, false) and writes up to min(n_tokens, trailing_snapshots) states most-recent first. // Selected-token mode requires (trailing_snapshots, selected_token, reserve_input) == (0, [0, n_tokens), false) and writes only the state after selected_token to slot 0. - // Reserved-input mode requires (trailing_snapshots, selected_token, reserve_input) == (min(n_tokens, K - 1), -1, true) and writes the input state to slot K - 1 and trailing states most-recent first from slot 0. + // Reserved-input mode requires (trailing_snapshots + leading_snapshots, selected_token, reserve_input) == (min(n_tokens, K - 1), -1, true) and writes the input state to slot K - 1, trailing states most-recent first from slot 0, and the state after token t < leading_snapshots to slot K - 2 - t. + // Other modes require leading_snapshots == 0. // State slots not selected by these modes stay untouched. // K == 1 requires (trailing_snapshots, selected_token, reserve_input) == (1, -1, false) and writes the final state to slot 0. GGML_API struct ggml_tensor * ggml_gated_delta_net( @@ -2668,7 +2669,8 @@ extern "C" { struct ggml_tensor * tensor, int32_t trailing_snapshots, int32_t selected_token, - bool reserve_input); + bool reserve_input, + int32_t leading_snapshots); GGML_API bool ggml_gated_delta_net_has_default_snapshot_params( const struct ggml_tensor * tensor); diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 5f28f22414b0..c701064835c3 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11010,6 +11010,7 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( const int32_t trailing_snapshots = ggml_get_op_params_i32(dst, 1); const int32_t selected_token = ggml_get_op_params_i32(dst, 2); const bool reserve_input = ggml_get_op_params_i32(dst, 3) != 0; + const int32_t leading_snapshots = ggml_get_op_params_i32(dst, 4); GGML_ASSERT(K >= 1); // per-seq stride in floats (seq s starts at state + s * seq_stride) const int64_t state_seq_stride = src_state->nb[3] / sizeof(float); @@ -11114,9 +11115,12 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( attn_data += S_v * H; // advance to next token if (K > 1) { - const int64_t target_slot = selected_token >= 0 ? (t == selected_token ? 0 : -1) : n_tokens - 1 - t; + int64_t target_slot = selected_token >= 0 ? (t == selected_token ? 0 : -1) : n_tokens - 1 - t; const int32_t snapshot_slots = selected_token >= 0 ? 1 : trailing_snapshots; - if (target_slot >= 0 && target_slot < snapshot_slots) { + if (target_slot >= snapshot_slots) { + target_slot = t < leading_snapshots ? K - 2 - t : -1; + } + if (target_slot >= 0) { float * curr_state_o = state_out_base + target_slot * state_size_per_snap + (iv3 * H + iv1) * S_v * S_v; memcpy(curr_state_o, s_out, S_v * S_v * sizeof(float)); diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 7812ba2f85e1..add0057530c1 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -65,7 +65,8 @@ gated_delta_net_cuda(const float * q, int K, int trailing_snapshots, int selected_token, - bool reserve_input) { + bool reserve_input, + int leading_snapshots) { const uint32_t h_idx = blockIdx.x; const uint32_t sequence = blockIdx.y; // each warp owns one column, using warp-level primitives to reduce across rows @@ -192,11 +193,14 @@ gated_delta_net_cuda(const float * q, if constexpr (keep_rs_t) { // Slots not selected by the snapshot parameters are caller-owned. - const int target_slot = selected_token >= 0 + int target_slot = selected_token >= 0 ? (t == selected_token ? 0 : -1) : (int) n_tokens - 1 - t; const int snapshot_slots = selected_token >= 0 ? 1 : trailing_snapshots; - if (target_slot >= 0 && target_slot < snapshot_slots) { + if (target_slot >= snapshot_slots) { + target_slot = t < leading_snapshots ? K - 2 - t : -1; + } + if (target_slot >= 0) { float * curr_state = state + target_slot * state_slot_stride; #pragma unroll for (int r = 0; r < rows_per_lane; r++) { @@ -227,7 +231,7 @@ static void launch_gated_delta_net( int64_t sb1, int64_t sb2, int64_t sb3, int64_t neqk1, int64_t rq3, float scale, int64_t state_slot_stride, int K, - int trailing_snapshots, int selected_token, bool reserve_input, cudaStream_t stream) { + int trailing_snapshots, int selected_token, bool reserve_input, int leading_snapshots, cudaStream_t stream) { //TODO: Add chunked kernel for even faster pre-fill const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; const int num_warps = 4; @@ -244,21 +248,21 @@ static void launch_gated_delta_net( q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, - trailing_snapshots, selected_token, reserve_input); + trailing_snapshots, selected_token, reserve_input, leading_snapshots); break; case 32: ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t>, launch_params, q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, - trailing_snapshots, selected_token, reserve_input); + trailing_snapshots, selected_token, reserve_input, leading_snapshots); break; case 64: { ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t>, launch_params, q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, - trailing_snapshots, selected_token, reserve_input); + trailing_snapshots, selected_token, reserve_input, leading_snapshots); break; } case 128: { @@ -266,7 +270,7 @@ static void launch_gated_delta_net( q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, - trailing_snapshots, selected_token, reserve_input); + trailing_snapshots, selected_token, reserve_input, leading_snapshots); break; } default: @@ -344,6 +348,7 @@ static void ggml_cuda_op_gated_delta_net_impl( const int trailing_snapshots = ggml_get_op_params_i32(dst, 1); const int selected_token = ggml_get_op_params_i32(dst, 2); const bool reserve_input = ggml_get_op_params_i32(dst, 3) != 0; + const int leading_snapshots = ggml_get_op_params_i32(dst, 4); const bool keep_rs = K > 1; // recurrent state -> gdn_out tail (after attention scores), or the cache when fusing @@ -359,24 +364,24 @@ static void ggml_cuda_op_gated_delta_net_impl( launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, - trailing_snapshots, selected_token, reserve_input, stream); + trailing_snapshots, selected_token, reserve_input, leading_snapshots, stream); } else { launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, - trailing_snapshots, selected_token, reserve_input, stream); + trailing_snapshots, selected_token, reserve_input, leading_snapshots, stream); } } else { if (keep_rs) { launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, - trailing_snapshots, selected_token, reserve_input, stream); + trailing_snapshots, selected_token, reserve_input, leading_snapshots, stream); } else { launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, - trailing_snapshots, selected_token, reserve_input, stream); + trailing_snapshots, selected_token, reserve_input, leading_snapshots, stream); } } } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 4e302d9f31c0..b670c3cc4212 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2822,6 +2822,7 @@ static int ggml_cuda_try_gdn_cache_fusion( const int32_t trailing = ggml_get_op_params_i32(gdn, 1); const int32_t selected = ggml_get_op_params_i32(gdn, 2); const bool reserve = ggml_get_op_params_i32(gdn, 3) != 0; + const int32_t leading = ggml_get_op_params_i32(gdn, 4); const int64_t n_written = selected >= 0 ? 1 : reserve ? trailing : std::min(n_tokens, K); // snapshot tail starts right after the attention scores @@ -2869,11 +2870,12 @@ static int ggml_cuda_try_gdn_cache_fusion( if (reserve) { const ggml_tensor * reserve_cpy = nullptr; int reserve_cpy_idx = 0; - if (!find_cpy(cpy_idx + 1, &reserve_cpy, &reserve_cpy_idx) || !match_cpy(reserve_cpy, K - 1, 1)) { + // the leading states and the input state share one contiguous cpy + if (!find_cpy(cpy_idx + 1, &reserve_cpy, &reserve_cpy_idx) || !match_cpy(reserve_cpy, K - 1 - leading, leading + 1)) { return 0; } const ggml_tensor * reserve_dst = reserve_cpy->src[1]; - if ((float *) reserve_dst->data != fused_state_cpy.data + (K - 1) * fused_state_cpy.slot_stride || + if ((float *) reserve_dst->data != fused_state_cpy.data + (K - 1 - leading) * fused_state_cpy.slot_stride || reserve_dst->nb[2] != dst->nb[2]) { return 0; } diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 594dccb5f92d..aab29adb27dc 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2016,6 +2016,7 @@ struct vk_op_gated_delta_net_push_constants { uint32_t trailing_snapshots; int32_t selected_token; uint32_t reserve_input; + uint32_t leading_snapshots; }; static_assert(sizeof(vk_op_gated_delta_net_push_constants) <= 128); @@ -13311,6 +13312,7 @@ static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& s const int32_t trailing_snapshots = ggml_get_op_params_i32(dst, 1); const int32_t selected_token = ggml_get_op_params_i32(dst, 2); const uint32_t reserve_input = ggml_get_op_params_i32(dst, 3) != 0; + const int32_t leading_snapshots = ggml_get_op_params_i32(dst, 4); const vk_op_gated_delta_net_push_constants pc = { H, n_tokens, n_seqs, s_off, sq1, sq2, sq3, @@ -13322,6 +13324,7 @@ static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& s (uint32_t) trailing_snapshots, selected_token, reserve_input, + (uint32_t) leading_snapshots, }; ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/gated_delta_net.comp b/ggml/src/ggml-vulkan/vulkan-shaders/gated_delta_net.comp index 9796a9a050cd..03e18ace0357 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/gated_delta_net.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/gated_delta_net.comp @@ -35,6 +35,7 @@ layout(push_constant) uniform Parameters { uint trailing_snapshots; int selected_token; uint reserve_input; + uint leading_snapshots; }; layout(binding = 0) readonly buffer QBuf { FLOAT_TYPE data_q[]; }; @@ -182,9 +183,13 @@ void main() { if (K > 1u) { const bool selected = selected_token >= 0; - const uint target_slot = selected ? 0u : n_tokens - 1u - t; + uint target_slot = selected ? 0u : n_tokens - 1u - t; const uint snapshot_slots = selected ? 1u : trailing_snapshots; - if ((!selected || t == uint(selected_token)) && target_slot < snapshot_slots) { + const bool leading = !selected && target_slot >= snapshot_slots && t < leading_snapshots; + if (leading) { + target_slot = K - 2u - t; + } + if ((!selected || t == uint(selected_token)) && (target_slot < snapshot_slots || leading)) { const uint slot_base = s_off + uint(target_slot) * state_size_per_snap + state_out_base; [[unroll]] for (uint r = 0; r < ROWS_PER_LANE; r++) { data_dst[slot_base + col * S_V + r * LANES_PER_COLUMN + lane] = s_shard[r]; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index ceb28be462af..c63daf1e0e6a 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6450,6 +6450,7 @@ bool ggml_gated_delta_net_validate(const struct ggml_tensor * tensor) { const int32_t trailing_snapshots = ggml_get_op_params_i32(tensor, 1); const int32_t selected_token = ggml_get_op_params_i32(tensor, 2); const int32_t reserve_input = ggml_get_op_params_i32(tensor, 3); + const int32_t leading_snapshots = ggml_get_op_params_i32(tensor, 4); if (!ggml_are_same_shape(q, k)) { return false; @@ -6494,10 +6495,12 @@ bool ggml_gated_delta_net_validate(const struct ggml_tensor * tensor) { } if (trailing_snapshots < 0 || trailing_snapshots > K || selected_token < -1 || selected_token >= n_tokens || + leading_snapshots < 0 || leading_snapshots > K || (K == 1 && (trailing_snapshots != 1 || selected_token != -1 || reserve_input != 0)) || (selected_token >= 0 && trailing_snapshots > 0) || (reserve_input != 0 && reserve_input != 1) || - (reserve_input == 1 && (selected_token >= 0 || K <= 1 || trailing_snapshots != MIN(n_tokens, K - 1)))) { + (reserve_input == 0 && leading_snapshots != 0) || + (reserve_input == 1 && (selected_token >= 0 || K <= 1 || trailing_snapshots + leading_snapshots != MIN(n_tokens, K - 1)))) { return false; } @@ -6556,7 +6559,7 @@ struct ggml_tensor * ggml_gated_delta_net( result->src[4] = beta; result->src[5] = state; - ggml_gated_delta_net_set_snapshots(result, (int32_t) K, -1, false); + ggml_gated_delta_net_set_snapshots(result, (int32_t) K, -1, false, 0); return result; } @@ -6565,7 +6568,8 @@ void ggml_gated_delta_net_set_snapshots( struct ggml_tensor * tensor, int32_t trailing_snapshots, int32_t selected_token, - bool reserve_input) { + bool reserve_input, + int32_t leading_snapshots) { GGML_ASSERT(tensor != NULL && tensor->op == GGML_OP_GATED_DELTA_NET); const int32_t K = ggml_get_op_params_i32(tensor, 0); @@ -6574,12 +6578,14 @@ void ggml_gated_delta_net_set_snapshots( GGML_ASSERT(selected_token >= -1 && selected_token < n_tokens); GGML_ASSERT(!(selected_token >= 0 && trailing_snapshots > 0)); GGML_ASSERT(!reserve_input || selected_token < 0); - GGML_ASSERT(!reserve_input || (K > 1 && trailing_snapshots == MIN(n_tokens, K - 1))); + GGML_ASSERT(leading_snapshots >= 0 && (reserve_input || leading_snapshots == 0)); + GGML_ASSERT(!reserve_input || (K > 1 && trailing_snapshots + leading_snapshots == MIN(n_tokens, K - 1))); GGML_ASSERT(K != 1 || (trailing_snapshots == 1 && selected_token == -1 && !reserve_input)); ggml_set_op_params_i32(tensor, 1, trailing_snapshots); ggml_set_op_params_i32(tensor, 2, selected_token); ggml_set_op_params_i32(tensor, 3, reserve_input ? 1 : 0); + ggml_set_op_params_i32(tensor, 4, leading_snapshots); } bool ggml_gated_delta_net_has_default_snapshot_params( @@ -6588,7 +6594,8 @@ bool ggml_gated_delta_net_has_default_snapshot_params( return ggml_get_op_params_i32(tensor, 1) == ggml_get_op_params_i32(tensor, 0) && ggml_get_op_params_i32(tensor, 2) == -1 && - ggml_get_op_params_i32(tensor, 3) == 0; + ggml_get_op_params_i32(tensor, 3) == 0 && + ggml_get_op_params_i32(tensor, 4) == 0; } // ggml_lightning_indexer diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 6c4ecfd75b96..4fae4ffec140 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4702,9 +4702,10 @@ struct test_gated_delta_net : public test_case { const int32_t trailing_snapshots; const int32_t selected_token; const bool reserve_input; + const int32_t leading_snapshots; std::string vars() override { - return VARS_TO_STR12(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K, trailing_snapshots, selected_token, reserve_input); + return VARS_TO_STR13(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K, trailing_snapshots, selected_token, reserve_input, leading_snapshots); } std::string op_desc(ggml_tensor * t) override { @@ -4715,10 +4716,10 @@ struct test_gated_delta_net : public test_case { test_gated_delta_net(ggml_type type = GGML_TYPE_F32, int64_t head_count = 4, int64_t head_size = 16, int64_t n_seq_tokens = 1, int64_t n_seqs = 1, int v_repeat = 1, bool permuted = false, bool kda = false, int64_t K = 1, - int32_t trailing_snapshots = -1, int32_t selected_token = -1, bool reserve_input = false) + int32_t trailing_snapshots = -1, int32_t selected_token = -1, bool reserve_input = false, int32_t leading_snapshots = 0) : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), v_repeat(v_repeat), permuted(permuted), kda(kda), K(K), trailing_snapshots(trailing_snapshots), - selected_token(selected_token), reserve_input(reserve_input) {} + selected_token(selected_token), reserve_input(reserve_input), leading_snapshots(leading_snapshots) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * q; @@ -4749,7 +4750,7 @@ struct test_gated_delta_net : public test_case { k = ggml_l2_norm(ctx, k, 1e-6f); ggml_tensor * out = ggml_gated_delta_net(ctx, q, k, v, g, beta, state, K); if (trailing_snapshots >= 0) { - ggml_gated_delta_net_set_snapshots(out, trailing_snapshots, selected_token, reserve_input); + ggml_gated_delta_net_set_snapshots(out, trailing_snapshots, selected_token, reserve_input, leading_snapshots); if (mode == MODE_SUPPORT) { return out; } @@ -4846,11 +4847,12 @@ static bool test_gated_delta_net_rejection() { ggml_backend_dev_t cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); GGML_ASSERT(cpu != nullptr); - const auto set_params = [&](int32_t k, int32_t trailing, int32_t selected, int32_t reserve) { + const auto set_params = [&](int32_t k, int32_t trailing, int32_t selected, int32_t reserve, int32_t leading) { memcpy(out->op_params + 0, &k, sizeof(k)); memcpy(out->op_params + 1, &trailing, sizeof(trailing)); memcpy(out->op_params + 2, &selected, sizeof(selected)); memcpy(out->op_params + 3, &reserve, sizeof(reserve)); + memcpy(out->op_params + 4, &leading, sizeof(leading)); }; const auto rejects_op = [&](ggml_tensor * op, const char * name) { if (ggml_backend_dev_supports_op(cpu, op)) { @@ -4891,20 +4893,24 @@ static bool test_gated_delta_net_rejection() { int32_t trailing; int32_t selected; int32_t reserve; + int32_t leading; const char * name; }; const malformed_params malformed[] = { - { 0, 0, -1, 0, "K" }, - { 3, 3, -1, 0, "K geometry" }, - { 4, 5, -1, 0, "trailing snapshots" }, - { 4, 0, 4, 0, "selected token" }, - { 4, 1, 1, 0, "selected and trailing snapshots" }, - { 4, 3, -1, 2, "reserve input value" }, - { 4, 4, -1, 1, "reserve input geometry" }, - { 4, 2, -1, 1, "reserve input trailing count" }, + { 0, 0, -1, 0, 0, "K" }, + { 3, 3, -1, 0, 0, "K geometry" }, + { 4, 5, -1, 0, 0, "trailing snapshots" }, + { 4, 0, 4, 0, 0, "selected token" }, + { 4, 1, 1, 0, 0, "selected and trailing snapshots" }, + { 4, 3, -1, 2, 0, "reserve input value" }, + { 4, 4, -1, 1, 0, "reserve input geometry" }, + { 4, 2, -1, 1, 0, "reserve input trailing count" }, + { 4, 3, -1, 0, 1, "leading snapshots without reserved input" }, + { 4, 2, -1, 1, 2, "reserve input leading count" }, + { 4, 1, -1, 1, -1, "negative leading snapshots" }, }; for (const auto & test_params : malformed) { - set_params(test_params.k, test_params.trailing, test_params.selected, test_params.reserve); + set_params(test_params.k, test_params.trailing, test_params.selected, test_params.reserve, test_params.leading); if (!rejects(test_params.name)) { return false; } @@ -4923,9 +4929,9 @@ static bool test_gated_delta_net_rejection() { memcpy(k1_out->op_params + 3, &reserve, sizeof(reserve)); }; const malformed_params malformed_k1[] = { - { 1, 0, -1, 0, "K=1 trailing snapshots" }, - { 1, 0, 0, 0, "K=1 selected token" }, - { 1, 1, -1, 1, "K=1 reserve input" }, + { 1, 0, -1, 0, 0, "K=1 trailing snapshots" }, + { 1, 0, 0, 0, 0, "K=1 selected token" }, + { 1, 1, -1, 1, 0, "K=1 reserve input" }, }; for (const auto & test_params : malformed_k1) { set_k1_params(test_params.trailing, test_params.selected, test_params.reserve); @@ -4934,7 +4940,7 @@ static bool test_gated_delta_net_rejection() { } } - set_params(4, 4, -1, 0); + set_params(4, 4, -1, 0, 0); if (!rejects_change(out->ne[1], out->ne[1] + 1, "output geometry")) { return false; } @@ -5196,7 +5202,7 @@ static int test_gated_delta_net_rpc_client(const char * endpoint, int marker_fd) } ggml_tensor probe = *out; for (int mode = 0; mode < 3; ++mode) { - ggml_gated_delta_net_set_snapshots(&probe, mode == 2 ? 0 : 3, mode == 2 ? 0 : -1, mode == 1); + ggml_gated_delta_net_set_snapshots(&probe, mode == 2 ? 0 : 3, mode == 2 ? 0 : -1, mode == 1, 0); if (ggml_backend_supports_op(backend.get(), &probe)) { fprintf(stderr, "RPC GDN wire test: unsupported snapshot mode %d was accepted\n", mode); return 9; @@ -5314,12 +5320,13 @@ struct test_gated_delta_net_fused_cache : public test_case { const int32_t trailing_snapshots; const int32_t selected_token; const bool reserve_input; + const int32_t leading_snapshots; ggml_tensor * state_cpy = nullptr; ggml_tensor * reserve_state_cpy = nullptr; std::string vars() override { - return VARS_TO_STR8(head_count, head_size, n_seq_tokens, n_seqs, K, trailing_snapshots, selected_token, reserve_input); + return VARS_TO_STR9(head_count, head_size, n_seq_tokens, n_seqs, K, trailing_snapshots, selected_token, reserve_input, leading_snapshots); } std::string op_desc(ggml_tensor * t) override { @@ -5345,10 +5352,10 @@ struct test_gated_delta_net_fused_cache : public test_case { test_gated_delta_net_fused_cache( int64_t head_count = 4, int64_t head_size = 16, int64_t n_seq_tokens = 4, int64_t n_seqs = 1, int64_t K = 3, - int32_t trailing_snapshots = -1, int32_t selected_token = -1, bool reserve_input = false) + int32_t trailing_snapshots = -1, int32_t selected_token = -1, bool reserve_input = false, int32_t leading_snapshots = 0) : head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), K(K), trailing_snapshots(trailing_snapshots), - selected_token(selected_token), reserve_input(reserve_input) { + selected_token(selected_token), reserve_input(reserve_input), leading_snapshots(leading_snapshots) { GGML_ASSERT(K > 0); } @@ -5371,7 +5378,7 @@ struct test_gated_delta_net_fused_cache : public test_case { ggml_tensor * gdn_out = ggml_gated_delta_net(ctx, q, k, v, g, beta, state, K); if (trailing_snapshots >= 0) { - ggml_gated_delta_net_set_snapshots(gdn_out, trailing_snapshots, selected_token, reserve_input); + ggml_gated_delta_net_set_snapshots(gdn_out, trailing_snapshots, selected_token, reserve_input, leading_snapshots); } const int64_t D = head_size * head_size * head_count; @@ -5398,7 +5405,7 @@ struct test_gated_delta_net_fused_cache : public test_case { const int64_t n_written = trailing_snapshots < 0 ? std::min(n_seq_tokens, K) : selected_token >= 0 ? 1 : trailing_snapshots; state_cpy = copy_snapshots(0, n_written, "state_cpy"); if (reserve_input) { - reserve_state_cpy = copy_snapshots(K - 1, 1, "reserve_state_cpy"); + reserve_state_cpy = copy_snapshots(K - 1 - leading_snapshots, leading_snapshots + 1, "reserve_state_cpy"); } return ggml_view_4d( @@ -11735,8 +11742,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net_fused_cache(4, 16, 1, 2, 3)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, 4, 0, 3, false)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, 4, 3, -1, true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 2, 1, false, false, 5, 1, -1, true, 3)); test_cases.emplace_back(new test_gated_delta_net_fused_cache(4, 32, 8, 1, 4, 0, 3, false)); test_cases.emplace_back(new test_gated_delta_net_fused_cache(4, 32, 8, 1, 4, 3, -1, true)); + test_cases.emplace_back(new test_gated_delta_net_fused_cache(4, 32, 8, 2, 4, 1, -1, true, 2)); test_cases.emplace_back(new test_gated_delta_net_fused_cache(4, 32, 1, 1, 4, 1, -1, true)); test_cases.emplace_back(new test_gated_delta_net_fused_cache(4, 32, 2, 1, 4, 2, -1, true)); From d3b96e1f096b53b7a774ffc16e6d58de76e45dca Mon Sep 17 00:00:00 2001 From: piggidragon Date: Sat, 19 Sep 2026 03:17:07 +0200 Subject: [PATCH 4/4] speculative : keep the early rollback planes for DFlash The capped planes kept the newest states, which suits MTP. DFlash rejects most drafts early, so it replayed 76% of the verifications. --spec-draft-rs-planes-lead sets how many planes keep the states after the first draft tokens. One plane still keeps the newest state, which a fully accepted draft continues from. The default picks the layout per mode: all but one leading plane for DFlash, none for MTP. With capped planes a rollback without replay now always selects a kept plane, instead of falling back to the checkpoint restore. Assisted-by: Claude Opus 5 --- common/arg.cpp | 7 ++++++ common/common.h | 13 ++++++++++ common/speculative.cpp | 14 +++++++++++ docs/speculative.md | 7 +++++- src/llama-context.cpp | 15 ++++++++---- src/llama-ext.h | 2 +- src/llama-graph.h | 8 ++++++- src/llama-memory-hybrid.cpp | 4 ++-- src/llama-memory-hybrid.h | 2 +- src/llama-memory-recurrent.cpp | 14 +++++++---- src/llama-memory-recurrent.h | 2 +- src/llama-memory.h | 2 +- src/models/delta-net-base.cpp | 14 +++++++---- tests/test-arg-parser.cpp | 25 +++++++++++++++++++ tests/test-recurrent-state-rollback.cpp | 14 +++++------ tools/cli/README.md | 1 + tools/server/README.md | 1 + tools/server/server-context.cpp | 32 ++++++++++++++++--------- 18 files changed, 139 insertions(+), 38 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 5ed110aac3f7..b8ba3085441d 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4214,6 +4214,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.rs_planes = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_RS_PLANES")); + add_opt(common_arg( + {"--spec-draft-rs-planes-lead"}, "N", + "capped recurrent planes that keep the states after the first draft tokens; the others keep the newest states (default: -1, auto: spec-draft-rs-planes - 2 for draft-dflash, 0 for draft-mtp)", + [](common_params & params, int value) { + params.speculative.rs_planes_lead = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_RS_PLANES_LEAD")); add_opt(common_arg( {"--spec-draft-ubatch-size", "--ubatch-size-draft", "-ubd"}, "N", "physical maximum batch size for the draft context (default: 0, inherit target ubatch); " diff --git a/common/common.h b/common/common.h index 3604ee9bed78..2f797d43e052 100644 --- a/common/common.h +++ b/common/common.h @@ -386,6 +386,7 @@ struct common_params_speculative { common_params_speculative_ngram_cache ngram_cache; int32_t rs_planes = 0; // total target recurrent planes for draft-mtp or draft-dflash (0 = draft.n_max + 1) + int32_t rs_planes_lead = -1; // capped planes that keep the states after the first draft tokens (-1 = auto) bool has_dft() const { return !draft.mparams.empty(); @@ -395,6 +396,18 @@ struct common_params_speculative { return rs_planes > 0 && int64_t(rs_planes) < int64_t(draft.n_max) + 1; } + int32_t get_rs_planes_lead() const { + if (!is_rs_capped()) { + return 0; + } + if (rs_planes_lead >= 0) { + return rs_planes_lead; + } + // DFlash rejects most drafts early, so keep one plane for full acceptance and the rest for the first tokens + const bool has_dflash = std::find(types.begin(), types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) != types.end(); + return has_dflash ? rs_planes - 2 : 0; + } + bool has_synth() const { return synth_len != -1.0 || !synth_rates.empty(); } diff --git a/common/speculative.cpp b/common/speculative.cpp index d1e952b0477a..a55d0cc4df25 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2348,6 +2348,14 @@ void common_validate_speculative_params( params.draft.n_ubatch, target_ubatch_raw)); } + if (params.rs_planes_lead < -1) { + throw std::invalid_argument("spec-draft-rs-planes-lead must be -1 (auto) or >= 0"); + } + + if (params.rs_planes_lead > 0 && !params.is_rs_capped()) { + throw std::invalid_argument("spec-draft-rs-planes-lead requires spec-draft-rs-planes below spec-draft-n-max + 1"); + } + if (params.rs_planes == 0) { return; } @@ -2377,6 +2385,12 @@ void common_validate_speculative_params( type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; }); + if (params.is_rs_capped() && params.rs_planes_lead > params.rs_planes - 2) { + throw std::invalid_argument(string_format( + "spec-draft-rs-planes-lead must be in [0, %d] for spec-draft-rs-planes=%d, one plane keeps the newest state", + params.rs_planes - 2, params.rs_planes)); + } + if (params.is_rs_capped() && n_recurrent_modes > 1) { throw std::invalid_argument( "spec-draft-rs-planes cannot be combined with another speculative mode that requires recurrent rollback"); diff --git a/docs/speculative.md b/docs/speculative.md index f5668f465f64..44b82c151c36 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -250,6 +250,9 @@ Use exactly one of these options: --spec-draft-rs-planes, --spec-mtp-rs-planes N total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, use spec-draft-n-max + 1) (env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) +--spec-draft-rs-planes-lead N + capped recurrent planes that keep the states after the first draft tokens; the others keep the newest states (default: -1, auto: spec-draft-rs-planes - 2 for draft-dflash, 0 for draft-mtp) + (env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES_LEAD) --spec-draft-p-split, --draft-p-split P speculative decoding split probability (default: 0.10) (env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) @@ -274,7 +277,9 @@ The effective target ubatch is `min(batch-size, ubatch-size)`, or `batch-size` w Every device selected for the recurrent graph operations must support sparse snapshots. The server rejects an unsupported device layout at startup; it does not move the operations to another device or fall back to CPU. -Plane 0 contains the newest trailing state. The following planes contain older trailing states, and plane `K - 1` preserves the state from before verification. If a verification batch has fewer than `K - 1` tokens, the gap between its trailing states and the reserved input plane remains untouched, including the recorded positions. Rollbacks inside the retained trailing range select a plane directly. A deeper partial acceptance restores the reserved pre-verification plane and replays the full verification batch while writing only the selected boundary. +Plane `K - 1` preserves the state from before verification. The other `K - 1` planes are split into `L = --spec-draft-rs-planes-lead` leading planes and `K - 1 - L` trailing planes. Plane 0 contains the newest state and the following trailing planes contain older states. The leading planes contain the states after the first `L` tokens of the verification batch, which are the states needed when the first drafts are rejected. At least one trailing plane is required, because a fully accepted draft continues from the newest state. If a verification batch has at most `K - 1` tokens, every state is kept in the trailing planes and the gap up to the reserved input plane remains untouched, including the recorded positions. + +A rollback to a kept state selects its plane directly. Any other partial acceptance restores the reserved pre-verification plane and replays the full verification batch while writing only the selected boundary. The replay costs a second target decode, so the layout should match where the drafts are usually rejected. MTP usually accepts most of the draft and keeps the default `L = 0`. DFlash usually rejects early, so by default it keeps one trailing plane and uses all other planes for leading states. For example, with `--spec-draft-n-max 7 --spec-draft-rs-planes 4`, DFlash keeps the states after 0, 1 and 7 accepted drafts. A sparse post-decode failure invalidates every slot that participated in the batch because their state and position planes were committed together. Per-slot recovery would require a separate pre-decode copy of the same recurrent state. diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9aaa0f2a5cf6..7181119ad0dd 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -580,12 +580,19 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 ggml_tensor probe = *op; ggml_gated_delta_net_set_snapshots( - &probe, (int32_t) std::min(n_tokens, n_snapshots - 1), -1, true); + &probe, (int32_t) std::min(n_tokens, n_snapshots - 1), -1, true, 0); if (!ggml_backend_dev_supports_op(device, &probe)) { return false; } - ggml_gated_delta_net_set_snapshots(&probe, 0, 0, false); + if (n_tokens > n_snapshots - 1 && n_snapshots > 2) { + ggml_gated_delta_net_set_snapshots(&probe, 1, -1, true, n_snapshots - 2); + if (!ggml_backend_dev_supports_op(device, &probe)) { + return false; + } + } + + ggml_gated_delta_net_set_snapshots(&probe, 0, 0, false, 0); return ggml_backend_dev_supports_op(device, &probe); }; @@ -4256,11 +4263,11 @@ bool llama_recurrent_sparse_snapshots_supported(const llama_context * ctx) { return ctx != nullptr && ctx->recurrent_sparse_snapshots_supported(); } -bool llama_recurrent_set_sparse_snapshot_mode(llama_context * ctx, bool enabled, int32_t selected_token) { +bool llama_recurrent_set_sparse_snapshot_mode(llama_context * ctx, bool enabled, int32_t selected_token, int32_t n_leading) { if (ctx == nullptr || ctx->get_memory() == nullptr || (enabled && !ctx->recurrent_sparse_snapshots_supported())) { return false; } - return ctx->get_memory()->recurrent_set_sparse_snapshot_mode(enabled, selected_token); + return ctx->get_memory()->recurrent_set_sparse_snapshot_mode(enabled, selected_token, n_leading); } const llama_model * llama_get_model(const llama_context * ctx) { diff --git a/src/llama-ext.h b/src/llama-ext.h index a7465db5da60..5de96a5342eb 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -18,7 +18,7 @@ LLAMA_API struct ggml_cgraph * llama_graph_reserve( LLAMA_API bool llama_recurrent_sparse_snapshots_supported(const struct llama_context * ctx); LLAMA_API bool llama_recurrent_set_sparse_snapshot_mode( - struct llama_context * ctx, bool enabled, int32_t selected_token); + struct llama_context * ctx, bool enabled, int32_t selected_token, int32_t n_leading); // Get the default ggml_type for a given ftype. LLAMA_API ggml_type llama_ftype_get_default_type(llama_ftype ftype); diff --git a/src/llama-graph.h b/src/llama-graph.h index 5466352c5f49..3e891c16e6a4 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -37,9 +37,15 @@ class llama_memory_hybrid_iswa_context; struct llama_recurrent_snapshot_mode { bool sparse = false; int32_t selected_token = -1; + int32_t n_leading = 0; + + // leading token states that a sparse verification of n_tokens keeps in K planes + int64_t leading(int64_t n_tokens, int64_t K) const { + return sparse && selected_token < 0 && n_tokens > K - 1 ? n_leading : 0; + } bool operator==(const llama_recurrent_snapshot_mode & other) const { - return sparse == other.sparse && selected_token == other.selected_token; + return sparse == other.sparse && selected_token == other.selected_token && n_leading == other.n_leading; } }; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 9d236fdc4f2d..bde83e547b20 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -149,8 +149,8 @@ bool llama_memory_hybrid::recurrent_sparse_snapshots_supported() const { return mem_recr->recurrent_sparse_snapshots_supported(); } -bool llama_memory_hybrid::recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token) { - return mem_recr->recurrent_set_sparse_snapshot_mode(enabled, selected_token); +bool llama_memory_hybrid::recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token, int32_t n_leading) { + return mem_recr->recurrent_set_sparse_snapshot_mode(enabled, selected_token, n_leading); } void llama_memory_hybrid::clear(bool data) { diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 393e12c44de6..d1fed7e83ebd 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -64,7 +64,7 @@ class llama_memory_hybrid : public llama_memory_i { bool get_can_shift() const override; bool recurrent_sparse_snapshots_supported() const override; - bool recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token) override; + bool recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token, int32_t n_leading) override; bool get_supports_partial_kv() const override { return true; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 543f34be1a97..66ec44f8a9e1 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -518,13 +518,15 @@ bool llama_memory_recurrent::recurrent_sparse_snapshots_supported() const { return n_rs_seq > 0 && !ctxs_bufs.empty() && graph_supports_sparse_snapshots; } -bool llama_memory_recurrent::recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token) { - if (selected_token < -1) { +bool llama_memory_recurrent::recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token, int32_t n_leading) { + // at least one trailing plane must keep the final state + if (selected_token < -1 || n_leading < 0 || (n_leading > 0 && (!enabled || selected_token >= 0 || n_leading + 2 > (int32_t) n_rs_seq + 1))) { return false; } next_snapshot_mode = { /*.sparse =*/ enabled, /*.selected_token =*/ enabled ? selected_token : -1, + /*.n_leading =*/ n_leading, }; return true; } @@ -1460,11 +1462,15 @@ bool llama_memory_recurrent_context::apply() { GGML_ASSERT((uint32_t) snapshot_mode.selected_token < ubatch.n_seq_tokens); plane_positions[0] = ubatch.pos[i + snapshot_mode.selected_token]; } else { - const uint32_t n_written = std::min(ubatch.n_seq_tokens, n_planes - 1); - for (uint32_t plane = 0; plane < n_written; ++plane) { + const uint32_t n_leading = (uint32_t) snapshot_mode.leading(ubatch.n_seq_tokens, n_planes); + const uint32_t n_trailing = std::min(ubatch.n_seq_tokens, n_planes - 1) - n_leading; + for (uint32_t plane = 0; plane < n_trailing; ++plane) { const int32_t token = (int32_t) ubatch.n_seq_tokens - 1 - (int32_t) plane; plane_positions[plane] = ubatch.pos[i + token]; } + for (uint32_t token = 0; token < n_leading; ++token) { + plane_positions[n_planes - 2 - token] = ubatch.pos[i + token]; + } plane_positions[n_planes - 1] = ubatch.pos[i] - 1; } mem->rs_plane_pos_sparse[seq_id] = true; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index c23ed2bb4ed7..8fd678fcaf13 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -61,7 +61,7 @@ class llama_memory_recurrent : public llama_memory_i { bool get_can_shift() const override; bool recurrent_sparse_snapshots_supported() const override; - bool recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token) override; + bool recurrent_set_sparse_snapshot_mode(bool enabled, int32_t selected_token, int32_t n_leading) override; // state write/load diff --git a/src/llama-memory.h b/src/llama-memory.h index 72ec7340d93e..bb5f85e96ff6 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -119,7 +119,7 @@ struct llama_memory_i { virtual bool get_can_shift() const = 0; virtual bool recurrent_sparse_snapshots_supported() const { return false; } - virtual bool recurrent_set_sparse_snapshot_mode(bool, int32_t) { return false; } + virtual bool recurrent_set_sparse_snapshot_mode(bool, int32_t, int32_t) { return false; } virtual bool get_supports_partial_kv() const { return false; diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index 823c96117931..b8df941e7184 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -523,10 +523,14 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( if (sparse && selected >= 0) { copy_state(selected + 1, 0); } else if (sparse) { - const int64_t n_trailing = std::min(n_seq_tokens, K - 1); + const int64_t n_leading = snapshot_mode.leading(n_seq_tokens, K); + const int64_t n_trailing = std::min(n_seq_tokens, K - 1) - n_leading; for (int64_t t = 0; t < n_trailing; ++t) { copy_state(n_seq_tokens - n_trailing + t + 1, n_trailing - 1 - t); } + for (int64_t t = 0; t < n_leading; ++t) { + copy_state(t + 1, K - 2 - t); + } copy_state(0, K - 1); } else { for (int64_t t = 0; t < K; ++t) { @@ -582,9 +586,11 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output. ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K); + const int64_t n_leading = snapshot_mode.leading(n_seq_tokens, K); if (sparse) { GGML_ASSERT(selected < n_seq_tokens); - ggml_gated_delta_net_set_snapshots(gdn_out, selected >= 0 ? 0 : (int32_t) std::min(n_seq_tokens, K - 1), selected, selected < 0); + const int32_t n_trailing = selected >= 0 ? 0 : (int32_t) (std::min(n_seq_tokens, K - 1) - n_leading); + ggml_gated_delta_net_set_snapshots(gdn_out, n_trailing, selected, selected < 0, (int32_t) n_leading); } if (n_seq_tokens > 1) { res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il}); @@ -622,9 +628,9 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( }; if (sparse && selected < 0) { - const int64_t n_trailing = std::min(n_seq_tokens, K - 1); + const int64_t n_trailing = std::min(n_seq_tokens, K - 1) - n_leading; copy_snapshots(0, n_trailing); - copy_snapshots(K - 1, 1); + copy_snapshots(K - 1 - n_leading, n_leading + 1); } else { copy_snapshots(0, sparse ? 1 : std::min(n_seq_tokens, K)); } diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 8651de0a3617..6a7e75fe8208 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -288,6 +288,31 @@ static void test(void) { } } + struct rs_lead_case { + std::vector args; + bool valid; + int32_t lead; + }; + const rs_lead_case rs_lead_cases[] = { + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "7", "--spec-draft-rs-planes", "4" }, true, 2 }, + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "7" }, true, 0 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "7", "--spec-draft-rs-planes", "4" }, true, 0 }, + { { "--spec-type", "draft-mtp", "--spec-draft-n-max", "7", "--spec-draft-rs-planes", "4", "--spec-draft-rs-planes-lead", "2" }, true, 2 }, + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "7", "--spec-draft-rs-planes", "4", "--spec-draft-rs-planes-lead", "0" }, true, 0 }, + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "7", "--spec-draft-rs-planes", "4", "--spec-draft-rs-planes-lead", "3" }, false, -1 }, + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "7", "--spec-draft-rs-planes", "8", "--spec-draft-rs-planes-lead", "1" }, false, -1 }, + { { "--spec-type", "draft-dflash", "--spec-draft-n-max", "7", "--spec-draft-rs-planes", "4", "--spec-draft-rs-planes-lead", "-2" }, false, -1 }, + }; + for (const auto & test_case : rs_lead_cases) { + params = common_params(); + argv = { "binary_name" }; + argv.insert(argv.end(), test_case.args.begin(), test_case.args.end()); + assert(test_case.valid == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SERVER)); + if (test_case.valid) { + assert(params.speculative.get_rs_planes_lead() == test_case.lead); + } + } + params = common_params(); params.model.path = "model_file.gguf"; diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index b7de44d5bb7c..9559e9f367da 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -340,7 +340,7 @@ static bool test_sparse_selected_then_trailing(const common_params & params, lla common_batch_add(split_batch, (11*pos + 38) % n_vocab, pos, { 1 }, pos == 4); } - const bool mode_ok = llama_recurrent_set_sparse_snapshot_mode(ctx_split, true, -1); + const bool mode_ok = llama_recurrent_set_sparse_snapshot_mode(ctx_split, true, -1, 0); const int split_ret = mode_ok ? llama_decode(ctx_split, split_batch) : 0; llama_batch_free(split_batch); if (!mode_ok || split_ret == 0) { @@ -373,9 +373,9 @@ static bool test_sparse_selected_then_trailing(const common_params & params, lla if (ctx_normal == nullptr || ctx_ref == nullptr) { return fail("normal transition context init"); } - if (!llama_recurrent_set_sparse_snapshot_mode(ctx_normal, true, -1) || + if (!llama_recurrent_set_sparse_snapshot_mode(ctx_normal, true, -1, 0) || !decode_ranges(ctx_normal, { { 0, 0 } }, 4) || - !llama_recurrent_set_sparse_snapshot_mode(ctx_normal, false, -1)) { + !llama_recurrent_set_sparse_snapshot_mode(ctx_normal, false, -1, 0)) { return fail("sparse-to-normal setup"); } @@ -406,7 +406,7 @@ static bool test_sparse_selected_then_trailing(const common_params & params, lla fprintf(stderr, "%s : sparse-to-normal rollback matched reference output\n", __func__); } - if (!llama_recurrent_set_sparse_snapshot_mode(ctx, true, -1)) { + if (!llama_recurrent_set_sparse_snapshot_mode(ctx, true, -1, 0)) { return fail("enable trailing mode"); } if (!decode_ranges(ctx, { { 0, 0 }, { 1, 0 } }, 4) || !decode_ranges(ctx, { { 0, 4 }, { 1, 4 } }, 5)) { @@ -416,14 +416,14 @@ static bool test_sparse_selected_then_trailing(const common_params & params, lla const std::pair selected_cases[] = { { 0, 2 }, { 1, 1 } }; for (const auto & [seq_id, selected] : selected_cases) { if (!llama_memory_seq_rm(llama_get_memory(ctx), seq_id, 4, -1) || - !llama_recurrent_set_sparse_snapshot_mode(ctx, true, selected) || + !llama_recurrent_set_sparse_snapshot_mode(ctx, true, selected, 0) || !decode_ranges(ctx, { { (uint32_t) seq_id, 4 } }, 5) || !llama_memory_seq_rm(llama_get_memory(ctx), seq_id, 5 + selected, -1)) { return fail(seq_id == 0 ? "sequence 0 selected replay" : "sequence 1 selected replay"); } } - if (!llama_recurrent_set_sparse_snapshot_mode(ctx, true, -1) || !decode_ranges(ctx, { { 0, 7 }, { 1, 6 } }, 3)) { + if (!llama_recurrent_set_sparse_snapshot_mode(ctx, true, -1, 0) || !decode_ranges(ctx, { { 0, 7 }, { 1, 6 } }, 3)) { return fail("selected-to-trailing co-batch"); } if (!llama_memory_seq_rm(llama_get_memory(ctx), 0, 7, -1) || @@ -481,7 +481,7 @@ static bool test_sparse_short_trailing_gaps(const common_params & params, llama_ return false; } - const bool ok = llama_recurrent_set_sparse_snapshot_mode(ctx_sparse, true, -1) && + const bool ok = llama_recurrent_set_sparse_snapshot_mode(ctx_sparse, true, -1, 0) && decode_range(ctx_sparse, 0, 3) && decode_range(ctx_sparse, 3, trailing) && decode_range(ctx_ref, 0, (uint32_t) rollback_pos + 1) && llama_memory_seq_rm(llama_get_memory(ctx_sparse), 0, rollback_pos + 1, -1) && diff --git a/tools/cli/README.md b/tools/cli/README.md index 2556717f6993..8d9e040963d4 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -201,6 +201,7 @@ | `--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-rs-planes, --spec-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, use spec-draft-n-max + 1); see [capped recurrent planes](../../docs/speculative.md#capped-recurrent-planes)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) | +| `--spec-draft-rs-planes-lead N` | capped recurrent planes that keep the states after the first draft tokens; the others keep the newest states (default: -1, auto: spec-draft-rs-planes - 2 for draft-dflash, 0 for draft-mtp)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES_LEAD) | | `--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 1926fdb7eefd..e91fa2bce0bb 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -262,6 +262,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--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-rs-planes, --spec-mtp-rs-planes N` | total target recurrent-state planes for draft-mtp or draft-dflash, including the current state (default: 0, use spec-draft-n-max + 1); see [capped recurrent planes](../../docs/speculative.md#capped-recurrent-planes)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES) | +| `--spec-draft-rs-planes-lead N` | capped recurrent planes that keep the states after the first draft tokens; the others keep the newest states (default: -1, auto: spec-draft-rs-planes - 2 for draft-dflash, 0 for draft-mtp)
(env: LLAMA_ARG_SPEC_DRAFT_RS_PLANES_LEAD) | | `--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/server-context.cpp b/tools/server/server-context.cpp index 9ae69f19f315..ea2ee97184f7 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1285,9 +1285,10 @@ struct server_context_impl { if (has_spec && llama_n_rs_seq(ctx_tgt) > 0) { const uint32_t total_planes = llama_n_rs_seq(ctx_tgt) + 1; - const uint32_t direct_rollback = capped_rs ? total_planes - 2 : total_planes - 1; - SRV_INF("recurrent-plane policy: draft depth = %d, total planes = %u, direct rollback horizon = %u, full-shape GPU replay = %s\n", - params_base.speculative.draft.n_max, total_planes, direct_rollback, + const int32_t n_lead = params_base.speculative.get_rs_planes_lead(); + const uint32_t n_trail = capped_rs ? total_planes - 1 - n_lead : total_planes - 1; + SRV_INF("recurrent-plane policy: draft depth = %d, total planes = %u, leading planes = %d, trailing planes = %u, full-shape GPU replay = %s\n", + params_base.speculative.draft.n_max, total_planes, n_lead, n_trail, capped_rs ? "enabled" : "disabled"); } @@ -3894,7 +3895,8 @@ struct server_context_impl { bool snapshot_mode_restored = true; queue_tasks.yield_to_queue([&]() { bool snapshot_mode_enabled = false; - if (sparse_snapshots && !llama_recurrent_set_sparse_snapshot_mode(ctx_tgt, true, selected_token)) { + const int32_t n_lead = selected_token < 0 ? params_base.speculative.get_rs_planes_lead() : 0; + if (sparse_snapshots && !llama_recurrent_set_sparse_snapshot_mode(ctx_tgt, true, selected_token, n_lead)) { snapshot_mode_restored = false; return; } @@ -3906,12 +3908,12 @@ struct server_context_impl { } } catch (...) { if (snapshot_mode_enabled) { - snapshot_mode_restored = llama_recurrent_set_sparse_snapshot_mode(ctx_tgt, false, -1); + snapshot_mode_restored = llama_recurrent_set_sparse_snapshot_mode(ctx_tgt, false, -1, 0); } throw; } if (snapshot_mode_enabled) { - snapshot_mode_restored = llama_recurrent_set_sparse_snapshot_mode(ctx_tgt, false, -1); + snapshot_mode_restored = llama_recurrent_set_sparse_snapshot_mode(ctx_tgt, false, -1, 0); } }); @@ -4187,8 +4189,15 @@ struct server_context_impl { GGML_ASSERT(accepted.size() >= 1); const uint32_t n_rollback = slot.spec_draft.size() + 1 - accepted.size(); - const uint32_t direct_horizon = llama_n_rs_seq(ctx_tgt) > 0 ? llama_n_rs_seq(ctx_tgt) - 1 : 0; - const bool use_gpu_replay = slot.spec_replay.gpu_snapshots_armed() && n_rollback > direct_horizon; + bool use_gpu_replay = false; + if (slot.spec_replay.gpu_snapshots_armed()) { + // same plane layout as the sparse snapshot graph: leading states, trailing states, input state + const uint32_t n_planes = llama_n_rs_seq(ctx_tgt) + 1; + const uint32_t n_tokens = slot.spec_draft.size() + 1; + const uint32_t n_lead = n_tokens > n_planes - 1 ? params_base.speculative.get_rs_planes_lead() : 0; + const uint32_t n_trail = std::min(n_tokens, n_planes - 1) - n_lead; + use_gpu_replay = n_rollback >= n_trail && accepted.size() - 1 >= n_lead; + } if (use_gpu_replay) { const uint32_t n_accepted = accepted.size() - 1; @@ -4206,9 +4215,10 @@ struct server_context_impl { return; } - const bool use_ckpt_tgt = - ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || - (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && n_rollback > llama_n_rs_seq(ctx_tgt)); + // with capped planes, a rollback without replay selects a kept plane + const bool use_ckpt_tgt = !slot.spec_replay.gpu_snapshots_armed() && + (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || + (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && n_rollback > llama_n_rs_seq(ctx_tgt))); // check for partial draft acceptance if (n_rollback > 0) {