transformers

606e6e8b - Skip the expert-parallel sentinel masking when expert parallelism is off (#48201)

Skip the expert-parallel sentinel masking when expert parallelism is off (#48201) * Skip the expert-parallel sentinel masking when expert parallelism is off `grouped_mm_experts_forward` builds a sentinel mask and runs two `masked_fill_` on tensors the size of the expert activations, on every forward and again on the backward. Sentinel ids only ever come from `RouterParallel`, i.e. only under expert parallelism; without it the mask is all-False and the work is pure memory traffic. Default the flag off where the experts module is built and set it where `MoeExpertsParallel` installs its forward. * Drop the explanatory comments; the reasoning is in the PR * Pass is_expert_parallel through to the experts TP style * Gate the batched_mm sentinel clamp and extend the EP-off skip to finegrained-fp8 and deepgemm * Apply ruff format to the deepgemm call sites * Carry is_expert_parallel on the test experts stand-ins The finegrained-fp8 and deepgemm forwards now read `self.is_expert_parallel`, which real experts modules get from `use_experts_implementation.__init__`. The kernel tests build their stand-ins as a `SimpleNamespace` and bypass the decorator, so all nine of them raised `AttributeError`. `_build_experts` carries the flag like the decorator does, and the two tests that exercise sentinels set it to `True`, since sentinels only exist under expert parallelism. * Set is_expert_parallel at the call site instead of threading a parameter * Set the flag where the experts are sharded, and keep the router path as is Wiring is_expert_parallel through the TP styles reached further than intended. It is read at tensor_parallel.py:601, where a true value skips the _AllReduceBackward on top_k_weights, so the router gradients would stop being summed across the group under expert parallelism. Nothing else sums them. That path now behaves exactly as on main. The flag also has to reach the experts module, and a module-level plan entry does not: llama4's expert-parallel plan has ep_router and grouped_gemm but no moe_tp_experts, so the experts would never have been marked. Set it in MoEParamShard.shard_param under shards_expert_dim, which is the grouped_gemm entry every expert-parallel plan carries, and which already rewrites the num_experts the sentinel mask compares against. The masking now tests the mask rather than the flag, so the two halves, forty lines apart in finegrained_fp8, cannot disagree. * Drop the dead default on _dispatch_routed_input Both call sites pass it. * Keep the expert-parallel flag off the public module surface The experts forwards read it at call time so it has to live on the module, but nothing outside them needs it, so it is `_is_expert_parallel` now.