spec : parallel drafting support

[ggerganov]

Overview

cont #22787 (rebased on top of it)

With these changes:

• The draft context can generate speculative drafts for multiple sequences in parallel
• The draft "sees" the multimodal data
• Can chain multiple speculative types, f.ex --spec-type ngram-mod,mtp

server_context

• Single common_speculative context for all slots, capable of handling multiple sequence ids
• Extract the drafting logic from server_slot::update_batch() and parallelize it across the active slots
• Disable adjusting of speculative decoding params (such as n_min, n_max, p_min) from the HTTP API

common/speculative

• Rework the API to support drafting for multiple sequences at once
• Add struct common_speculative_draft_params
• Update draft-model based implementation to decode multiple draft sequences in parallel
• Add common_speculative_process(spec, batch) for feeding the prompt tokens through the speculative context
• Remove common_speculative_n_max
• Remove common_speculative_n_min

llama_context

• Avoid extra backend buffer allocations during llama_state_seq_get_data_ext() with the LLAMA_STATE_SEQ_FLAGS_ON_DEVICE flag (this can be merged directly to master)

common/args

• spec : allow for multiple spec types (chains of speculators) #22546

TODOs

• Clean-up

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: NO

[ggerganov]

@am17an @ruixiang63 This refactor is ready. I also went ahead and introduced the common_speculative_process() API. We can now write specialized logic for the different speculative decoding methods in common/speculative.cpp by implementing the process() method of each different implementations. The most basic one is the draft-model based one - it simply processes the same batch as the target context, without any extra data from it:

llama.cpp/common/speculative.cpp

Lines 227 to 241 in db8e326

	bool process(const llama_batch & batch) override {
	auto * ctx_dft = params.ctx_dft;
	const int ret = llama_decode(ctx_dft, batch);
	if (ret != 0) {
	LOG_ERR("%s: failed to decode draft batch, ret = %d\n", __func__, ret);
	return false;
	}
	return true;
	}

I'll start looking next into hooking the target embeddings. I'll probably start from the MTP PR #22673 as it seems to be the easiest speculative approach. The Eagle3 has some extra complexities:

• More target features to extract
• Encoder step
• Different vocabs

So I think after we fit the MTP implementation into the new speculative architecture, it will be easier to build the Eagle3 from there.

I'm already having some ideas how to efficiently move the target embeddings between the contexts. Inside the process() method, we can analyze the contents of the batch and decide which embeddings to store for later (i.e. the embeddings for the last token of each sequence in the batch) and which to copy directly to the draft context for processing with the current batch. Anyway, I think I am starting to wrap my head around it, but it still not 100% clear.

If you are prototyping stuff, it would be great to base it on top of this branch. There are comments in the code to clarify the logic, but if something is unclear - let me know.

[ggerganov]

I'm already having some ideas how to efficiently move the target embeddings between the contexts.

For now, we should definitely focus on initial implementation that uses host roundtrips. The device-to-device version is possible, but it requires first a refactoring that introduces the concept of "on-device input/output batch embeddings" in the llama_context. I now have quite a good idea how to do this - will try to spec it soon. But IMO we should do this in next iterations as it's a relatively big change.

For the initial implementation, my current suggestion is to mirror the existing llama_get_embeddings_* API in order to at least utilize the host pinned memory. This should perform better than the current direct calls to ggml_backend_tensor_get() used in #22673. Basically, add respective "pre-norm output embedding" calls in llama-ext for the following functions:

// llama-ext.h
// (not in llama.h since this is a stop-gap solution to get something mergeable)

    // mirrors:
    // LLAMA_API void llama_set_embeddings(struct llama_context * ctx, bool embeddings);
    LLAMA_API void llama_set_embeddings_pre_norm(struct llama_context * ctx, bool value);

    // mirrors:
    // LLAMA_API float * llama_get_embeddings(struct llama_context * ctx);
    LLAMA_API float * llama_get_embeddings_pre_norm(struct llama_context * ctx);

    // mirrors:
    // LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i);
    LLAMA_API float * llama_get_embeddings_pre_norm_ith(struct llama_context * ctx, int32_t i);

The llama_context should reserve extra pinned-host buffer space for the pre-norm embeddings (see llama_context::output_reserve()), async extract after every ubatch (see llama_context::decode(), .., i.e. the entire logic mirrors the existing "output embeddings" logic.

With this API, utlize it to extract both the hidden target embeddings and the MTP-generated embeddings. There is not need to distinguish explicitly between the 2 types of embeddings. I.e. this code:

    // MTP related inputs/outputs
    ggml_tensor * t_h_pre_norm  = nullptr; // [n_embd, n_outputs] hidden state required for MTP
    ggml_tensor * t_mtp_out     = nullptr; // [n_embd, n_tokens]

Should be just:

    ggml_tensor * t_embd_pre_norm  = nullptr; // [n_embd, n_outputs] embeddings before the output norm

This new llama-ext API should be enough for everything to get the MTP working. There is no need to access the hidden state tensors from the user code, so there should be no changes to llama.h at all.

The rest of the changes from #22673 should be moved into a new common/speculative MTP implementation in the respective process()/draft() methods.

[am17an]

@ggerganov okay let me work on this, will send you my working branch once I have something.

[ggerganov]

Ok great, I think the only a bit complicated thing would be to add the logic for analyzing the contents of the batch in process() and draft() to determine which sequences participate and to map the embeddings offsets correctly. Basically, the "pending" logic that you already have, should be extended to work per sequence id.

An initial simple version could be to assert that there is single-sequence id inside the batch and when this is working, we can extend it to multi-sequence.

[am17an]

@ggerganov can you review am17an#6? It seems to work with these changes. cc @ngxson as well

[ruixiang63]

Thanks very much for the work! @ggerganov
I will also start an early prototype based on this refactoring, though it will not be finalized yet because unified feature extraction from the target model is still not finalized.

So I think after we fit the MTP implementation into the new speculative architecture, it will be easier to build the Eagle3 from there.

Makes sense to me!

Inside the process() method, we can analyze the contents of the batch and decide which embeddings to store for later, i.e. the embeddings for the last token of each sequence in the batch, and which to copy directly to the draft context for processing with the current batch.

Regarding “the embeddings for the last token of each sequence in the batch”: Eagle3 needs embeddings for every prompt token, rather than only the last token. Will this be supported in the current design?

With this API, utilize it to extract both the hidden target embeddings and the MTP-generated embeddings. There is no need to distinguish explicitly between the two types of embeddings.

Will this apply to Eagle3 as well? By the way, Eagle3 also uses the pre_norm embedding from the draft model as input for the next draft-token generation.

[ggerganov]

Regarding “the embeddings for the last token of each sequence in the batch”: Eagle3 needs embeddings for every prompt token, rather than only the last token. Will this be supported in the current design?

Yes. Take a look at how it is done in am17an#6.

Will this apply to Eagle3 as well? By the way, Eagle3 also uses the pre_norm embedding from the draft model as input for the next draft-token generation.

Yes, for the Eagle3 pre_norm embeddings do the same as we do in the MTP code.

[ggerganov]

Merging after CI is green.

[petersid2022]

thanks for including some of my changes to this PR! should i amend/reword the last commit to remove the out-of-date description that i originally left?

[ggerganov]

No need, it's not a problem if it is a bit outdated.

The base branch was changed.