diff --git a/common/common.cpp b/common/common.cpp index 6af85e7ab7..6030ead189 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1421,6 +1421,7 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) { } if (llama_n_rs_seq(ctx) > 0) { + LOG_INF("%s: the context supports bounded partial sequence removal\n", __func__); res = COMMON_CONTEXT_SEQ_RM_TYPE_PART_BOUNDED; goto done; } @@ -1496,15 +1497,14 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.n_ctx = params.n_ctx; cparams.n_seq_max = params.n_parallel; { - // TODO: add for MTP - bool has_spec = params.speculative.has_dft(); - for (auto t : params.speculative.types) { - if (t != COMMON_SPECULATIVE_TYPE_NONE) { - has_spec = true; - break; - } - } - cparams.n_rs_seq = has_spec ? (uint32_t) params.speculative.draft.n_max : 0u; + // Since MTP has a low number of draft tokens, enable recurrent checkpointing + // for hybrid attn models + // TODO: figure out how to make it place nicely with other speculative techniques + bool has_mtp = std::any_of(params.speculative.types.begin(), + params.speculative.types.end(), [&](auto t) { + return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP; + }); + cparams.n_rs_seq = has_mtp ? (uint32_t) params.speculative.draft.n_max : 0u; } cparams.n_batch = params.n_batch; cparams.n_ubatch = params.n_ubatch; diff --git a/common/speculative.cpp b/common/speculative.cpp index 3914a5e245..b932c86ded 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -394,6 +394,16 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { std::vector i_batch_beg; std::vector i_batch_end; + // Hidden rows from the most recent target verification batch, grouped by seq. + // Row 0 corresponds to the sampled token, row N to the Nth accepted draft token. + std::vector> verify_h; + std::vector verify_h_rows; + + // Per-seq draft length from the last draft() call, used in accept() to + // roll back ctx_dft's recurrent state past the AR draft's redundant + // pre-advancement before process() mirrored the verify batch. + std::vector last_n_drafted; + common_speculative_state_draft_mtp(const common_params_speculative & params, uint32_t n_seq) : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq) , params(params.draft) @@ -414,7 +424,7 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { for (auto & s : smpls) { common_params_sampling sparams; sparams.no_perf = false; - sparams.top_k = 10; + sparams.top_k = 1; sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K }; s.reset(common_sampler_init(llama_get_model(ctx_dft), sparams)); } @@ -426,6 +436,11 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { i_batch_beg.assign(n_seq, -1); i_batch_end.assign(n_seq, -1); + + verify_h.assign(n_seq, {}); + verify_h_rows.assign(n_seq, 0); + + last_n_drafted.assign(n_seq, 0); } ~common_speculative_state_draft_mtp() override { @@ -535,8 +550,17 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { continue; } - const float * h_last = llama_get_embeddings_pre_norm_ith(ctx_tgt, i_batch_end[seq_id]); - std::memcpy(pending_h[seq_id].data(), h_last, row_bytes); + const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1; + verify_h_rows[seq_id] = n_rows; + verify_h[seq_id].resize((size_t) n_rows * n_embd); + + for (int32_t i = 0; i < n_rows; ++i) { + const float * h = llama_get_embeddings_pre_norm_ith(ctx_tgt, i_batch_beg[seq_id] + i); + std::memcpy(verify_h[seq_id].data() + (size_t) i * n_embd, h, row_bytes); + } + + std::memcpy(pending_h[seq_id].data(), + verify_h[seq_id].data() + (size_t) (n_rows - 1) * n_embd, row_bytes); } return true; @@ -606,14 +630,6 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { // add drafted token for each sequence const llama_token id = cur_p->data[0].id; - // only collect very high-confidence draft tokens - if (cur_p->data[0].p < params.p_min) { - drafting[seq_id] = false; - n_drafting--; - - continue; - } - common_sampler_accept(smpl, id, true); auto & dp = dparams.at(seq_id); @@ -621,8 +637,7 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { result.push_back(id); - if ((params.n_max <= (int) result.size()) || - (dp.n_max > 0 && dp.n_max <= (int) result.size())) { + if (params.n_max <= (int) result.size()) { drafting[seq_id] = false; n_drafting--; continue; @@ -646,7 +661,8 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { ++i; } - for (auto & dp : dparams) { + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + auto & dp = dparams[seq_id]; if (!dp.drafting) { continue; } @@ -654,10 +670,24 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl { if (dp.result->size() < (size_t) params.n_min) { dp.result->clear(); } + + last_n_drafted[seq_id] = (uint16_t) dp.result->size(); } } - void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/) override { + void accept(llama_seq_id seq_id, uint16_t n_accepted) override { + if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { + return; + } + + const int32_t n_rows = verify_h_rows[seq_id]; + if (n_rows <= 0) { + return; + } + + const int32_t i_h = std::min(n_accepted, n_rows - 1); + const size_t row_bytes = (size_t) n_embd * sizeof(float); + std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes); } bool need_embd() const override { @@ -1460,10 +1490,6 @@ void common_speculative_draft(common_speculative * spec) { } void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, uint16_t n_accepted) { - if (n_accepted == 0) { - return; - } - common_speculative_impl * impl = spec->impl_last[seq_id]; GGML_ASSERT(impl); diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index 5325bcc9e3..6848de988d 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -82,7 +82,8 @@ int main(int argc, char ** argv) { } // check if the context supports partial sequence removal - const bool use_ckpt_tgt = (common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL); + const bool use_ckpt_tgt = (common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL) + || (common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_PART_BOUNDED); const bool use_ckpt_dft = (common_context_can_seq_rm(ctx_dft.get()) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL); if (use_ckpt_tgt) { diff --git a/include/llama.h b/include/llama.h index 0f711ab29e..75095b22d0 100644 --- a/include/llama.h +++ b/include/llama.h @@ -537,7 +537,7 @@ extern "C" { LLAMA_API uint32_t llama_n_batch (const struct llama_context * ctx); LLAMA_API uint32_t llama_n_ubatch (const struct llama_context * ctx); LLAMA_API uint32_t llama_n_seq_max (const struct llama_context * ctx); - LLAMA_API uint32_t llama_n_rs_seq (const struct llama_context * ctx); + LLAMA_API uint32_t llama_n_rs_seq (const struct llama_context * ctx); DEPRECATED(LLAMA_API int32_t llama_n_ctx_train(const struct llama_model * model), "use llama_model_n_ctx_train instead"); DEPRECATED(LLAMA_API int32_t llama_n_embd (const struct llama_model * model), "use llama_model_n_embd instead"); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4a000f59f1..a481757e38 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -225,6 +225,7 @@ llama_context::llama_context( 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); + LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq); if (cparams.n_ctx_seq < hparams.n_ctx_train) { LLAMA_LOG_WARN("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n", diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 1913e9414a..084c5d9ea4 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -120,8 +120,8 @@ llama_memory_recurrent::llama_memory_recurrent( const size_t memory_size_r = size_r_bytes(); const size_t memory_size_s = size_s_bytes(); - LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs), R (%s): %7.2f MiB, S (%s): %7.2f MiB\n", __func__, - (float)(memory_size_r + memory_size_s) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs %2u rs_seq), R (%s): %7.2f MiB, S (%s): %7.2f MiB\n", __func__, + (float)(memory_size_r + memory_size_s) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq, ggml_type_name(type_r), (float)memory_size_r / (1024.0f * 1024.0f), ggml_type_name(type_s), (float)memory_size_s / (1024.0f * 1024.0f)); } diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index 081f490c5a..852d5bc88e 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -547,7 +547,10 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( } const int64_t D = S_v * S_v * H_v; - const int64_t K = (int64_t) cparams.n_rs_seq; + // Memory has 1 + n_rs_seq slots (slot 0 = current, slots 1..n_rs_seq = rollback distances). + // The snapshot buffer must match — otherwise the deepest rollback (= n_rs_seq) reads + // uninitialized memory and corrupts the recurrent state. + const int64_t K = (int64_t) cparams.n_rs_seq + 1; ggml_tensor * state_3d = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, K, n_seqs); ggml_tensor * slot_0 = ggml_view_2d(ctx0, state_3d, D, n_seqs, state_3d->nb[2], 0);