Skip to content

Commit 75dc25e

Browse files
SamuelOliveiradsSamuelOliveirads
committed
mtp-batch (wip): organize batch for mtp cache
1 parent 3da7e7f commit 75dc25e

4 files changed

Lines changed: 33 additions & 46 deletions

File tree

‎common/speculative.cpp‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -373,7 +373,7 @@ llama_token mtp_speculative_gen_draft(
373373
if (!smpl) {
374374
return -1;
375375
}
376-
const float * draft_input_hidden_state = llama_get_embeddings_ith(ctx, -1);
376+
const float * draft_input_hidden_state = llama_get_embeddings(ctx);
377377
llama_set_draft_input_hidden_state(ctx, draft_input_hidden_state);
378378
LOG_INF("[DEBUG-DRAFT-STATE] Main model final embd pointer: %p, State being used for draft: %p\n",
379379
(void*)llama_get_embeddings(ctx), (void*)draft_input_hidden_state);
@@ -413,17 +413,24 @@ llama_token mtp_speculative_gen_draft(
413413
}
414414

415415

416-
void mtp_update_kv_cache(struct llama_context * ctx, std::vector<mtp_kv_update_data>& tokens) {
416+
void mtp_update_kv_cache(struct llama_context * ctx, std::vector<mtp_kv_update_data>& tokens, const char* tag) {
417417
if (tokens.empty()) {
418418
return;
419419
}
420420

421421
const size_t n_to_process = tokens.size();
422-
423-
LOG_DBG(
424-
"[MTP BATCHING] mtp_update_kv_cache call for %zu tokens.\n",
425-
n_to_process
426-
);
422+
std::string details_str;
423+
for (size_t i = 0; i < std::min((size_t)5, n_to_process); ++i) {
424+
details_str += " {id: " + std::to_string(tokens[i].id) + ", pos: " + std::to_string(tokens[i].n_past) + "}";
425+
}
426+
LOG_INF("[MTP-UPDATE|%s] Updating %zu tokens. Details:%s ...\n", tag, n_to_process, details_str.c_str());
427+
428+
// LOG_INF("[DEBUG-CHUNK] Warming up MTP model chunk. Batch size: %zu\n", n_to_process);
429+
// std::string positions_str;
430+
// for (size_t i = 0; i < std::min((size_t)5, n_to_process); ++i) {
431+
// positions_str += std::to_string(tokens[i].n_past) + " ";
432+
// }
433+
// LOG_INF("[DEBUG-CHUNK] MTP warm-up positions: %s...\n", positions_str.c_str());
427434
llama_batch mtp_batch = llama_batch_init(n_to_process, 0, 1);
428435

429436
for (size_t i = 0; i < n_to_process; ++i) {

‎common/speculative.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,4 +49,4 @@ llama_tokens common_speculative_gen_draft(
4949
const llama_tokens & prompt,
5050
llama_token id_last);
5151

52-
void mtp_update_kv_cache(struct llama_context * ctx, std::vector<mtp_kv_update_data>& tokens);
52+
void mtp_update_kv_cache(struct llama_context * ctx, std::vector<mtp_kv_update_data>& tokens, const char* tag);

‎src/llama-context.cpp‎

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -778,13 +778,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
778778
ggml_tensor* hidden_states_input = ggml_get_tensor(res->get_ctx(), target_tensor_name);
779779

780780
const float * source_hidden_state = nullptr;
781-
if (do_mtp_kv_update) {
782-
// Cache warming uses the entire embeddings buffer
783-
source_hidden_state = this->embd;
784-
} else {
785-
// Draft generation uses the specific state
786-
source_hidden_state = this->draft_input_hidden_state;
787-
}
781+
source_hidden_state = this->draft_input_hidden_state;
788782

789783
if (source_hidden_state != nullptr && hidden_states_input != nullptr) {
790784
ggml_backend_tensor_set(hidden_states_input, source_hidden_state, 0, ggml_nbytes(hidden_states_input));
@@ -1149,14 +1143,13 @@ int llama_context::decode(const llama_batch & batch_inp) {
11491143
if (do_mtp_kv_update) {
11501144
LLAMA_LOG_WARN("[DEBUG-MTP-UPDATE] MTP KV Update ubatch: n_tokens=%d\n", ubatch.n_tokens);
11511145
std::string positions_str;
1152-
for (int i = 0; i < ubatch.n_tokens; ++i) {
1146+
for (int i = 0; i < std::min((uint32_t)5, ubatch.n_tokens); ++i) {
11531147
positions_str += std::to_string(ubatch.pos[i]) + " ";
11541148
}
1155-
LLAMA_LOG_WARN("[DEBUG-MTP-UPDATE] Positions: %s\n", positions_str.c_str());
1149+
LLAMA_LOG_WARN("[DEBUG-MTP-UPDATE] Positions: %s...\n", positions_str.c_str());
11561150
}
11571151
ggml_status status;
11581152
const auto * res = process_ubatch(ubatch, LLM_GRAPH_TYPE_DECODER, mctx.get(), status, do_mtp_kv_update, use_mtp_head);
1159-
11601153
if (!res) {
11611154
// the last ubatch failed or was aborted -> remove all positions of that ubatch from the KV cache
11621155
llama_pos pos_min[LLAMA_MAX_SEQ];

‎tools/server/server.cpp‎

Lines changed: 15 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -3387,14 +3387,8 @@ struct server_context {
33873387
slot.n_prompt_tokens_processed += n_pos;
33883388
}
33893389

3390-
const size_t n_to_log = slot.mtp_kv_update_batch.size();
3391-
if (n_to_log > 0) {
3392-
SLT_INF(slot,
3393-
"DEBUG-KV-REQ Cache Warm-up: Requesting KV update for %zu tokens. Positions: %d ... %d\n",
3394-
n_to_log,
3395-
slot.mtp_kv_update_batch.front().n_past,
3396-
slot.mtp_kv_update_batch.back().n_past
3397-
);
3390+
if (slot.has_mtp) {
3391+
slot.mtp_kv_update_batch.clear();
33983392
}
33993393
// add prompt tokens for processing in the current batch
34003394
while (slot.n_past < slot.n_prompt_tokens && batch.n_tokens < n_batch) {
@@ -3484,6 +3478,7 @@ struct server_context {
34843478
batch.seq_id + i,
34853479
batch.logits + i,
34863480
};
3481+
LOG_INF("\n[DEBUG-CHUNK] Processing main model chunk. Batch size: %d\n", n_tokens);
34873482

34883483
const int ret = llama_decode(ctx, batch_view);
34893484

@@ -3525,16 +3520,18 @@ struct server_context {
35253520

35263521
continue; // continue loop of n_batch
35273522
}
3523+
3524+
// This should only trigger on a non-empty update batch once, after prompt processing but not during token generation
3525+
// Aquece o cache MTP para os pedaços do prompt que acabaram de ser processados.
3526+
// Esta lógica SÓ deve ser executada durante o processamento do prompt.
35283527
for (auto & slot : slots) {
3529-
if (slot.has_mtp && slot.n_past == slot.n_prompt_tokens) {
3530-
SLT_INF(slot, "Prompt processing finished. Warming up MTP KV cache for %d tokens.\n", slot.n_prompt_tokens);
3531-
slot.mtp_kv_update_batch.clear();
3532-
3533-
for (int j = 0; j < slot.n_prompt_tokens; ++j) {
3534-
slot.mtp_kv_update_batch.push_back({ slot.prompt_tokens[j], (llama_pos)j, j });
3535-
}
3536-
3537-
mtp_update_kv_cache(ctx, slot.mtp_kv_update_batch);
3528+
if (slot.state == SLOT_STATE_PROCESSING_PROMPT && slot.has_mtp && !slot.mtp_kv_update_batch.empty()) {
3529+
SLT_INF(slot, "DEBUG-KV-REQ: Warming up MTP cache for prompt chunk of size %zu. Positions: %d ... %d\n",
3530+
slot.mtp_kv_update_batch.size(),
3531+
slot.mtp_kv_update_batch.front().n_past,
3532+
slot.mtp_kv_update_batch.back().n_past
3533+
);
3534+
mtp_update_kv_cache(ctx, slot.mtp_kv_update_batch, "PROMPT_WARMUP");
35383535
}
35393536
}
35403537

@@ -3581,11 +3578,6 @@ struct server_context {
35813578

35823579
common_sampler_accept(slot.smpl, id, true);
35833580

3584-
// This should only trigger on a non-empty update batch once, after prompt processing but not during token generation
3585-
//if (slot.has_mtp) {
3586-
// mtp_update_kv_cache(ctx, slot.mtp_kv_update_batch);
3587-
//}
3588-
35893581
slot.n_decoded += 1;
35903582

35913583
const int64_t t_current = ggml_time_us();
@@ -3670,11 +3662,6 @@ struct server_context {
36703662
draft = common_speculative_gen_draft(slot.spec, params_spec, cached_text_tokens, id);
36713663
}
36723664

3673-
//llama_token draft_id = mtp_speculative_gen_draft(slot.smpl, ctx, id, slot.n_past, slot.last_tok_idx);
3674-
//llama_tokens draft;
3675-
//draft.reserve(1);
3676-
//draft.push_back(draft_id);
3677-
36783665
// ignore small drafts
36793666
if (slot.params.speculative.n_min > (int)draft.size()) {
36803667
SLT_DBG(slot, "ignoring small draft: %d < %d\n", (int)draft.size(), slot.params.speculative.n_min);
@@ -3706,7 +3693,7 @@ struct server_context {
37063693
for (int32_t i = 0; i < ids.size(); ++i) {
37073694
slot.mtp_kv_update_batch.push_back({ ids[i], slot.n_past + i, i });
37083695
}
3709-
mtp_update_kv_cache(ctx, slot.mtp_kv_update_batch);
3696+
mtp_update_kv_cache(ctx, slot.mtp_kv_update_batch, "GEN_ACCEPTED");
37103697
}
37113698

37123699
slot.n_past += ids.size();

0 commit comments

Comments
 (0)