Expert choice actually happens per layer, not just per token. It's not a limitation when doing inference at scale since all experts are then loaded in fast VRAM

Expert choice actually happens per layer, not just per token. It's not a limitation when doing inference at scale since all experts are then loaded in fast VRAM anyway. It's mostly just a way to enforce some kind of model sparsity and save on compute. The new version of Apple Foundation Model (AFM) Core Advanced is an exception, it actually routes experts per prompt (with sporadic rerouting perhaps?) which is more in line with what you're proposing. But this will theoretically make the model less smart than a similar one where experts are picked per layer.