# Applies to llama.cpp commit a3b1eff # Single cumulative patch: Eaman ROCm/MTP/VEC/MoE work plus MTP recurrent-state rollback # Includes Eaman cumulative ROCm/MTP/VEC work through commit 5850e07 # Includes experimental MTP recurrent-state rollback work through commit 732dd501 # Tested with llama.cpp version 1184 # Tested on ROCm 7.14.60850-0000000 / Vulkan # Experimental: --spec-mtp-rs-depth remains opt-in and isolated from production eaman diff --git a/common/arg.cpp b/common/arg.cpp index 6f5fe377d5c4f0e7da5f76a4947f77b07b1d964a..f231a4473fd4e93a2d91eafae1831d12bf4d1930 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1292,6 +1292,9 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e ctx_arg.params = params_org; return false; } + if (ctx_arg.params.speculative.draft.n_rs_seq > ctx_arg.params.speculative.draft.n_max) { + throw std::invalid_argument("--spec-mtp-rs-depth must not exceed --spec-draft-n-max"); + } if (ctx_arg.params.usage) { common_params_print_usage(ctx_arg); if (ctx_arg.print_usage) { @@ -1756,6 +1759,21 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str())); } }).set_env("LLAMA_ARG_FLASH_ATTN")); + add_opt(common_arg({ "--pipeline-parallel" }, "[on|off|auto]", + string_format("set pipeline parallelism ('on', 'off', or 'auto', default: '%s')", + llama_pipeline_parallel_type_name(params.pipeline_parallel_type)), + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_ENABLED; + } else if (is_falsey(value)) { + params.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED; + } else if (is_autoy(value)) { + params.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_AUTO; + } else { + throw std::runtime_error( + string_format("error: unknown value for --pipeline-parallel: '%s'\n", value.c_str())); + } + }).set_env("LLAMA_ARG_PIPELINE_PARALLEL")); add_opt(common_arg( {"-p", "--prompt"}, "PROMPT", "prompt to start generation with; for system message, use -sys", @@ -4083,6 +4101,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.n_max = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX")); + add_opt(common_arg( + {"--spec-mtp-rs-depth"}, "N", + "number of target recurrent-state rollback snapshots for MTP; lower values save memory but replay accepted tokens after deep rejection (default: --spec-draft-n-max)", + [](common_params & params, int value) { + if (value < 1) { + throw std::invalid_argument("--spec-mtp-rs-depth must be at least 1"); + } + params.speculative.draft.n_rs_seq = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_MTP_RS_DEPTH")); add_opt(common_arg( {"--spec-draft-n-min"}, "N", string_format("minimum number of draft tokens to use for speculative decoding (default: %d)", params.speculative.draft.n_min), diff --git a/common/common.cpp b/common/common.cpp index 25ca838dff0ad0a20800b50bdf452b925d9abb36..68fc14e12db9957b5099a5797293c1107e378362 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1716,6 +1716,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.pooling_type = params.pooling_type; cparams.attention_type = params.attention_type; cparams.flash_attn_type = params.flash_attn_type; + cparams.pipeline_parallel_type = params.pipeline_parallel_type; cparams.cb_eval = params.cb_eval; cparams.cb_eval_user_data = params.cb_eval_user_data; cparams.offload_kqv = !params.no_kv_offload; @@ -2244,6 +2245,8 @@ void common_prompt_checkpoint::clear() { data_tgt.clear(); data_dft.clear(); data_spec.clear(); + data_tgt_on_device = false; + data_dft_on_device = false; } void common_prompt_checkpoint::update_pos( @@ -2255,6 +2258,50 @@ void common_prompt_checkpoint::update_pos( this->pos_max = pos_max; } +static void common_prompt_checkpoint_save( + std::vector & data, + bool & on_device, + llama_context * ctx, + llama_seq_id seq_id, + llama_state_seq_flags flags, + const char * label) { + const bool requested_on_device = flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; + on_device = false; + + auto save = [&](llama_state_seq_flags save_flags) { + const size_t ckpt_size = llama_state_seq_get_size_ext(ctx, seq_id, save_flags); + if (ckpt_size == 0) { + return false; + } + + std::vector saved(ckpt_size); + const size_t n = llama_state_seq_get_data_ext(ctx, saved.data(), ckpt_size, seq_id, save_flags); + if (n != ckpt_size) { + return false; + } + + data.swap(saved); + return true; + }; + + bool saved = save(flags); + if (requested_on_device && !saved) { + COM_WRN("%s: ON_DEVICE %s checkpoint save failed; retrying with host storage\n", + __func__, label); + flags &= ~LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; + saved = save(flags); + } + + if (!saved) { + if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { + GGML_ABORT("checkpoint device save failed for %s\n", label); + } + GGML_ABORT("checkpoint size mismatch while saving %s\n", label); + } + + on_device = flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; +} + void common_prompt_checkpoint::update_tgt( llama_context * ctx, llama_seq_id seq_id, @@ -2263,14 +2310,7 @@ void common_prompt_checkpoint::update_tgt( return; } - const size_t ckpt_size = llama_state_seq_get_size_ext(ctx, seq_id, flags); - - data_tgt.resize(ckpt_size); - - const size_t n = llama_state_seq_get_data_ext(ctx, data_tgt.data(), ckpt_size, seq_id, flags); - if (n != ckpt_size) { - GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", ckpt_size, n); - } + common_prompt_checkpoint_save(data_tgt, data_tgt_on_device, ctx, seq_id, flags, "target"); } void common_prompt_checkpoint::update_dft( @@ -2281,14 +2321,7 @@ void common_prompt_checkpoint::update_dft( return; } - const size_t ckpt_size = llama_state_seq_get_size_ext(ctx, seq_id, flags); - - data_dft.resize(ckpt_size); - - const size_t n = llama_state_seq_get_data_ext(ctx, data_dft.data(), ckpt_size, seq_id, flags); - if (n != ckpt_size) { - GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", ckpt_size, n); - } + common_prompt_checkpoint_save(data_dft, data_dft_on_device, ctx, seq_id, flags, "draft"); } void common_prompt_checkpoint::load_tgt( @@ -2303,6 +2336,9 @@ void common_prompt_checkpoint::load_tgt( return; } + flags = (flags & ~LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) | + (data_tgt_on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : 0); + const size_t n = llama_state_seq_set_data_ext(ctx, data_tgt.data(), data_tgt.size(), seq_id, flags); if (n != data_tgt.size()) { GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", data_tgt.size(), n); @@ -2321,6 +2357,9 @@ void common_prompt_checkpoint::load_dft( return; } + flags = (flags & ~LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) | + (data_dft_on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : 0); + const size_t n = llama_state_seq_set_data_ext(ctx, data_dft.data(), data_dft.size(), seq_id, flags); if (n != data_dft.size()) { GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", data_dft.size(), n); @@ -2329,9 +2368,11 @@ void common_prompt_checkpoint::load_dft( void common_prompt_checkpoint::clear_tgt() { data_tgt.clear(); + data_tgt_on_device = false; } void common_prompt_checkpoint::clear_dft() { data_dft.clear(); data_spec.clear(); + data_dft_on_device = false; } diff --git a/common/common.h b/common/common.h index d8a16897b8405a6ac5ff5602a7abb47ec2d239a7..966953dc669b1035ae78311cfbd85734013b96c2 100644 --- a/common/common.h +++ b/common/common.h @@ -324,6 +324,7 @@ struct common_params_model { struct common_params_speculative_draft { int32_t n_max = 3; // maximum number of tokens to draft during speculative decoding int32_t n_min = 0; // minimum number of draft tokens to use for speculative decoding + int32_t n_rs_seq = -1; // MTP recurrent-state rollback depth (-1 = n_max) float p_split = 0.1f; // speculative decoding split probability float p_min = 0.0f; // minimum speculative decoding probability (greedy) @@ -384,11 +385,18 @@ struct common_params_speculative { } uint32_t need_n_rs_seq() const { - bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { - 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 needs_mtp = std::find(types.begin(), types.end(), COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != types.end(); + const bool needs_other_rs = std::any_of(types.begin(), types.end(), [&](auto t) { + return t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; }); - return needs_rs_seq ? draft.n_max : 0u; + if (needs_other_rs) { + return draft.n_max; + } + if (needs_mtp) { + return draft.n_rs_seq >= 0 ? draft.n_rs_seq : draft.n_max; + } + return 0u; } }; @@ -487,6 +495,7 @@ struct common_params { enum llama_pooling_type pooling_type = LLAMA_POOLING_TYPE_UNSPECIFIED; // pooling type for embeddings enum llama_attention_type attention_type = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings enum llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention + enum llama_pipeline_parallel_type pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_AUTO; // whether to use pipeline parallelism struct common_params_sampling sampling; struct common_params_speculative speculative; @@ -1145,6 +1154,11 @@ struct common_prompt_checkpoint { std::vector data_tgt; std::vector data_dft; + // The actual storage mode used for each checkpoint. This can differ from + // the flags passed to update_* when an on-device save is not possible. + bool data_tgt_on_device = false; + bool data_dft_on_device = false; + // (optional) speculative-decoding implementation state stashed with the checkpoint // (e.g. eagle3's deferred-boundary g_embd row) std::vector data_spec; diff --git a/common/fit.cpp b/common/fit.cpp index dd1f3ef76619cb5f56f657ab7583884a6d114006..b8bc1dfef4a3961c0ce6f4c387be31269510a1c6 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -185,6 +185,7 @@ static void common_params_fit_impl( constexpr int64_t MiB = 1024*1024; typedef std::vector dmds_t; const llama_model_params default_mparams = llama_model_default_params(); + const bool fixed_gpu_layers = mparams->n_gpu_layers != default_mparams.n_gpu_layers; std::vector devs; uint32_t hp_ngl = 0; // hparams.n_gpu_layers @@ -317,7 +318,7 @@ static void common_params_fit_impl( sum_used_target -= margins[id]; } } - if (nd > 1) { + if (nd > 1 && !fixed_gpu_layers) { // for multiple devices we need to be more conservative in terms of how much context we think can fit: // - for dense models only whole layers can be assigned to devices // - for MoE models only whole tensors can be assigned to devices, which we estimate to be <= 1/3 of a layer @@ -337,6 +338,30 @@ static void common_params_fit_impl( sum_projected_used_min_ctx += dmds_min_ctx[id].mb.total(); } } + if (nd > 1 && fixed_gpu_layers) { + uint32_t n_ctx_fit = hp_nct; + bool fits = true; + for (size_t id = 0; id < nd; id++) { + const int64_t used_full = dmds_full[id].mb.total(); + const int64_t used_min = dmds_min_ctx[id].mb.total(); + const int64_t target = dmds_full[id].free - margins[id]; + if (used_full <= target) { + continue; + } + uint32_t n_ctx_device = n_ctx_min; + fits = fits && target >= used_min; + if (target > used_min && used_full > used_min) { + n_ctx_device += (hp_nct - n_ctx_min) * (target - used_min) / (used_full - used_min); + } + n_ctx_fit = std::min(n_ctx_fit, n_ctx_device); + } + cparams->n_ctx = std::max(n_ctx_fit - n_ctx_fit % 256, n_ctx_min); + LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " using per-device limits\n", + __func__, hp_nct, cparams->n_ctx); + if (fits) { + return; + } + } if (sum_used_target > sum_projected_used_min_ctx) { // linear interpolation between minimum and maximum context size: cparams->n_ctx += (hp_nct - n_ctx_min) * (sum_used_target - sum_projected_used_min_ctx) @@ -374,7 +399,7 @@ static void common_params_fit_impl( throw common_params_fit_exception("was unable to fit model into system memory by reducing context, abort"); } - if (mparams->n_gpu_layers != default_mparams.n_gpu_layers) { + if (fixed_gpu_layers) { throw common_params_fit_exception("n_gpu_layers already set by user to " + std::to_string(mparams->n_gpu_layers) + ", abort"); } if (nd > 1) { diff --git a/common/speculative.cpp b/common/speculative.cpp index 89e9b2782c2ac470c6b2401769148bccee24b037..390b2c7c916e6694d12dfb4c8faaca15ead36b4d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2383,6 +2383,10 @@ common_speculative_init_result::common_speculative_init_result( if (spec_mtp) { cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED; + // use the same (possibly fit-reduced) context size as the target + // avoids MTP draft reserving flash attention workspace for full n_ctx + cparams.n_ctx = llama_n_ctx(ctx_tgt); } // note: for small models maybe we can set this to the maximum possible draft from all speculative types diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ab7a3b297c07c43dc2902de425b292167e1f8d11..980f21eb1d1e310d15ad28e8b122cd84e145211b 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -457,6 +457,14 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path. const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0; +#ifdef GGML_USE_HIP + // HIP quantized-KV TILE/MMA paths materialize large F16 temporary buffers; + // VEC dequantizes in-register and is also the safe path on RDNA2. + if ((ggml_is_quantized(K->type) || ggml_is_quantized(V->type)) && can_use_vector_kernel) { + return BEST_FATTN_KERNEL_VEC; + } +#endif + // If Turing tensor cores are available, use them: if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { if (can_use_vector_kernel) { diff --git a/include/llama.h b/include/llama.h index 177fc10a9139e7104f07decb8d4d5ded94279b06..4b1e818ce8270e3c755d19858db599c7cb55f19d 100644 --- a/include/llama.h +++ b/include/llama.h @@ -195,6 +195,14 @@ extern "C" { LLAMA_API const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_type); + enum llama_pipeline_parallel_type { + LLAMA_PIPELINE_PARALLEL_TYPE_AUTO = -1, + LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED = 0, + LLAMA_PIPELINE_PARALLEL_TYPE_ENABLED = 1, + }; + + LLAMA_API const char * llama_pipeline_parallel_type_name(enum llama_pipeline_parallel_type pipeline_parallel_type); + enum llama_split_mode { LLAMA_SPLIT_MODE_NONE = 0, // single GPU LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs @@ -364,6 +372,7 @@ extern "C" { enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id enum llama_attention_type attention_type; // attention type to use for embeddings enum llama_flash_attn_type flash_attn_type; // when to enable Flash Attention + enum llama_pipeline_parallel_type pipeline_parallel_type; // when to enable pipeline parallelism // ref: https://github.com/ggml-org/llama.cpp/pull/2054 float rope_freq_base; // RoPE base frequency, 0 = from model diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 52f8d53672a3b2ab76886d8febd358b4068315ab..8b9d6357dd4973fc2cc9eba1f33b4aa5cea8c7a4 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -310,6 +310,7 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); + LLAMA_LOG_INFO("%s: pipeline mode = %s\n", __func__, llama_pipeline_parallel_type_name(params.pipeline_parallel_type)); LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); @@ -426,6 +427,8 @@ llama_context::llama_context( // TODO: move these checks to ggml_backend_sched // enabling pipeline parallelism in the scheduler increases memory usage, so it is only done when necessary bool pipeline_parallel = + cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP && + params.pipeline_parallel_type != LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED && model.n_devices() > 1 && model.n_gpu_layers() > model.hparams.n_layer_all && model.split_mode() == LLAMA_SPLIT_MODE_LAYER && @@ -452,6 +455,11 @@ llama_context::llama_context( } } + if (cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP && + params.pipeline_parallel_type == LLAMA_PIPELINE_PARALLEL_TYPE_ENABLED && !pipeline_parallel) { + LLAMA_LOG_WARN("%s: pipeline parallelism requested but unavailable\n", __func__); + } + cparams.pipeline_parallel = pipeline_parallel; if (cparams.pipeline_parallel) { @@ -3526,6 +3534,7 @@ llama_context_params llama_context_default_params() { /*.pooling_type =*/ LLAMA_POOLING_TYPE_UNSPECIFIED, /*.attention_type =*/ LLAMA_ATTENTION_TYPE_UNSPECIFIED, /*.flash_attn_type =*/ LLAMA_FLASH_ATTN_TYPE_AUTO, + /*.pipeline_parallel_type =*/ LLAMA_PIPELINE_PARALLEL_TYPE_AUTO, /*.rope_freq_base =*/ 0.0f, /*.rope_freq_scale =*/ 0.0f, /*.yarn_ext_factor =*/ -1.0f, diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index ef82eb976ca76f1b999e225118320e7f1e379fcb..f7638544759a03f4c4d8cb682927e3ac0f23733b 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -788,7 +788,7 @@ void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq } if ((flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) && cell_ranges.size() > 1) { - GGML_ABORT("cannot save/load multiple ranges of cells to/from device memory\n"); + throw std::runtime_error("cannot save/load multiple ranges of cells to/from device memory"); } // DEBUG CHECK: Sum of cell counts in ranges should equal the total cell count diff --git a/src/llama.cpp b/src/llama.cpp index 1609fec88ddfe3032e677590794bbadd1c93a2f7..6db5163f350575343e4112de291b63c0eb886a55 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -46,6 +46,18 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty GGML_ABORT("fatal error"); } +const char * llama_pipeline_parallel_type_name(enum llama_pipeline_parallel_type pipeline_parallel_type) { + switch (pipeline_parallel_type) { + case LLAMA_PIPELINE_PARALLEL_TYPE_AUTO: + return "auto"; + case LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED: + return "disabled"; + case LLAMA_PIPELINE_PARALLEL_TYPE_ENABLED: + return "enabled"; + } + GGML_ABORT("fatal error"); +} + const char * llama_load_mode_name(enum llama_load_mode load_mode) { switch (load_mode) { case LLAMA_LOAD_MODE_AUTO: @@ -614,4 +626,3 @@ const char * llama_print_system_info(void) { return s.c_str(); } - diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index ba58f852eb4f772ffad61ab6925738b331709aea..3fad40d422320dabbc859ff3713f34f840eda918 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -3,6 +3,7 @@ #include "download.h" #include "llama.h" #include "speculative.h" +#include "ggml-backend.h" #include #include @@ -197,6 +198,132 @@ static void test(void) { assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SPECULATIVE)); assert(params.speculative.draft.n_max == 123); + { + common_params params_mtp; + argv = {"binary_name", "-m", "abc.gguf", "--spec-draft-n-max", "5", "--spec-mtp-rs-depth", "1"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params_mtp, LLAMA_EXAMPLE_SPECULATIVE)); + assert(params_mtp.speculative.draft.n_max == 5); + assert(params_mtp.speculative.draft.n_rs_seq == 1); + + params_mtp.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP }; + assert(params_mtp.speculative.need_n_rs_seq() == 1); + + params_mtp.speculative.types.push_back(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3); + assert(params_mtp.speculative.need_n_rs_seq() == 5); + } + + { + common_params params_mtp; + argv = {"binary_name", "-m", "abc.gguf", "--spec-draft-n-max", "5"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params_mtp, LLAMA_EXAMPLE_SPECULATIVE)); + assert(params_mtp.speculative.draft.n_rs_seq == -1); + + params_mtp.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP }; + assert(params_mtp.speculative.need_n_rs_seq() == 5); + } + + { + common_params params_mtp; + argv = {"binary_name", "-m", "abc.gguf", "--spec-mtp-rs-depth", "1", "--spec-draft-n-max", "5"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params_mtp, LLAMA_EXAMPLE_SPECULATIVE)); + assert(params_mtp.speculative.draft.n_rs_seq == 1); + } + + { + common_params params_mtp; + argv = {"binary_name", "-m", "abc.gguf", "--spec-draft-n-max", "5", "--spec-mtp-rs-depth", "0"}; + assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), params_mtp, LLAMA_EXAMPLE_SPECULATIVE)); + } + + { + common_params params_mtp; + argv = {"binary_name", "-m", "abc.gguf", "--spec-draft-n-max", "5", "--spec-mtp-rs-depth=-1"}; + assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), params_mtp, LLAMA_EXAMPLE_SPECULATIVE)); + } + + { + common_params params_mtp; + argv = {"binary_name", "-m", "abc.gguf", "--spec-mtp-rs-depth", "6", "--spec-draft-n-max", "5"}; + assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), params_mtp, LLAMA_EXAMPLE_SPECULATIVE)); + } + + { + printf("test-arg-parser: test recurrent checkpoint reservation inputs\n\n"); + + common_params fit_params; + assert(fit_params.fit_params); + assert(fit_params.fit_params_target.size() == llama_max_devices()); + for (const size_t target : fit_params.fit_params_target) { + assert(target == 1024 * 1024 * 1024ULL); + } + + argv = {"binary_name", "-m", "abc.gguf", "--fit", "off", "--split-mode", "none", "--main-gpu", "1", + "--fit-target", "256,64", "--tensor-split", "3,1"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), fit_params, LLAMA_EXAMPLE_COMMON)); + assert(!fit_params.fit_params); + assert(fit_params.split_mode == LLAMA_SPLIT_MODE_NONE); + assert(fit_params.main_gpu == 1); + assert(fit_params.fit_params_target[0] == 256 * 1024 * 1024ULL); + assert(fit_params.fit_params_target[1] == 64 * 1024 * 1024ULL); + assert(fit_params.tensor_split[0] == 3.0f); + assert(fit_params.tensor_split[1] == 1.0f); + + auto checkpoint_reservation_needed = [](const common_params & p) { + return p.speculative.need_n_rs_seq() < static_cast(p.speculative.draft.n_max); + }; + + common_params mtp_full; + argv = {"binary_name", "-m", "abc.gguf", "--spec-draft-n-max", "5"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), mtp_full, LLAMA_EXAMPLE_SPECULATIVE)); + mtp_full.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP }; + assert(mtp_full.speculative.need_n_rs_seq() == 5); + assert(!checkpoint_reservation_needed(mtp_full)); + + common_params mtp_reduced; + argv = {"binary_name", "-m", "abc.gguf", "--spec-draft-n-max", "5", "--spec-mtp-rs-depth", "2"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), mtp_reduced, LLAMA_EXAMPLE_SPECULATIVE)); + mtp_reduced.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP }; + assert(mtp_reduced.speculative.need_n_rs_seq() == 2); + assert(checkpoint_reservation_needed(mtp_reduced)); + + common_params mtp_with_other_rs; + argv = {"binary_name", "-m", "abc.gguf", "--spec-draft-n-max", "5", "--spec-mtp-rs-depth", "2"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), mtp_with_other_rs, LLAMA_EXAMPLE_SPECULATIVE)); + mtp_with_other_rs.speculative.types = { + COMMON_SPECULATIVE_TYPE_DRAFT_MTP, + COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, + }; + assert(mtp_with_other_rs.speculative.need_n_rs_seq() == 5); + assert(!checkpoint_reservation_needed(mtp_with_other_rs)); + + ggml_backend_load_all(); + std::vector gpu_devices; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { + gpu_devices.push_back(dev); + } + } + if (gpu_devices.size() >= 2) { + const std::string device_arg = string_format("%s,%s", + ggml_backend_dev_name(gpu_devices[0]), ggml_backend_dev_name(gpu_devices[1])); + common_params multi_gpu; + argv = {"binary_name", "-m", "abc.gguf", "--device", device_arg, "--split-mode", "layer", + "--fit-target", "128,256", "--tensor-split", "3,1"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), multi_gpu, LLAMA_EXAMPLE_COMMON)); + assert(multi_gpu.devices.size() == 3); + assert(multi_gpu.devices[0] == gpu_devices[0]); + assert(multi_gpu.devices[1] == gpu_devices[1]); + assert(multi_gpu.devices[2] == nullptr); + assert(multi_gpu.fit_params_target[0] == 128 * 1024 * 1024ULL); + assert(multi_gpu.fit_params_target[1] == 256 * 1024 * 1024ULL); + assert(multi_gpu.tensor_split[0] == 3.0f); + assert(multi_gpu.tensor_split[1] == 1.0f); + } else { + printf("test-arg-parser: skip asymmetric multi-GPU input test (fewer than two non-CPU devices)\n\n"); + } + } + argv = {"binary_name", "-lm", "none"}; assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); assert(params.load_mode == LLAMA_LOAD_MODE_NONE); diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index 5d1f0140b623cb9a44476f8e15a0279e9e3a9887..90ef3522d3af21dab92a040725633755fdfee7d9 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -8,15 +8,60 @@ #include #include -static llama_context * make_ctx(const common_params & params, llama_model * model) { +static llama_context * make_ctx( + const common_params & params, + llama_model * model, + uint32_t n_rs_seq = 8, + uint32_t n_seq_max = 1) { auto cparams = common_context_params_to_llama(params); - cparams.n_seq_max = 1; - cparams.n_rs_seq = 8; + cparams.n_seq_max = n_seq_max; + cparams.n_rs_seq = n_rs_seq; cparams.n_batch = std::max(cparams.n_batch, (uint32_t) (cparams.n_rs_seq + 1)); cparams.n_ubatch = std::max(cparams.n_ubatch, (uint32_t) (cparams.n_rs_seq + 1)); return llama_init_from_model(model, cparams); } +static bool test_fragmented_on_device_fallback( + const common_params & params, + llama_model * model, + const std::vector & tokens) { + llama_context * ctx = make_ctx(params, model, 1, 3); + if (ctx == nullptr) { + fprintf(stderr, "%s : failed to initialize fragmented context\n", __func__); + return false; + } + + llama_batch batch = llama_batch_init(3, 0, 3); + for (llama_seq_id seq_id = 0; seq_id < 3; ++seq_id) { + common_batch_add(batch, tokens[seq_id % tokens.size()], 0, { seq_id }, seq_id == 2); + } + bool ok = llama_decode(ctx, batch) == 0; + llama_batch_free(batch); + + if (ok) { + ok = llama_memory_seq_rm(llama_get_memory(ctx), 1, -1, -1); + } + + common_prompt_checkpoint ckpt; + constexpr llama_state_seq_flags flags = + LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; + if (ok) { + ckpt.update_tgt(ctx, -1, flags); + if (ckpt.data_tgt_on_device) { + fprintf(stderr, "%s : fragmented checkpoint did not fall back to host storage\n", __func__); + ok = false; + } + } + if (ok) { + // The caller still requests ON_DEVICE here. load_tgt() must use the + // actual host mode recorded when the checkpoint was saved. + ckpt.load_tgt(ctx, -1, flags); + } + + llama_free(ctx); + return ok; +} + static bool decode_tokens(llama_context * ctx, const std::vector & tokens, uint32_t count) { llama_batch batch = llama_batch_init(count, 0, 1); for (uint32_t pos = 0; pos < count; ++pos) { @@ -35,6 +80,156 @@ static bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos) { return ok; } +static bool decode_range( + llama_context * ctx, + const std::vector & tokens, + uint32_t first, + uint32_t count) { + if (count == 0) { + return true; + } + + llama_batch batch = llama_batch_init(count, 0, 1); + for (uint32_t i = 0; i < count; ++i) { + common_batch_add(batch, tokens[first + i], first + i, { 0 }, i + 1 == count); + } + const bool ok = llama_decode(ctx, batch) == 0; + llama_batch_free(batch); + return ok; +} + +static bool compare_logits( + llama_context * ctx_full, + llama_context * ctx_depth, + int n_vocab, + uint32_t n_accepted, + const char * stage, + float eps) { + const float * logits_full = llama_get_logits(ctx_full); + const float * logits_depth = llama_get_logits(ctx_depth); + if (logits_full == nullptr || logits_depth == nullptr) { + fprintf(stderr, "%s : missing %s logits after accepting %u draft tokens\n", __func__, stage, n_accepted); + return false; + } + + int argmax_full = 0; + int argmax_depth = 0; + float max_diff = 0.0f; + for (int token = 0; token < n_vocab; ++token) { + if (logits_full[token] > logits_full[argmax_full]) { + argmax_full = token; + } + if (logits_depth[token] > logits_depth[argmax_depth]) { + argmax_depth = token; + } + + const float diff = std::fabs(logits_full[token] - logits_depth[token]); + max_diff = std::max(max_diff, diff); + if (eps >= 0.0f && diff > eps) { + fprintf(stderr, "%s : %s logits mismatch after accepting %u draft tokens, token %d (%g != %g)\n", + __func__, stage, n_accepted, token, (double) logits_full[token], (double) logits_depth[token]); + return false; + } + } + if (eps < 0.0f && argmax_full != argmax_depth) { + fprintf(stderr, "%s : %s greedy token mismatch after accepting %u draft tokens (%d != %d, max logit diff %g)\n", + __func__, stage, n_accepted, argmax_full, argmax_depth, (double) max_diff); + return false; + } + if (eps < 0.0f) { + fprintf(stderr, "%s : %s greedy token %d preserved after accepting %u draft tokens (max logit diff %g)\n", + __func__, stage, argmax_full, n_accepted, (double) max_diff); + } + return true; +} + +static bool test_recompute_fallback( + const common_params & params, + llama_model * model, + const std::vector & input_tokens, + int n_vocab) { + constexpr uint32_t n_prefix = 4; + constexpr uint32_t n_draft = 5; + constexpr uint32_t n_verify = n_draft + 1; + + std::vector tokens = input_tokens; + tokens.resize(n_prefix + n_verify + 2, tokens.back()); + + for (uint32_t n_accepted = 0; n_accepted < n_draft; ++n_accepted) { + llama_context * ctx_full = make_ctx(params, model, n_draft); + llama_context * ctx_depth = make_ctx(params, model, 1); + if (ctx_full == nullptr || ctx_depth == nullptr) { + fprintf(stderr, "%s : failed to initialize fallback contexts\n", __func__); + llama_free(ctx_full); + llama_free(ctx_depth); + return false; + } + + bool ok = decode_range(ctx_full, tokens, 0, n_prefix) && + decode_range(ctx_depth, tokens, 0, n_prefix); + if (ok) { + ok = compare_logits(ctx_full, ctx_depth, n_vocab, n_accepted, "prefix", 1e-5f); + } + + constexpr llama_state_seq_flags partial_flags = + LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; + common_prompt_checkpoint ckpt; + if (ok) { + ckpt.update_tgt(ctx_depth, 0, partial_flags); + if (!ckpt.data_tgt_on_device) { + fprintf(stderr, "%s : contiguous recurrent checkpoint unexpectedly fell back to host storage\n", __func__); + ok = false; + } + } + if (ok) { + ok = decode_range(ctx_full, tokens, n_prefix, n_verify) && + decode_range(ctx_depth, tokens, n_prefix, n_verify); + } + + const llama_pos rollback_pos = n_prefix + 1 + n_accepted; + const uint32_t n_rollback = n_draft - n_accepted; + if (ok) { + ok = llama_memory_seq_rm(llama_get_memory(ctx_full), 0, rollback_pos, -1); + } + if (ok && n_rollback <= 1) { + ok = llama_memory_seq_rm(llama_get_memory(ctx_depth), 0, rollback_pos, -1); + } else if (ok) { + ckpt.load_tgt(ctx_depth, 0, partial_flags); + ok = llama_memory_seq_rm(llama_get_memory(ctx_depth), 0, n_prefix, -1); + } + + const llama_pos replacement_pos = rollback_pos; + if (ok && n_rollback <= 1) { + ok = decode_one(ctx_full, tokens[n_prefix + n_verify], replacement_pos) && + decode_one(ctx_depth, tokens[n_prefix + n_verify], replacement_pos); + } else if (ok) { + std::vector replay_tokens = tokens; + replay_tokens[replacement_pos] = tokens[n_prefix + n_verify]; + ok = decode_one(ctx_full, tokens[n_prefix + n_verify], replacement_pos) && + decode_range(ctx_depth, replay_tokens, n_prefix, 2 + n_accepted); + } + if (ok) { + ok = compare_logits(ctx_full, ctx_depth, n_vocab, n_accepted, "replacement", -1.0f); + } + if (ok) { + ok = decode_one(ctx_full, tokens[n_prefix + n_verify + 1], replacement_pos + 1) && + decode_one(ctx_depth, tokens[n_prefix + n_verify + 1], replacement_pos + 1) && + compare_logits(ctx_full, ctx_depth, n_vocab, n_accepted, "continuation", -1.0f); + } + + llama_free(ctx_full); + llama_free(ctx_depth); + + if (!ok) { + fprintf(stderr, "%s : fallback validation failed after accepting %u draft tokens\n", __func__, n_accepted); + return false; + } + } + + fprintf(stderr, "%s : depth-1 recomputation preserves depth-5 greedy decisions at every rejection position\n", __func__); + return true; +} + int main(int argc, char ** argv) { std::setlocale(LC_NUMERIC, "C"); @@ -65,6 +260,23 @@ int main(int argc, char ** argv) { const llama_vocab * vocab = llama_model_get_vocab(model); const int n_vocab = llama_vocab_n_tokens(vocab); + std::vector tokens; + if (llama_vocab_type(vocab) == LLAMA_VOCAB_TYPE_NONE) { + tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 }; + } else { + tokens = common_tokenize(vocab, "The quick brown fox jumps over the lazy dog", true); + } + if (tokens.empty()) { + fprintf(stderr, "%s : not enough prompt tokens\n", __func__); + return 1; + } + if (!test_fragmented_on_device_fallback(params, model, tokens)) { + return 1; + } + if (!test_recompute_fallback(params, model, tokens, n_vocab)) { + return 1; + } + llama_context * ctx_src = make_ctx(params, model); llama_context * ctx_dst = make_ctx(params, model); if (ctx_src == nullptr || ctx_dst == nullptr) { @@ -79,12 +291,6 @@ int main(int argc, char ** argv) { return 0; } - std::vector tokens; - if (llama_vocab_type(vocab) == LLAMA_VOCAB_TYPE_NONE) { - tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 }; - } else { - tokens = common_tokenize(ctx_src, "The quick brown fox jumps over the lazy dog", true); - } const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src); constexpr uint32_t n_rollback = 3; if (n_rs_seq < n_rollback) { @@ -93,10 +299,6 @@ int main(int argc, char ** argv) { llama_free(ctx_dst); return 0; } - if (tokens.empty()) { - fprintf(stderr, "%s : not enough prompt tokens\n", __func__); - return 1; - } tokens.resize(n_rs_seq + 1, tokens.back()); const uint32_t n_tokens = tokens.size(); @@ -161,6 +363,10 @@ int main(int argc, char ** argv) { constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; common_prompt_checkpoint ckpt_partial; ckpt_partial.update_tgt(ctx_src, 0, partial_flags); + if (ckpt_partial.data_tgt_on_device) { + fprintf(stderr, "%s : host recurrent checkpoint recorded device storage\n", __func__); + return 1; + } ckpt_partial.load_tgt(ctx_dst, 0, partial_flags); if (!replay_and_compare("partial")) { diff --git a/tools/server/README.md b/tools/server/README.md index b63a0e6dac00334650536616b74d7c33d9c7f82a..77c1d33113329440e946b6f790d8168d7c6b4e33 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -254,6 +254,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-cpu-moe, -cmoed, --cpu-moe-draft` | keep all Mixture of Experts (MoE) weights in the CPU for the draft model
(env: LLAMA_ARG_SPEC_DRAFT_CPU_MOE) | | `--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-mtp-rs-depth N` | number of target recurrent-state rollback snapshots for MTP; lower values save memory but replay accepted tokens after deep rejection (default: `--spec-draft-n-max`)
(env: LLAMA_ARG_SPEC_MTP_RS_DEPTH) | | `--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-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 585f65e83c655d3b8b7e398e8bf76552dc846f36..e57548fce8f7c9a3873e9556e0366c6b9d538214 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -83,6 +83,10 @@ json server_slot_stats::to_json() const { base["draft_n"] = n_draft_tokens; base["draft_n_accepted"] = n_draft_accepted; } + if (n_draft_replay_count > 0) { + base["draft_replay_count"] = n_draft_replay_count; + base["draft_replay_n"] = n_draft_replay_tokens; + } return base; } diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 6488be344c6ad5d7100cae854b6047f5728f4989..5a3d135f5f514d8362f884b4aa858dac0b1725a0 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -353,6 +353,8 @@ struct server_slot_stats { uint64_t n_draft_tokens = 0; uint64_t n_draft_accepted = 0; uint64_t n_draft_verif_steps = 0; + uint64_t n_draft_replay_count = 0; + uint64_t n_draft_replay_tokens = 0; // these are absolute timestamps (in us) // note: must be signed - they are subtracted before the later ones are set diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 21ff783941382346293429b30e2ff1501516a491..73bed4ff249876f5b23e322a89a32fa9e73d5ea7 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -636,6 +636,12 @@ struct server_slot { " acc per pos = (%s)\n", acceptance_rates_per_pos.c_str()); } + if (stats.n_draft_replay_count > 0) { + SLT_INF(*this, + " MTP replays = %10" PRIu64 " events / %5" PRIu64 " tokens\n", + stats.n_draft_replay_count, stats.n_draft_replay_tokens); + } + common_speculative_print_stats(spec); } @@ -1039,6 +1045,85 @@ private: } } + // The on-device recurrent speculative checkpoint is allocated lazily, so reserve + // the additional recurrent-state group before fitting the target model. + if (params_base.fit_params && spec_mtp) { + const uint32_t n_rs_seq = params_base.speculative.need_n_rs_seq(); + const int32_t draft_n_max = params_base.speculative.draft.n_max; + if (draft_n_max < 0) { + SRV_ERR("%s", "[spec] invalid negative MTP draft.n_max while fitting\n"); + return false; + } + if (n_rs_seq < static_cast(draft_n_max)) { + try { + common_params params_rs = params_base; + auto mparams_rs = common_model_params_to_llama(params_rs); + auto cparams_rs = common_context_params_to_llama(params_rs); + + std::vector devs_rs; + uint32_t hp_ngl_rs = 0; + uint32_t hp_nct_rs = 0; + uint32_t hp_nex_rs = 0; + cparams_rs.n_rs_seq = n_rs_seq; + const auto dmd_rs = common_get_device_memory_data( + params_base.model.path.c_str(), &mparams_rs, &cparams_rs, + devs_rs, hp_ngl_rs, hp_nct_rs, hp_nex_rs, GGML_LOG_LEVEL_ERROR); + + std::vector devs_rs_next; + uint32_t hp_ngl_rs_next = 0; + uint32_t hp_nct_rs_next = 0; + uint32_t hp_nex_rs_next = 0; + cparams_rs.n_rs_seq = n_rs_seq + 1; + const auto dmd_rs_next = common_get_device_memory_data( + params_base.model.path.c_str(), &mparams_rs, &cparams_rs, + devs_rs_next, hp_ngl_rs_next, hp_nct_rs_next, hp_nex_rs_next, GGML_LOG_LEVEL_ERROR); + + if (devs_rs.size() != devs_rs_next.size() || dmd_rs.size() < devs_rs.size() || + dmd_rs_next.size() < devs_rs_next.size()) { + throw std::runtime_error("target context device count changed during recurrent-state measurement"); + } + + // common_fit_params() consumes fit_params_target in the measured model-device + // order. The configured device list can contain a nullptr sentinel, CPU/ACCEL + // devices, or be reduced by split-mode none, so it is not a safe index map. + if (params_base.fit_params_target.size() < devs_rs.size()) { + throw std::runtime_error("fit_params_target has no entry for every target device"); + } + + size_t total = 0; + for (size_t j = 0; j < devs_rs.size(); ++j) { + const auto next_dev = std::find(devs_rs_next.begin(), devs_rs_next.end(), devs_rs[j]); + if (next_dev == devs_rs_next.end()) { + throw std::runtime_error("target context device mapping changed during recurrent-state measurement"); + } + const size_t next_index = next_dev - devs_rs_next.begin(); + if (next_index != j) { + throw std::runtime_error("target context device order changed during recurrent-state measurement"); + } + if (dmd_rs_next[j].context < dmd_rs[j].context) { + throw std::runtime_error("recurrent-state context measurement decreased"); + } + const size_t delta = dmd_rs_next[j].context - dmd_rs[j].context; + params_base.fit_params_target[j] += delta; + total += delta; + SRV_INF("[spec] recurrent checkpoint fit reservation: device %s, %.2f MiB " + "(context %.2f -> %.2f MiB)\n", + ggml_backend_dev_name(devs_rs[j]), delta / (1024.0 * 1024.0), + dmd_rs[j].context / (1024.0 * 1024.0), + dmd_rs_next[j].context / (1024.0 * 1024.0)); + } + if (total == 0) { + throw std::runtime_error("recurrent-state measurement produced no device reservation"); + } + SRV_INF("[spec] recurrent checkpoint fit reservation: %.2f MiB total\n", + total / (1024.0 * 1024.0)); + } catch (const std::exception & e) { + SRV_ERR("[spec] failed to reserve recurrent-state memory before fitting: %s\n", e.what()); + return false; + } + } + } + // optionally reserve VRAM for the draft / MTP context before fitting the target model if (params_base.fit_params) { if (has_spec) { @@ -1051,6 +1136,25 @@ private: auto cparams_dft = common_context_params_to_llama(params_dft); if (spec_mtp) { cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams_dft.pipeline_parallel_type = LLAMA_PIPELINE_PARALLEL_TYPE_DISABLED; + // MTP draft uses the target's (possibly fit-reduced) n_ctx at runtime. + // Estimate with the same n_ctx to avoid over-reserving fit_params_target. + { + // Fit against a scratch copy. The fitter may write an automatically + // generated tensor split and tensor overrides; publishing those to + // params_base here would make the refinement and final target fit + // mistake them for user-provided placement constraints. + common_params params_fit = params_base; + auto mparams_tgt = common_model_params_to_llama(params_fit); + auto cparams_tgt = common_context_params_to_llama(params_fit); + common_fit_params(params_dft.model.path.c_str(), &mparams_tgt, &cparams_tgt, + params_fit.tensor_split, + params_fit.tensor_buft_overrides.data(), + params_fit.fit_params_target.data(), + params_fit.fit_params_min_ctx, + GGML_LOG_LEVEL_ERROR); + cparams_dft.n_ctx = cparams_tgt.n_ctx; + } } cparams_dft.n_rs_seq = 0; @@ -1074,6 +1178,7 @@ private: } } + std::vector measured_bytes(tgt_devices.size(), 0); for (size_t j = 0; j < devs.size(); ++j) { const size_t bytes = (measure_model_bytes ? dmd[j].model : 0) + dmd[j].context + dmd[j].compute; total += bytes; @@ -1082,6 +1187,7 @@ private: SRV_DBG("[spec] adding %.2f MiB to fit_params_target for device %s\n", bytes / (1024.0 * 1024.0), ggml_backend_dev_name(devs[j])); params_base.fit_params_target[i] += bytes; + measured_bytes[i] += bytes; break; } } @@ -1089,6 +1195,51 @@ private: SRV_TRC("[spec] estimated memory usage of %s is %.2f MiB\n", has_draft ? "draft model" : "MTP context", total / (1024.0 * 1024.0)); + + if (spec_mtp) { + for (int pass = 0; pass < 2; ++pass) { + // Each refinement starts without placement produced by an earlier + // estimate. Only the final common_init_from_params() call should + // commit fitted placement to params_base. + common_params params_fit = params_base; + auto mparams_tgt = common_model_params_to_llama(params_fit); + auto cparams_tgt = common_context_params_to_llama(params_fit); + common_fit_params(params_dft.model.path.c_str(), &mparams_tgt, &cparams_tgt, + params_fit.tensor_split, + params_fit.tensor_buft_overrides.data(), + params_fit.fit_params_target.data(), + params_fit.fit_params_min_ctx, + GGML_LOG_LEVEL_ERROR); + if (cparams_dft.n_ctx == cparams_tgt.n_ctx) { + break; + } + + cparams_dft.n_ctx = cparams_tgt.n_ctx; + auto dmd_refined = common_get_device_memory_data( + params_dft.model.path.c_str(), &mparams_dft, &cparams_dft, + devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR); + + std::vector refined_bytes(tgt_devices.size(), 0); + total = 0; + for (size_t j = 0; j < devs.size(); ++j) { + const size_t bytes = (measure_model_bytes ? dmd_refined[j].model : 0) + dmd_refined[j].context + dmd_refined[j].compute; + total += bytes; + for (size_t i = 0; i < tgt_devices.size(); ++i) { + if (tgt_devices[i] == devs[j]) { + refined_bytes[i] += bytes; + break; + } + } + } + for (size_t i = 0; i < tgt_devices.size(); ++i) { + params_base.fit_params_target[i] -= measured_bytes[i]; + params_base.fit_params_target[i] += refined_bytes[i]; + } + measured_bytes = std::move(refined_bytes); + SRV_TRC("[spec] refined MTP memory estimate at n_ctx=%" PRIu32 " is %.2f MiB\n", + cparams_dft.n_ctx, total / (1024.0 * 1024.0)); + } + } } catch (const std::exception & e) { SRV_WRN("[spec] failed to measure %s memory: %s\n", has_draft ? "draft model" : "MTP context", e.what()); @@ -3005,9 +3156,18 @@ private: (ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft)); if (use_ckpt_tgt) { + llama_state_seq_flags ckpt_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; + if (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS) { + // Phase 1 proof of concept: avoid copying the recurrent fallback checkpoint + // through host memory on every speculative round. The device buffer is + // allocated lazily; load_model() reserves one measured recurrent-state + // group for it when automatic fitting and reduced RS depth are enabled. + ckpt_flags |= LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; + } + //const int64_t t_start = ggml_time_us(); - ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_tgt(ctx_tgt, slot.id, ckpt_flags); //const int64_t t_total = ggml_time_us() - t_start; //printf("checkpoint total: %f ms\n", t_total / 1000.0); @@ -3851,6 +4011,10 @@ private: 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)); + const bool use_rs_replay = + 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) { if (use_ckpt_tgt) { @@ -3862,11 +4026,20 @@ private: slot.spec_is_replay = true; slot.spec_draft = std::move(accepted); + if (use_rs_replay) { + slot.stats.n_draft_replay_count += 1; + slot.stats.n_draft_replay_tokens += slot.spec_draft.size() + 1; + } + const auto & ckpt = slot.spec_ckpt; SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size()); - ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + llama_state_seq_flags ckpt_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; + if (use_rs_replay) { + ckpt_flags |= LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; + } + ckpt.load_tgt(slot.ctx_tgt, slot.id, ckpt_flags); if (slot.ctx_dft) { ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index 9171dbc02977e26ce5e2e9b876d73457583314e0..6e6ca5b82f399a0d3e30739a5de48fa4a5e2cbe1 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -99,6 +99,7 @@ class ServerProcess: spec_type: str | None = None spec_draft_n_min: int | None = None spec_draft_n_max: int | None = None + spec_mtp_rs_depth: int | None = None no_ui: bool | None = None jinja: bool | None = None reasoning_format: Literal['deepseek', 'none', 'nothink'] | None = None @@ -245,6 +246,8 @@ class ServerProcess: server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max]) if self.spec_draft_n_min: server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min]) + if self.spec_mtp_rs_depth: + server_args.extend(["--spec-mtp-rs-depth", self.spec_mtp_rs_depth]) if self.no_ui: server_args.append("--no-ui") if self.no_models_autoload: