# Applies to llama.cpp commit 73a43d1 (latest master, 42 commits after 0df017d6) # Same content as latest_rocm_improvement_0df017d6.patch, rebased onto current master (zero conflicts) # Includes eaman commit 7f07026c (complete Eaman ROCm fitting/pipeline improvements, MTP Compact Rollback with adaptive MTP and multi-ubatch synchronization, and opt-in HIP VEC forcing) # Verified: full build compiles clean on 73a43d1; test-arg-parser passes diff --git a/common/arg.cpp b/common/arg.cpp index 015196c..d446833 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1296,6 +1296,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-cr-depth must not exceed --spec-draft-n-max"); + } if (ctx_arg.params.usage) { common_params_print_usage(ctx_arg); if (ctx_arg.print_usage) { @@ -1768,6 +1771,34 @@ 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({ "--hip-fa-force-vec" }, "[on|off]", + string_format("force the HIP quantized-KV Flash Attention VEC path when supported (default: '%s')", + params.hip_fa_force_vec ? "on" : "off"), + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.hip_fa_force_vec = true; + } else if (is_falsey(value)) { + params.hip_fa_force_vec = false; + } else { + throw std::runtime_error( + string_format("error: unknown value for --hip-fa-force-vec: '%s'\n", value.c_str())); + } + }).set_env("LLAMA_ARG_HIP_FA_FORCE_VEC")); + 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", @@ -4164,6 +4195,23 @@ 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-cr-depth"}, "N", + "MTP Compact Rollback depth; 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-cr-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_CR_DEPTH")); + add_opt(common_arg( + {"--spec-draft-adaptive"}, + string_format("size each draft from measured acceptance rather than always drafting --spec-draft-n-max (default: %s)", params.speculative.draft.adaptive ? "true" : "false"), + [](common_params & params) { + params.speculative.draft.adaptive = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_ADAPTIVE")); 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 d162a38..19f48fd 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1307,6 +1307,7 @@ common_init_result::common_init_result(common_params & params, bool model_only) 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; } cparams_dft.n_rs_seq = 0; @@ -1741,6 +1742,8 @@ 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.hip_fa_force_vec = params.hip_fa_force_vec; + 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; @@ -2269,6 +2272,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( @@ -2280,6 +2285,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, @@ -2288,14 +2337,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( @@ -2306,14 +2348,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( @@ -2328,6 +2363,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); @@ -2346,6 +2384,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); @@ -2354,9 +2395,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 63d0bad..da8a243 100644 --- a/common/common.h +++ b/common/common.h @@ -325,6 +325,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 Compact Rollback depth (-1 = n_max) float p_split = 0.1f; // speculative decoding split probability float p_min = 0.0f; // minimum speculative decoding probability (greedy) @@ -341,6 +342,9 @@ struct common_params_speculative_draft { ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V + // size each draft from measured acceptance instead of always drafting n_max + bool adaptive = false; + common_cpu_params cpuparams; common_cpu_params cpuparams_batch; @@ -392,11 +396,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; } }; @@ -497,6 +508,8 @@ 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 + bool hip_fa_force_vec = false; // force HIP quantized-KV FA onto VEC when supported + 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; @@ -1174,6 +1187,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 c601fe4..d6169ff 100644 --- a/common/fit.cpp +++ b/common/fit.cpp @@ -34,7 +34,8 @@ static std::vector common_get_device_memory_data_impl( uint32_t & hp_ngl, uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert, - ggml_log_level log_level) { + ggml_log_level log_level, + std::vector * checkpoint_sizes = nullptr) { struct user_data_t { struct { ggml_log_callback callback; @@ -96,6 +97,28 @@ static std::vector common_get_device_memory_data_impl( } } + if (checkpoint_sizes) { + checkpoint_sizes->assign(nd + 1, 0); + const auto checkpoint_breakdown = llama_get_state_seq_device_buffer_sizes(ctx); + for (const auto & [buft, size] : checkpoint_breakdown) { + if (ggml_backend_buft_is_host(buft)) { + checkpoint_sizes->back() += size; + continue; + } + + ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft); + if (!dev) { + continue; + } + for (size_t i = 0; i < nd; ++i) { + if (dev == llama_model_get_device(model, i)) { + (*checkpoint_sizes)[i] += size; + break; + } + } + } + } + { ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); if (cpu_dev == nullptr) { @@ -161,8 +184,10 @@ common_device_memory_data_vec common_get_device_memory_data( uint32_t & hp_n_ctx_train, uint32_t & hp_n_expert, ggml_log_level log_level) { + std::vector checkpoint_sizes; std::vector impl = common_get_device_memory_data_impl( - path_model, mparams, cparams, devs, hp_ngl, hp_n_ctx_train, hp_n_expert, log_level); + path_model, mparams, cparams, devs, hp_ngl, hp_n_ctx_train, hp_n_expert, log_level, + &checkpoint_sizes); common_device_memory_data_vec ret(impl.size()); for (size_t i = 0; i < impl.size(); i++) { @@ -171,6 +196,7 @@ common_device_memory_data_vec common_get_device_memory_data( ret[i].model = impl[i].mb.model; ret[i].context = impl[i].mb.context; ret[i].compute = impl[i].mb.compute; + ret[i].checkpoint = checkpoint_sizes[i]; } return ret; } @@ -185,6 +211,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 @@ -399,7 +426,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 @@ -420,6 +447,31 @@ 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 = n_ctx_max; + 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_total; + fits = fits && target >= used_min; + if (target > used_min && used_full > used_min) { + n_ctx_device += (n_ctx_max - n_ctx_min_total) * (target - used_min) / (used_full - used_min); + } + n_ctx_fit = std::min(n_ctx_fit, n_ctx_device); + } + const uint32_t align = 256 * n_streams; + cparams->n_ctx = std::max(n_ctx_fit - n_ctx_fit % align, n_ctx_min_total); + LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " using per-device limits\n", + __func__, n_ctx_max, 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 += (n_ctx_max - n_ctx_min_total) * (sum_used_target - sum_projected_used_min_ctx) @@ -459,7 +511,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/fit.h b/common/fit.h index 824d386..7022746 100644 --- a/common/fit.h +++ b/common/fit.h @@ -51,6 +51,7 @@ struct common_device_memory_data { size_t model; size_t context; size_t compute; + size_t checkpoint; }; using common_device_memory_data_vec = std::vector; diff --git a/common/speculative.cpp b/common/speculative.cpp index 851a47b..19b331b 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -159,8 +159,97 @@ struct common_speculative_impl { int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. - common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) : type(type), n_seq(n_seq), n_max(n_max) {} + // Adaptive draft length. The configured n_max is the known-good steady + // state. Ordinary tail rejection does not justify shortening because a + // wider verification batch is usually cheaper than another target decode + // round. Two consecutive drafts accepting fewer than three tokens identify + // a real phase change and step the limit down by one. A fully accepted draft + // steps it back up by one. + std::vector adaptive_limit; + std::vector adaptive_floor; + std::vector adaptive_ceiling; + std::vector n_last_draft; + std::vector early_reject_streak; + std::vector adaptive_limit_reached; + bool adaptive_n = false; + + static constexpr int32_t adaptive_productive = 3; + static constexpr uint8_t adaptive_patience = 2; + + void update_adaptive_limit(llama_seq_id seq_id, int32_t n_accepted) { + if (!adaptive_n || seq_id < 0 || (size_t) seq_id >= adaptive_limit.size()) { + return; + } + const int32_t n_drafted = n_last_draft[seq_id]; + if (n_drafted <= 0) { + return; + } + const bool limit_reached = adaptive_limit_reached[seq_id] != 0; + adaptive_limit_reached[seq_id] = 0; + if (!limit_reached) { + // Confidence stopping (p_min), a per-slot cap, or a decode failure + // already terminated this draft. Do not reinterpret that shorter + // observation as evidence for changing the phase-level ceiling. + early_reject_streak[seq_id] = 0; + n_last_draft[seq_id] = 0; + return; + } + + if (n_accepted >= n_drafted) { + early_reject_streak[seq_id] = 0; + adaptive_limit[seq_id] = std::min(adaptive_ceiling[seq_id], adaptive_limit[seq_id] + 1); + } else if (n_accepted >= std::min(adaptive_productive, n_drafted)) { + early_reject_streak[seq_id] = 0; + } else if (++early_reject_streak[seq_id] >= adaptive_patience) { + early_reject_streak[seq_id] = 0; + adaptive_limit[seq_id] = std::max(adaptive_floor[seq_id], adaptive_limit[seq_id] - 1); + } + n_last_draft[seq_id] = 0; + } + + // Reset on a new prompt / reused server slot; never mid-generation, since + // tracking phase changes within a response is the point of adaptation. + void reset_adaptive_limit(llama_seq_id seq_id) { + if (seq_id >= 0 && (size_t) seq_id < adaptive_limit.size()) { + adaptive_limit[seq_id] = 0; + adaptive_floor[seq_id] = 1; + adaptive_ceiling[seq_id] = 0; + n_last_draft[seq_id] = 0; + early_reject_streak[seq_id] = 0; + adaptive_limit_reached[seq_id] = 0; + } + } + + void record_adaptive_draft(llama_seq_id seq_id, size_t n_drafted) { + if (adaptive_n && seq_id >= 0 && (size_t) seq_id < n_last_draft.size()) { + n_last_draft[seq_id] = static_cast(n_drafted); + adaptive_limit_reached[seq_id] = n_drafted >= (size_t) adaptive_limit[seq_id]; + } + } + // effective draft length for this step, never above the configured n_max + int32_t adaptive_n_draft(llama_seq_id seq_id, int32_t n_cfg, int32_t n_min) { + if (!adaptive_n || n_cfg <= 0 || seq_id < 0 || (size_t) seq_id >= adaptive_limit.size()) { + return n_cfg; + } + adaptive_floor[seq_id] = std::clamp(n_min, 1, n_cfg); + adaptive_ceiling[seq_id] = n_cfg; + if (adaptive_limit[seq_id] <= 0) { + adaptive_limit[seq_id] = n_cfg; + } + adaptive_limit[seq_id] = std::clamp(adaptive_limit[seq_id], adaptive_floor[seq_id], n_cfg); + return adaptive_limit[seq_id]; + } + + common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) + : type(type), n_seq(n_seq), n_max(n_max) { + adaptive_limit.assign(n_seq, 0); + adaptive_floor.assign(n_seq, 1); + adaptive_ceiling.assign(n_seq, 0); + n_last_draft.assign(n_seq, 0); + early_reject_streak.assign(n_seq, 0); + adaptive_limit_reached.assign(n_seq, 0); + } virtual ~common_speculative_impl() = default; virtual void begin(llama_seq_id seq_id, const llama_tokens & prompt) = 0; @@ -195,6 +284,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { } SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n"); + adaptive_n = this->params.adaptive; SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min); SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", this->params.n_gpu_layers, @@ -255,8 +345,8 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { llama_batch_free(batch); } - void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override { - // noop + void begin(llama_seq_id seq_id, const llama_tokens & /*prompt*/) override { + reset_adaptive_limit(seq_id); } bool process(const llama_batch & batch) override { @@ -319,6 +409,8 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { auto * smpl = smpls[seq_id].get(); + const int32_t n_draft_eff = adaptive_n_draft(seq_id, params.n_max, params.n_min); + common_sampler_sample(smpl, ctx_dft, i_batch, true); ++i_batch; @@ -348,7 +440,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { result.push_back(id); - if ((params.n_max <= (int) result.size()) || + if ((n_draft_eff <= (int) result.size()) || (dp.n_max > 0 && dp.n_max <= (int) result.size())) { drafting[seq_id] = false; n_drafting--; @@ -458,6 +550,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { , params(params.draft) { SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n"); + adaptive_n = this->params.adaptive; SPC_TRC("- n_max=%d, n_min=%d, p_min=%f, backend_sampling=%d\n", params.draft.n_max, params.draft.n_min, params.draft.p_min, (int) params.draft.backend_sampling); auto * ctx_tgt = this->params.ctx_tgt; @@ -554,6 +647,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { } void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { + reset_adaptive_limit(seq_id); + const int32_t N = (int32_t) prompt.size(); if (N <= 0) { return; @@ -779,6 +874,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { auto * smpl = smpls[seq_id].get(); + const int32_t n_draft_eff = adaptive_n_draft(seq_id, params.n_max, params.n_min); + common_sampler_sample(smpl, ctx_dft, i_batch, true); // pre-norm hidden state of this position becomes g_embd for the next step const float * prenorm = llama_get_embeddings_nextn_ith(ctx_dft, i_batch); @@ -810,7 +907,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { result.push_back(id); - if (params.n_max <= (int) result.size()) { + if (n_draft_eff <= (int) result.size()) { drafting[seq_id] = false; n_drafting--; continue; @@ -992,6 +1089,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); + adaptive_n = this->params.adaptive; LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__, block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false"); @@ -1070,6 +1168,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { + + reset_adaptive_limit(seq_id); if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { return; } @@ -1193,7 +1293,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n = (int32_t) dp.n_past; - const int32_t n_draft = params.n_max; + // DFlash decodes the whole block in one pass, so a shorter block does + // not save draft time -- but it does shrink the target's verification + // batch, which is where the cost actually is. + const int32_t n_draft = adaptive_n_draft(seq_id, params.n_max, params.n_min); const int32_t n_block_tokens = n_draft + (is_dspark && sample_from_anchor ? 0 : 1); i_block_beg[seq_id] = batch.n_tokens; @@ -1371,6 +1474,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { n_mtp_layers = std::max(1, (int) llama_model_n_layer_nextn(llama_get_model(ctx_dft))); SPC_TRC("%s", "adding speculative implementation 'draft-mtp'\n"); + adaptive_n = this->params.adaptive; SPC_TRC("- n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling); SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", this->params.n_gpu_layers, @@ -1458,6 +1562,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { + + reset_adaptive_limit(seq_id); const int32_t N = (int32_t) prompt.size(); if (N <= 0) { return; @@ -1662,6 +1768,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { auto * smpl = smpls[seq_id].get(); + // MTP drafts sequentially, so a shorter draft saves draft passes + // as well as target verification work. + const int32_t n_draft_eff = adaptive_n_draft(seq_id, params.n_max, params.n_min); + common_sampler_sample(smpl, ctx_dft, i_last[seq_id], true); const float * h_row = llama_get_embeddings_nextn_ith(ctx_dft, i_last[seq_id]); @@ -1691,7 +1801,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { result.push_back(id); - if (params.n_max <= (int) result.size()) { + if (n_draft_eff <= (int) result.size()) { drafting[seq_id] = false; n_drafting--; continue; @@ -2527,6 +2637,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); } // the draft context holds as many tokens per sequence as the target context @@ -2846,6 +2960,7 @@ void common_speculative_draft(common_speculative * spec) { // remember which implementation was used spec->impl_last[seq_id] = impl.get(); + impl->record_adaptive_draft(seq_id, result.size()); impl->n_gen_drafts++; impl->n_gen_tokens += result.size(); @@ -2872,7 +2987,11 @@ void common_speculative_draft(common_speculative * spec) { } } -void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, uint16_t n_accepted) { +void common_speculative_accept( + common_speculative * spec, + llama_seq_id seq_id, + uint16_t n_accepted, + int32_t n_accepted_observed) { common_speculative_impl * impl = spec->impl_last[seq_id]; if (impl == nullptr) { @@ -2896,6 +3015,7 @@ void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, u impl->n_acc_tokens += n_accepted; } + impl->update_adaptive_limit(seq_id, n_accepted_observed >= 0 ? n_accepted_observed : n_accepted); impl->accept(seq_id, n_accepted, false); impl->n_call_accept++; } diff --git a/common/speculative.h b/common/speculative.h index 2250589..681ff3b 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -82,8 +82,14 @@ bool common_speculative_process(common_speculative * spec, const llama_batch & b // generate drafts for the sequences specified with `common_speculative_get_draft_params` void common_speculative_draft(common_speculative * spec); -// informs the speculative context that n_accepted tokens were accepted by the target model -void common_speculative_accept(common_speculative * spec, llama_seq_id, uint16_t n_accepted); +// informs the speculative context that n_accepted tokens were accepted by the target model. +// n_accepted_observed can preserve the original acceptance measurement when a +// rollback implementation has to replay a different number of physical tokens. +void common_speculative_accept( + common_speculative * spec, + llama_seq_id seq_id, + uint16_t n_accepted, + int32_t n_accepted_observed = -1); // (optional) get/set internal state bool common_speculative_get_state(common_speculative * spec, llama_seq_id seq_id, std::vector & data); diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index b88b7e5..b822ede 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2453,6 +2453,13 @@ extern "C" { GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a); + GGML_API void ggml_flash_attn_ext_set_force_vec( + struct ggml_tensor * a, + bool force_vec); + + GGML_API bool ggml_flash_attn_ext_get_force_vec( + const struct ggml_tensor * a); + // Use finite mask entries as a sparse K/V set. Set 0 to disable. // n_kv_max must bound the number of finite entries in every mask row. GGML_API void ggml_flash_attn_ext_set_n_kv_max( diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ae217fb..1ff1dad 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -590,6 +590,15 @@ 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_flash_attn_ext_get_force_vec(dst) && + (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/ggml/src/ggml.c b/ggml/src/ggml.c index 6257cdb..f416a47 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5507,6 +5507,18 @@ enum ggml_prec ggml_flash_attn_ext_get_prec( return (enum ggml_prec) prec_i32; } +void ggml_flash_attn_ext_set_force_vec( + struct ggml_tensor * a, + bool force_vec) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + ggml_set_op_params_i32(a, 5, force_vec ? 1 : 0); +} + +bool ggml_flash_attn_ext_get_force_vec(const struct ggml_tensor * a) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + return ggml_get_op_params_i32(a, 5) != 0; +} + void ggml_flash_attn_ext_set_n_kv_max( struct ggml_tensor * a, int32_t n_kv_max) { diff --git a/include/llama.h b/include/llama.h index ef7a012..382a297 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 @@ -372,6 +380,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 @@ -416,6 +425,8 @@ extern "C" { // a source/target/parent context // can be utilized in various ways, for example by sharing results or llama_memory between 2 contexts struct llama_context * ctx_other; + + bool hip_fa_force_vec; // force HIP quantized-KV FA onto VEC when supported }; struct llama_model_tensor_override { @@ -929,6 +940,11 @@ extern "C" { llama_seq_id seq_id, llama_state_seq_flags flags); + // Allocate and pin all configured per-sequence device checkpoint buffers. + // Returns false when the context has no device checkpoint layout or an + // allocation fails. + LLAMA_API bool llama_state_seq_reserve_device_buffers(struct llama_context * ctx); + LLAMA_API size_t llama_state_seq_set_data_ext( struct llama_context * ctx, const uint8_t * src, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c1ef12f..70459c8 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -228,6 +228,7 @@ llama_context::llama_context( cparams.flash_attn = params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED; cparams.auto_fa = params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO; + cparams.hip_fa_force_vec = params.hip_fa_force_vec; cparams.fused_gdn_ar = true; cparams.fused_gdn_ch = true; @@ -310,6 +311,8 @@ 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: hip_fa_force_vec = %s\n", __func__, cparams.hip_fa_force_vec ? "true" : "false"); + 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 +429,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 +457,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) { @@ -1803,6 +1813,9 @@ int llama_context::decode(const llama_batch & batch_inp) { int64_t n_outputs_prev = 0; int64_t n_tokens_prev = 0; + bool has_next_ubatch = false; + bool mtp_multi_ubatch = false; + do { const auto & ubatch = mctx->get_ubatch(); @@ -1978,7 +1991,15 @@ int llama_context::decode(const llama_batch & batch_inp) { n_outputs_prev += n_outputs; n_tokens_prev += ubatch.n_tokens; - } while (mctx->next()); + + has_next_ubatch = mctx->next(); + mtp_multi_ubatch |= has_next_ubatch; + + // MTP ubatches update the same KV cache and must complete in order. + if (cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && mtp_multi_ubatch) { + synchronize(); + } + } while (has_next_ubatch); // set to total number of outputs in the batch, for use in llama_get_logits_ith n_outputs = n_outputs_all; @@ -2732,10 +2753,12 @@ private: class llama_io_write_device : public llama_io_write_i { public: - llama_io_write_device(uint8_t * p, size_t len, llama_memory_buffers & mbufs) : ptr(p), buf_size(len), mbufs(mbufs) { + llama_io_write_device( + uint8_t * p, size_t len, llama_memory_buffers & mbufs, bool copy_tensors = true) : + ptr(p), buf_size(len), mbufs(mbufs), copy_tensors(copy_tensors) { } - ~llama_io_write_device() { + void finish() { llama_memory_buffers mbufs_new; for (const auto & winfo : winfos) { @@ -2797,11 +2820,21 @@ public: if (need_alloc) { if (!mbuf_cur.buf || mbuf_cur.total_size != mbuf.total_size) { - mbuf_cur = std::move(mbuf); + ggml_backend_buffer_ptr buf { + ggml_backend_alloc_ctx_tensors_from_buft(mbuf.ctx.get(), buft) + }; + if (!buf) { + throw std::runtime_error( + std::string("failed to allocate device checkpoint buffer from '") + + ggml_backend_buft_name(buft) + "'"); + } - mbuf_cur.buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(mbuf_cur.ctx.get(), buft)); + const size_t allocated_size = ggml_backend_buffer_get_size(buf.get()); + mbuf.buf = std::move(buf); + mbuf_cur = std::move(mbuf); - LLAMA_LOG_INFO("%s: allocated '%s' buffer %.3f MiB\n", __func__, ggml_backend_buft_name(buft), mbuf.total_size/1024.0/1024.0); + LLAMA_LOG_INFO("%s: allocated '%s' buffer %.3f MiB\n", __func__, + ggml_backend_buft_name(buft), allocated_size/1024.0/1024.0); } else { //LLAMA_LOG_INFO("%s: reallocating tensors in '%s' buffer %.3f MiB\n", __func__, ggml_backend_buft_name(buft), mbuf.total_size/1024.0/1024.0); @@ -2821,8 +2854,10 @@ public: } } - for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { - ggml_backend_tensor_copy(mbuf_cur.org[i], mbuf_cur.cpy[i]); + if (copy_tensors) { + for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { + ggml_backend_tensor_copy(mbuf_cur.org[i], mbuf_cur.cpy[i]); + } } } } @@ -2831,14 +2866,16 @@ public: if (size > buf_size) { throw std::runtime_error("unexpectedly reached end of buffer"); } - memcpy(ptr, src, size); - ptr += size; + if (ptr != nullptr) { + memcpy(ptr, src, size); + ptr += size; + } size_written += size; buf_size -= size; } void write_tensor(ggml_tensor * tensor, size_t offset, size_t size) override { - // save the write for later during destruction + // save the write until finish(), after all state tensors are known winfos.push_back({tensor, ptr, size, offset}); } @@ -2860,6 +2897,7 @@ private: std::vector winfos; llama_memory_buffers & mbufs; + const bool copy_tensors; }; class llama_io_read_device : public llama_io_read_i { @@ -3079,8 +3117,11 @@ size_t llama_context::state_seq_get_size(llama_seq_id seq_id, llama_state_seq_fl size_t llama_context::state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, size_t size, llama_state_seq_flags flags) { std::unique_ptr io; + llama_io_write_device * io_device = nullptr; if (flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) { - io = std::make_unique(dst, size, mem_storage[seq_id]); + auto device = std::make_unique(dst, size, mem_storage[seq_id]); + io_device = device.get(); + io = std::move(device); } else { io = std::make_unique(dst, size); } @@ -3089,7 +3130,11 @@ size_t llama_context::state_seq_get_data(llama_seq_id seq_id, uint8_t * dst, siz io->write(&io_magic, sizeof(io_magic)); io->write(&seq_id, sizeof(seq_id)); - return state_seq_write_data(*io, seq_id, flags); + const size_t n = state_seq_write_data(*io, seq_id, flags); + if (io_device) { + io_device->finish(); + } + return n; } catch (const std::exception & err) { LLAMA_LOG_ERROR("%s: error saving state: %s\n", __func__, err.what()); return 0; @@ -3382,6 +3427,30 @@ llama_memory_breakdown llama_context::memory_breakdown() const { return ret; } +std::map llama_context::state_seq_device_buffer_sizes() const { + return memory ? memory->state_seq_device_buffer_sizes() + : std::map {}; +} + +bool llama_context::state_seq_reserve_device_buffers() { + if (!memory || memory->state_seq_device_buffer_sizes().empty()) { + return false; + } + + try { + for (llama_seq_id seq_id = 0; seq_id < static_cast(cparams.n_seq_max); ++seq_id) { + llama_io_write_device io(nullptr, std::numeric_limits::max(), mem_storage[seq_id], false); + memory->state_seq_write_device_layout(io); + io.finish(); + } + return true; + } catch (const std::exception & err) { + mem_storage.clear(); + LLAMA_LOG_ERROR("%s: failed to reserve device checkpoint buffers: %s\n", __func__, err.what()); + return false; + } +} + // // training // @@ -3629,6 +3698,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, @@ -3652,6 +3722,7 @@ llama_context_params llama_context_default_params() { /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, /*.ctx_other =*/ nullptr, + /*.hip_fa_force_vec =*/ false, }; return result; @@ -4325,6 +4396,15 @@ llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * c return ctx->memory_breakdown(); } +std::map llama_get_state_seq_device_buffer_sizes( + const struct llama_context * ctx) { + return ctx->state_seq_device_buffer_sizes(); +} + +bool llama_state_seq_reserve_device_buffers(struct llama_context * ctx) { + return ctx->state_seq_reserve_device_buffers(); +} + llama_context * llama_get_ctx_other(struct llama_context * ctx) { return ctx->get_cparams().ctx_other; } diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa..8327323 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -188,6 +188,8 @@ struct llama_context { void perf_reset(); llama_memory_breakdown memory_breakdown() const; + std::map state_seq_device_buffer_sizes() const; + bool state_seq_reserve_device_buffers(); // // training diff --git a/src/llama-cparams.h b/src/llama-cparams.h index b592de1..ba833f9 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -39,6 +39,7 @@ struct llama_cparams { bool offload_kqv; bool flash_attn; bool auto_fa; + bool hip_fa_force_vec; bool fused_gdn_ar; // use fused gated delta net (autoregressive) bool fused_gdn_ch; // use fused gated delta net (chunked) bool auto_fgdn; diff --git a/src/llama-ext.h b/src/llama-ext.h index 92a759b..e011d6f 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -90,6 +90,11 @@ LLAMA_API ggml_backend_dev_t llama_model_get_device(const struct llama_model * m LLAMA_API llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * ctx); +// Exact backend-buffer sizes needed for one device-resident sequence-state +// checkpoint per configured sequence. +LLAMA_API std::map llama_get_state_seq_device_buffer_sizes( + const struct llama_context * ctx); + // Set whether the context outputs nextn embeddings or not // If masked == true, output the embeddings only for the tokens with batch.logits != 0 // If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8ea441f..18b7ec9 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2585,6 +2585,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( GGML_ASSERT(n_kv_max >= 0 && n_kv_max <= INT32_MAX); ggml_flash_attn_ext_set_n_kv_max(cur, static_cast(n_kv_max)); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); + ggml_flash_attn_ext_set_force_vec(cur, cparams.hip_fa_force_vec); if (v_mla) { #if 0 diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index 06f7fd5..408982b 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -192,6 +192,14 @@ std::map llama_memory_hybrid_iswa::memory_br return mb; } +std::map llama_memory_hybrid_iswa::state_seq_device_buffer_sizes() const { + return mem_recr->state_seq_device_buffer_sizes(); +} + +void llama_memory_hybrid_iswa::state_seq_write_device_layout(llama_io_write_i & io) const { + mem_recr->state_seq_write_device_layout(io); +} + void llama_memory_hybrid_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { mem_attn->state_write(io, seq_id, flags); mem_recr->state_write(io, seq_id, flags); diff --git a/src/llama-memory-hybrid-iswa.h b/src/llama-memory-hybrid-iswa.h index c9d3f9f..02e5ce3 100644 --- a/src/llama-memory-hybrid-iswa.h +++ b/src/llama-memory-hybrid-iswa.h @@ -70,6 +70,8 @@ public: llama_pos seq_pos_max(llama_seq_id seq_id) const override; std::map memory_breakdown() const override; + std::map state_seq_device_buffer_sizes() const override; + void state_seq_write_device_layout(llama_io_write_i & io) const override; // state write/load diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381..b05cb1a 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -187,6 +187,14 @@ std::map llama_memory_hybrid::memory_breakdo return mb; } +std::map llama_memory_hybrid::state_seq_device_buffer_sizes() const { + return mem_recr->state_seq_device_buffer_sizes(); +} + +void llama_memory_hybrid::state_seq_write_device_layout(llama_io_write_i & io) const { + mem_recr->state_seq_write_device_layout(io); +} + void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_write(io, seq_id, flags); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb..5818fab 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -70,6 +70,8 @@ public: llama_pos seq_pos_max(llama_seq_id seq_id) const override; std::map memory_breakdown() const override; + std::map state_seq_device_buffer_sizes() const override; + void state_seq_write_device_layout(llama_io_write_i & io) const override; // state write/load diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 57919ac..b8db687 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -1,5 +1,6 @@ #include "llama-memory-recurrent.h" +#include "ggml-alloc.h" #include "ggml-backend.h" #include "llama-impl.h" #include "llama-io.h" @@ -425,6 +426,64 @@ std::map llama_memory_recurrent::memory_brea return ret; } +std::map llama_memory_recurrent::state_seq_device_buffer_sizes() const { + struct layout { + size_t n_tensors = 0; + std::vector> tensors; + }; + + std::map layouts; + const uint32_t n_layer = hparams.n_layer(); + + auto add_row = [&](ggml_tensor * tensor, uint32_t n_embd) { + if (tensor == nullptr) { + return; + } + const uint64_t row_size = ggml_row_size(tensor->type, n_embd); + auto * buft = ggml_backend_buffer_get_type(tensor->buffer); + auto & cur = layouts[buft]; + cur.n_tensors++; + cur.tensors.emplace_back(tensor->type, row_size / ggml_element_size(tensor)); + }; + + for (uint32_t il = 0; il < n_layer; ++il) { + add_row(r_l[il], hparams.n_embd_r()); + } + for (uint32_t il = 0; il < n_layer; ++il) { + add_row(s_l[il], hparams.n_embd_s()); + } + + std::map ret; + for (const auto & [buft, layout] : layouts) { + ggml_init_params params = { + /*.mem_size =*/ layout.n_tensors * ggml_tensor_overhead(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context_ptr ctx { ggml_init(params) }; + if (!ctx) { + throw std::runtime_error("failed to create recurrent checkpoint sizing context"); + } + for (const auto & [type, n] : layout.tensors) { + ggml_new_tensor_1d(ctx.get(), type, n); + } + + const size_t one_seq = ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft); + if (n_seq_max != 0 && one_seq > std::numeric_limits::max() / n_seq_max) { + throw std::runtime_error("recurrent checkpoint size overflow"); + } + ret[buft] = one_seq * n_seq_max; + } + return ret; +} + +void llama_memory_recurrent::state_seq_write_device_layout(llama_io_write_i & io) const { + // One logical recurrent row is the complete per-sequence checkpoint. The + // actual save may select a rollback plane, but its tensor shapes and buffer + // sizes are identical to row zero. + state_write_data(io, {{0, 1}}); +} + llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { balloc.split_reset(); @@ -822,7 +881,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-memory-recurrent.h b/src/llama-memory-recurrent.h index 4abb3f5..945e179 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -53,6 +53,8 @@ public: llama_pos seq_pos_max(llama_seq_id seq_id) const override; std::map memory_breakdown() const override; + std::map state_seq_device_buffer_sizes() const override; + void state_seq_write_device_layout(llama_io_write_i & io) const override; bool prepare(const std::vector & ubatches); diff --git a/src/llama-memory.h b/src/llama-memory.h index db82539..374d49d 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -118,6 +118,15 @@ struct llama_memory_i { virtual std::map memory_breakdown() const = 0; + // Exact backend-buffer sizes needed to keep one sequence-state checkpoint + // per configured sequence. Memory types without device checkpoint data + // return an empty map. + virtual std::map state_seq_device_buffer_sizes() const { return {}; } + + // Describe the tensor data for one device-resident sequence checkpoint. + // This is used to allocate the real checkpoint buffers before evaluation. + virtual void state_seq_write_device_layout(llama_io_write_i & io) const { GGML_UNUSED(io); } + // // state write/read // diff --git a/src/llama.cpp b/src/llama.cpp index ad8e443..e51fd10 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -47,6 +47,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: @@ -617,4 +629,3 @@ const char * llama_print_system_info(void) { return s.c_str(); } - diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp index 80260a6..12c3339 100644 --- a/src/models/minimax-m3.cpp +++ b/src/models/minimax-m3.cpp @@ -192,6 +192,7 @@ ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_fa( ggml_tensor * o = ggml_flash_attn_ext(ctx0, q, k, v, mask, kq_scale, hparams.f_max_alibi_bias, 0.0f); ggml_flash_attn_ext_set_prec(o, GGML_PREC_F32); + ggml_flash_attn_ext_set_force_vec(o, cparams.hip_fa_force_vec); cb(o, "msa_fattn", il); // [D, Gp, R, C] -> [D, Gp, C, R] -> [n_embd, T] diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c46377c..5f6b4c8 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -198,6 +198,13 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-llama-archs.cpp) + llama_test( + test-llama-archs + NAME test-mtp-ubatch-sync + LABEL main + ARGS --test-mtp-ubatch-sync + ) + set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index e090763..8ac51a9 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 @@ -186,6 +187,23 @@ static void test(void) { argv = {"binary_name", "-sm", "hello"}; assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); + { + common_params hip_fa_params; + assert(!hip_fa_params.hip_fa_force_vec); + assert(!llama_context_default_params().hip_fa_force_vec); + + argv = {"binary_name", "-m", "model.gguf", "--hip-fa-force-vec", "off"}; + assert(common_params_parse(argv.size(), list_str_to_char(argv).data(), hip_fa_params, LLAMA_EXAMPLE_COMMON)); + assert(!hip_fa_params.hip_fa_force_vec); + + argv = {"binary_name", "-m", "model.gguf", "--hip-fa-force-vec", "on"}; + assert(common_params_parse(argv.size(), list_str_to_char(argv).data(), hip_fa_params, LLAMA_EXAMPLE_COMMON)); + assert(hip_fa_params.hip_fa_force_vec); + + argv = {"binary_name", "-m", "model.gguf", "--hip-fa-force-vec", "maybe"}; + assert(!common_params_parse(argv.size(), list_str_to_char(argv).data(), hip_fa_params, LLAMA_EXAMPLE_COMMON)); + } + { common_params penalty_params; assert(penalty_params.sampling.penalty_last_n == 64); @@ -254,6 +272,136 @@ 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); + argv = {"binary_name", "--spec-draft-adaptive"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SPECULATIVE)); + assert(params.speculative.draft.adaptive); + + { + common_params params_mtp; + argv = {"binary_name", "-m", "abc.gguf", "--spec-draft-n-max", "5", "--spec-mtp-cr-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-cr-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-cr-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-cr-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-cr-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-cr-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-cr-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"); + } + } + { common_params synth_params; argv = {"binary_name", "--spec-synth-len", "3.4"}; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 0f3d1c7..463c0a9 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -65,7 +65,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } static void usage(char ** argv) { - printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help]\n", argv[0]); + printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help] [--test-mtp-ubatch-sync]\n", argv[0]); } static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ @@ -79,7 +79,7 @@ static std::vector get_tokens(const uint32_t n_tokens, const uint32 return ret; } -static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { +static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const bool mtp = false) { gguf_context_ptr ret(gguf_init_empty()); llama_model_saver ms(arch, ret.get()); const uint32_t n_ctx = 256; @@ -146,6 +146,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_FEATURES_LENGTH, n_embd); ms.add_kv(LLM_KV_BLOCK_COUNT, n_layer); ms.add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, uint32_t(1)); + if (mtp) { + ms.add_kv(LLM_KV_NEXTN_PREDICT_LAYERS, uint32_t(1)); + } if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { std::vector n_ff_per_layer; @@ -420,6 +423,83 @@ static std::pair get_model_and_ctx( return std::make_pair(std::move(model), std::move(lctx)); } +static bool mtp_sync_test_decode(llama_model * model, uint32_t n_ubatch) { + const int32_t n_tokens = 4; + const int32_t n_embd = llama_model_n_embd_out(model); + + llama_context_params ctx_params = llama_context_default_params(); + ctx_params.n_ctx = 8; + ctx_params.n_batch = n_tokens; + ctx_params.n_ubatch = n_ubatch; + ctx_params.n_threads = 4; + ctx_params.n_threads_batch = 4; + ctx_params.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + + llama_context_ptr ctx(llama_init_from_model(model, ctx_params)); + if (!ctx) { + throw std::runtime_error("failed to create MTP context"); + } + + std::vector token(n_tokens); + std::vector embd((size_t) n_tokens * n_embd, 1.0e-2f); + std::vector pos(n_tokens); + std::vector n_seq_id(n_tokens, 1); + std::vector seq_id_data(n_tokens, 0); + std::vector seq_id(n_tokens); + std::vector logits(n_tokens, 0); + + for (int32_t i = 0; i < n_tokens; ++i) { + token[i] = i; + pos[i] = i; + seq_id[i] = &seq_id_data[i]; + } + logits.back() = 1; + + llama_batch batch = { + /*.n_tokens =*/ n_tokens, + /*.token =*/ token.data(), + /*.embd =*/ embd.data(), + /*.pos =*/ pos.data(), + /*.n_seq_id =*/ n_seq_id.data(), + /*.seq_id =*/ seq_id.data(), + /*.logits =*/ logits.data(), + }; + + llama_perf_context_reset(ctx.get()); + const int32_t ret = llama_decode(ctx.get(), batch); + if (ret != 0) { + throw std::runtime_error("failed to decode MTP batch"); + } + + // synchronize() accounts the complete queued batch immediately. Without a + // synchronization, llama_perf_context() reports its minimum placeholder of 1. + return llama_perf_context(ctx.get()).n_p_eval >= n_tokens; +} + +static int test_mtp_ubatch_sync(const size_t seed) { + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_QWEN35, false, true); + llama_model_params model_params = llama_model_default_params(); + model_params.progress_callback = silent_model_load_progress; + model_params.load_mtp = true; + + size_t tmp = seed; + llama_model_ptr model(llama_model_init_from_user(gguf_ctx.get(), set_tensor_data, &tmp, model_params)); + if (!model) { + throw std::runtime_error("failed to create MTP model"); + } + + if (!mtp_sync_test_decode(model.get(), 2)) { + fprintf(stderr, "MTP ubatches were not synchronized\n"); + return 1; + } + if (mtp_sync_test_decode(model.get(), 4)) { + fprintf(stderr, "single MTP ubatch was synchronized\n"); + return 1; + } + + return 0; +} + static std::vector get_logits( llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false) { const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); @@ -817,6 +897,7 @@ int main(int argc, char ** argv) { llm_arch arch = LLM_ARCH_UNKNOWN; size_t seed = rd(); std::string out; + bool test_mtp_sync = false; int verbosity = LOG_LEVEL_ERROR; @@ -862,6 +943,9 @@ int main(int argc, char ** argv) { return 1; } } + if (strcmp(argv[i], "--test-mtp-ubatch-sync") == 0) { + test_mtp_sync = true; + } } printf("%s: using seed %zu\n", __func__, seed); @@ -869,6 +953,9 @@ int main(int argc, char ** argv) { if (!out.empty()) { return save_models(arch, seed, verbosity, out); } + if (test_mtp_sync) { + return test_mtp_ubatch_sync(seed); + } return test_backends(arch, seed, verbosity); } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what()); diff --git a/tests/test-recurrent-state-rollback.cpp b/tests/test-recurrent-state-rollback.cpp index c6f599e..ae97295 100644 --- a/tests/test-recurrent-state-rollback.cpp +++ b/tests/test-recurrent-state-rollback.cpp @@ -8,15 +8,112 @@ #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) { + // Some recurrent layouts compact the surviving sequences into one + // device range. That is valid and does not require host fallback. + fprintf(stderr, "%s : checkpoint remained device-contiguous\n", __func__); + } else { + fprintf(stderr, "%s : fragmented checkpoint fell back to host storage\n", __func__); + } + } + if (ok) { + // load_tgt() must use the actual storage mode recorded when the + // checkpoint was saved, whether the layout stayed contiguous or fell + // back to host storage. + ckpt.load_tgt(ctx, -1, flags); + } + + llama_free(ctx); + return ok; +} + +static bool test_preallocated_device_checkpoints( + const common_params & params, + llama_model * model, + const std::vector & tokens) { + constexpr int32_t n_seq = 2; + llama_context * ctx = make_ctx(params, model, 1, n_seq); + if (ctx == nullptr) { + fprintf(stderr, "%s : failed to initialize multi-sequence context\n", __func__); + return false; + } + + bool ok = llama_state_seq_reserve_device_buffers(ctx); + if (!ok) { + // Architectures without recurrent tensor rows have no device layout to + // reserve. The server handles this by using host checkpoints. + fprintf(stderr, "%s : no device checkpoint layout; using host checkpoints\n", __func__); + llama_free(ctx); + return true; + } + + llama_batch batch = llama_batch_init(n_seq, 0, n_seq); + for (llama_seq_id seq_id = 0; seq_id < n_seq; ++seq_id) { + common_batch_add(batch, tokens[seq_id % tokens.size()], 0, { seq_id }, true); + } + if (ok) { + ok = llama_decode(ctx, batch) == 0; + } + llama_batch_free(batch); + + constexpr llama_state_seq_flags flags = + LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE; + std::vector checkpoints(n_seq); + for (llama_seq_id seq_id = 0; ok && seq_id < n_seq; ++seq_id) { + checkpoints[seq_id].update_tgt(ctx, seq_id, flags); + if (!checkpoints[seq_id].data_tgt_on_device) { + fprintf(stderr, "%s : sequence %d did not reuse its preallocated device checkpoint\n", + __func__, seq_id); + ok = false; + } + } + for (llama_seq_id seq_id = 0; ok && seq_id < n_seq; ++seq_id) { + checkpoints[seq_id].load_tgt(ctx, seq_id, 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) { @@ -207,6 +304,156 @@ static bool test_multi_seq_split_replay(const common_params & params, llama_mode return true; } +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"); @@ -237,6 +484,26 @@ 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_preallocated_device_checkpoints(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) { @@ -251,12 +518,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) { @@ -265,10 +526,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(); @@ -333,6 +590,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 952d31e..37e5021 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -262,6 +262,8 @@ 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-cr-depth N` | MTP Compact Rollback depth; lower values save memory but replay accepted tokens after deep rejection (default: `--spec-draft-n-max`)
(env: LLAMA_ARG_SPEC_MTP_CR_DEPTH) | +| `--spec-draft-adaptive` | size each draft from measured acceptance rather than always drafting `--spec-draft-n-max` (default: false)
(env: LLAMA_ARG_SPEC_DRAFT_ADAPTIVE) | | `--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-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) | diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 2ac98b6..5bdd899 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 6c681a2..9712090 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -359,6 +359,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 f78cfb3..6adcdf3 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -256,6 +256,7 @@ struct server_slot { std::vector spec_i_batch; common_prompt_checkpoint spec_ckpt; bool spec_is_replay = false; + int32_t spec_n_accepted_observed = -1; std::mt19937 spec_synth_rng; // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state @@ -370,6 +371,7 @@ struct server_slot { SLT_DBG(*this, "%s", "\n"); spec_is_replay = false; + spec_n_accepted_observed = -1; last_nl_pos = 0; generated_text = ""; @@ -680,6 +682,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); } @@ -892,6 +900,7 @@ private: common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; + bool spec_mtp_device_checkpoint = false; common_speculative_ptr spec; @@ -1088,6 +1097,90 @@ private: } } + // The on-device recurrent speculative checkpoint is allocated lazily, so reserve + // the additional recurrent-state group before the joint target / MTP fit. + 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; + const size_t checkpoint = dmd_rs[j].checkpoint; + if (checkpoint == 0 && delta != 0) { + throw std::runtime_error("recurrent checkpoint sizing produced no device allocation"); + } + params_base.fit_params_target[j] += checkpoint; + total += checkpoint; + SRV_INF("[spec] recurrent checkpoint fit reservation: device %s, %.2f MiB " + "(exact backend allocation; recurrent-plane delta %.2f MiB, context %.2f -> %.2f MiB)\n", + ggml_backend_dev_name(devs_rs[j]), checkpoint / (1024.0 * 1024.0), + 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; + } + } + } + // note: the draft / MTP context is fitted together with the target model, see common_fit_extra_model // attach a progress callback @@ -1279,6 +1372,16 @@ private: model_dft = nullptr; } + spec_mtp_device_checkpoint = false; + if (spec && ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS) { + spec_mtp_device_checkpoint = llama_state_seq_reserve_device_buffers(ctx_tgt); + if (spec_mtp_device_checkpoint) { + SRV_INF("%s", "[spec] reserved device recurrent checkpoints before evaluation\n"); + } else { + SRV_WRN("%s", "[spec] device recurrent checkpoint reservation failed; using host checkpoints\n"); + } + } + if (!spec && params_base.speculative.has_synth()) { SRV_ERR("%s", "synthetic acceptance requires an initialized speculative decoding context\n"); return false; @@ -3060,9 +3163,23 @@ 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 && spec_mtp_device_checkpoint) { + // Phase 1 proof of concept: avoid copying the recurrent fallback checkpoint + // through host memory on every speculative round. The device buffer is + // allocated during load_model(), before evaluation can consume the + // fitted headroom reserved for the recurrent-state checkpoint. + 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); + + if ((ckpt_flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) && !ckpt.data_tgt_on_device) { + spec_mtp_device_checkpoint = false; + SRV_WRN("%s", "[spec] device recurrent checkpoint save failed; using host checkpoints for the rest of this server run\n"); + } //const int64_t t_total = ggml_time_us() - t_start; //printf("checkpoint total: %f ms\n", t_total / 1000.0); @@ -3911,6 +4028,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) { @@ -3918,15 +4039,32 @@ private: SLT_INF(slot, "accepted %2zu/%2zu draft tokens (restore checkpoint)\n", accepted.size() - 1, slot.spec_draft.size()); } + // Preserve the original observation for adaptive draft sizing. The + // replay batch also contains the target replacement token, so its + // physical acceptance count is one larger than the measurement that + // selected the rollback path. + if (!slot.spec_is_replay) { + slot.spec_n_accepted_observed = static_cast(accepted.size() - 1); + } + // partial acceptance is not supported by the context -> truncate the draft and restore the state 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); @@ -3945,7 +4083,12 @@ private: SLT_INF(slot, "accepted %2zu/%2zu draft tokens\n", accepted.size() - 1, n_draft); } - common_speculative_accept(spec.get(), slot.id, accepted.size() - 1); + common_speculative_accept( + spec.get(), + slot.id, + accepted.size() - 1, + slot.spec_is_replay ? slot.spec_n_accepted_observed : -1); + slot.spec_n_accepted_observed = -1; slot.spec_draft = std::move(accepted); } diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index 826aef2..268df76 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_cr_depth: int | None = None spec_synth_len: float | None = None spec_synth_rates: List[float] | None = None no_ui: bool | None = None @@ -247,6 +248,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_cr_depth: + server_args.extend(["--spec-mtp-cr-depth", self.spec_mtp_cr_depth]) if self.spec_synth_len is not None: server_args.extend(["--spec-synth-len", self.spec_synth_len]) if self.spec_synth_rates is not None: