What started as curiosity turned into a two-week rabbit hole of experiments and ended with a PoC that fully recovers the MTP PP overhead on GPU, above any expectation I had.

TL;DR: instead of processing the last layer MoE FFN for the entire ubatch tokens (usually 512-2048 tokens), this PoC processes only the output row (usually 1 token during prefill). The result is PP TPS is back to the same as with MTP disabled, keeping most of MTP's benefits to TG TPS, even with a slight drop in draft acceptance rate in one of the benchs.

I'm not opening a PR to llama.cpp because this is AI-generated code, which goes against their contribution policy, which I support. If you know C++ and llama.cpp internals, I invite to work together with me to open a PR with a more mature implementation.