upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-23 11:10:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix pending state

Browse Source
This commit is contained in:
Aman Gupta
2026-05-14 14:32:42 +08:00
parent 4e732e0a6c
commit 9a3a48722a
7 changed files with 64 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
common/common.cpp
View File

@@ -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;

64
common/speculative.cpp
View File

@@ -394,6 +394,16 @@ struct common_speculative_state_draft_mtp : public common_speculative_impl {
std::vector<int32_t> i_batch_beg;
std::vector<int32_t> 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<std::vector<float>> verify_h;
std::vector<int32_t> 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<uint16_t> 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<int32_t>(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);

3
examples/speculative-simple/speculative-simple.cpp
View File

@@ -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) {

2
include/llama.h
View File

@@ -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");

1
src/llama-context.cpp
View File

@@ -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",

4
src/llama-memory-recurrent.cpp
View File

@@ -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));
}

5
src/models/delta-net-base.cpp
View File

@@ -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);