Worth mentioning why this is harder than it looks.

There is a different set of experts at every layer, and each layer has a small router that decides which ones to use.

The router needs to look at the state produced by the experts below it.

Drafted tokens from the MTP head can be used to predict which experts the first layer will want, but not beyond that. To know what layer 10 experts needs, you have to run layers 1-9 which means loading their experts.

So, yes, instead of a next-token drafter like MTP, you'd want something trained to predict the expert activation across all layers at once.