A cross-model pass against the sibling PRs turned up that Qwen3-MoE was never actually reworked with the rest of the series. The branch still shipped the pre-fix sharding this whole effort exists to remove: the contracting hidden dim on the model axis, query heads and vocab on the batch axis, which is the ~2.7x-extra-attention-traffic bug and is also inconsistent with the layer's own correct attention_output rule. The test still asserted the old tuples too.
The fix flips the QKV kernels to (batch, model, None) and reverse_embeddings to (batch, model), leaves attention_output alone since it was already right, keeps GQA k/v replicated, and swaps the confused k/v comment for the Gemma-style rationale. The test is rebuilt on the shared helper (#2825) with mlp_only_layers=[1] so both routed-expert and dense-FFN rules get exercised, plus the mesh-shape and live-preset sweeps. Under #2807; PR #2818.
A cross-model pass against the sibling PRs turned up that Qwen3-MoE was never actually reworked with the rest of the series. The branch still shipped the pre-fix sharding this whole effort exists to remove: the contracting hidden dim on the model axis, query heads and vocab on the batch axis, which is the ~2.7x-extra-attention-traffic bug and is also inconsistent with the layer's own correct
attention_outputrule. The test still asserted the old tuples too.The fix flips the QKV kernels to
(batch, model, None)andreverse_embeddingsto(batch, model), leavesattention_outputalone since it was already right, keeps GQA k/v replicated, and swaps the confused k/v comment for the Gemma-style rationale. The test is rebuilt on the shared helper (#2825) withmlp_only_layers=[1]so both routed-expert and dense-FFN rules get exercised, plus the mesh-shape and live-preset sweeps. Under #2807; PR #2818.