Hmm ok I could be missing something here, I'm reusing the same set of input tokens for both N+1 and N+2 steps, while N+2 token need the sampled token from N+1

Not sure yet how we can do this

Would this implementation work on deepseev3?

@sorasoras Judging from this illustration, yes it's the same:

https://dataturbo.medium.com/deepseek-technical-analysis-3-multi-token-prediction-f8f3ea7eaf9c

Maybe this can be implemented by loading the MTP layers as a separate draft model and reusing the speculative decoding functionality. AFAICT, the predicted tokens from the MTP blocks are technically draft tokens and they have to be accepted by the main decoder.

Btw, based on these 2 diagrams, if they are correct, there is a small difference between DS and Mimo - Mimo uses the same h from the main decoder for all MTP blocks, while DS updates the h after each MTP block.

What is not clear to me is how big is N. In both diagrams, we have N = 4. Is this a parameter?

Ok thanks for the clue, it sounds like what you suggest is exactly what they did on vLLM implementation

I think the N is not important as the MTP layer has its own KV cache. The number of N token need to be corresponse to the number of embedding vectors from h. In other words, this is a way to implement residual connection, basically making the input token embedding to bypass the whole 36 "normal" layers

I'm thinking more about the implementation today, having 2 main ideas in my mind but both have downsides:

First idea is to have an API like llama_model_get_mtp(struct llama_model * model, int32_t i) which returns a shallow copy of the llama_model object, meaning only the pointer-to-tensors are copy, but not the actual data. The copied llama_model object will have a different layer index (or n_layer) to specify that it is a "child" model.

The downside is that:

• Because this is a shallow copy, llama_model_free will free tensors or both main and child models
• More importantly, this mean we need 2 different llama_context, and this mean we need yet another API to pass the hidden embeddings from one context to another

My second idea is to have something equivalent to llama_set_causal_attn, meaning the will be an attribute in llama_cparams to specify if we are about to run the llama_decode using main layers or MTP layers. However, managing KV cache in this case is a bit tricky and I still don't yet have any idea how to handle this.