transformers
e8e6b879 - deepgemm: only Mega MoE cannot apply a post-expert norm

Commit
44 days ago
deepgemm: only Mega MoE cannot apply a post-expert norm The guard refused a per-expert output norm on every DeepGEMM arm, but only Mega MoE has to: it fuses the routing-weighted reduce into its kernel, so there is no seam for the norm. The other two arms reduce in `_combine_routed_output`, so the norm rides the rows just before it — the same place the reference forwards apply it. `deepgemm_experts_guards` gains `post_expert_norm`, which only the Mega MoE arm sets False. A model naming a norm — Muse-Spark names `input_scaled_rms_norm` — was locked out of every DeepGEMM path and now loses only Mega MoE. The old behaviour had no test, which is why relaxing it broke nothing; the new one covers both halves so the Mega MoE restriction cannot silently disappear either. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Parents
Loading