From 9284f04a5f478d8f2c66e22c7cfa3801d3cb755c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Fri, 21 Aug 2026 18:20:49 +0000
Subject: [PATCH 01/10] 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.
---
 .../distributed/tensor_parallel.py             |  2 ++
 src/transformers/integrations/moe.py           | 18 ++++++++++++++----
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 4534ec46887e..88ed04f7c4f7 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -593,6 +593,8 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_
 
     def install_forward(self, module, mesh, *, is_expert_parallel=False):
         """Install the transforms but pass `is_expert_parallel` in the forward call."""
+        # The experts forward needs this too: it is what decides whether router ids can carry sentinels.
+        module.is_expert_parallel = is_expert_parallel
         original_forward = module.forward
         output_source = (
             Partial()
diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index b19a9ede47bc..8f3cbe2ce3a9 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -417,8 +417,13 @@ def grouped_mm_experts_forward(
     # In-place clamp on `expert_ids_g` keeps the per-row bias gather in-bounds (bias added at
     # sentinel positions falls in rows the kernel skips, so harmless). Safe to mutate now —
     # nothing downstream needs the sentinel info from `expert_ids_g` itself.
-    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
-    expert_ids_g.clamp_(max=self.num_experts - 1)
+    # Without expert parallelism the router never emits an id >= num_experts, so the mask is all-False
+    # and the two `masked_fill_` below are pure memory traffic on tensors the size of the expert
+    # activations -- twice in the forward, twice more in the backward.
+    sentinel_mask = None
+    if self.is_expert_parallel:
+        sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
+        expert_ids_g.clamp_(max=self.num_experts - 1)
 
     # Select expert weights and biases
     # NOTE: We keep all experts here and rely on offsets to target the active ones.
@@ -434,7 +439,8 @@ def grouped_mm_experts_forward(
         selected_biases = self.up_proj_bias[expert_ids_g] if self.has_bias else None
 
     # Pre-mask (bwd path).
-    selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
+    if sentinel_mask is not None:
+        selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
 
     # --- Up projection per expert (grouped) ---
     proj_out = _grouped_linear(
@@ -462,7 +468,8 @@ def grouped_mm_experts_forward(
     weighted_out = proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
 
     # Post-mask (fwd path).
-    weighted_out.masked_fill_(sentinel_mask, 0.0)
+    if sentinel_mask is not None:
+        weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Restore original order
     inv_perm = torch.empty_like(perm)
@@ -563,6 +570,9 @@ def __init__(self, config, *args, **kwargs):
             self.has_bias = has_bias
             self.is_transposed = is_transposed
             self.is_concatenated = is_concatenated
+            # Only expert parallelism makes the router emit ids >= num_experts; `MoeExpertsParallel`
+            # flips this on when it installs its forward.
+            self.is_expert_parallel = False
 
         @wraps(original_forward)
         def forward(self, *args, **kwargs):

From 9ef37162512739d7e69d8f92ce9bc21fe17203b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Fri, 21 Aug 2026 21:42:50 +0000
Subject: [PATCH 02/10] Drop the explanatory comments; the reasoning is in the
 PR

---
 src/transformers/distributed/tensor_parallel.py | 1 -
 src/transformers/integrations/moe.py            | 5 -----
 2 files changed, 6 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 88ed04f7c4f7..6151e2d2cd43 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -593,7 +593,6 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_
 
     def install_forward(self, module, mesh, *, is_expert_parallel=False):
         """Install the transforms but pass `is_expert_parallel` in the forward call."""
-        # The experts forward needs this too: it is what decides whether router ids can carry sentinels.
         module.is_expert_parallel = is_expert_parallel
         original_forward = module.forward
         output_source = (
diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index 8f3cbe2ce3a9..362527a552bd 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -417,9 +417,6 @@ def grouped_mm_experts_forward(
     # In-place clamp on `expert_ids_g` keeps the per-row bias gather in-bounds (bias added at
     # sentinel positions falls in rows the kernel skips, so harmless). Safe to mutate now —
     # nothing downstream needs the sentinel info from `expert_ids_g` itself.
-    # Without expert parallelism the router never emits an id >= num_experts, so the mask is all-False
-    # and the two `masked_fill_` below are pure memory traffic on tensors the size of the expert
-    # activations -- twice in the forward, twice more in the backward.
     sentinel_mask = None
     if self.is_expert_parallel:
         sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
@@ -570,8 +567,6 @@ def __init__(self, config, *args, **kwargs):
             self.has_bias = has_bias
             self.is_transposed = is_transposed
             self.is_concatenated = is_concatenated
-            # Only expert parallelism makes the router emit ids >= num_experts; `MoeExpertsParallel`
-            # flips this on when it installs its forward.
             self.is_expert_parallel = False
 
         @wraps(original_forward)

From a7b07d3ff37ec84adfc43add3f5c54f3cc92e771 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Thu, 27 Aug 2026 04:12:49 +0000
Subject: [PATCH 03/10] Pass is_expert_parallel through to the experts TP style

---
 src/transformers/distributed/tensor_parallel.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index bf8c6b0100a2..a24a234b1aa8 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -144,7 +144,7 @@ def context_around_forward(self, module, mesh):
     def transform_output_post_forward(self, module, output, mesh):
         return output
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         """Install pre / around / post transforms by replacing module.forward."""
         original_forward = module.forward
 
@@ -344,7 +344,7 @@ class ReplicatedWithGradAllReduce(TensorParallelLayer):
     summed across the mesh.
     """
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         # A module hook rather than `param.register_hook`: params are replaced during weight
         # loading, which happens after TP is applied, and would drop a param-level hook.
         def _all_reduce_grads(mod, grad_input, grad_output):
@@ -409,7 +409,7 @@ def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = True):
         self.sequence_dim = sequence_dim
         self.use_local_output = use_local_output
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         # Replicate the module's params (LayerNorm/RMSNorm ones-init → from_local is safe).
         for p_name, p in list(module.named_parameters(recurse=False)):
             module.register_parameter(
@@ -793,7 +793,9 @@ def apply_tensor_parallelism(model, tp_mesh):
                 # MLA needs to know the qk_rope_head_dim to split the projection output into KV and RoPE parts.
                 # TODO: Store qk_rope_head_dim on MLA projection modules when the models initialize them.
                 module.config = model.config.get_text_config()
-            ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh)
+            ALL_PARALLEL_STYLES[style_name].install_forward(
+                module, tp_mesh, is_expert_parallel=model.config.distributed_config.enable_expert_parallel
+            )
         module._is_hooked = True
 
     return model

From a2a96130bc737cb4062e6150bea23a58b428db0a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Thu, 3 Sep 2026 01:05:59 +0000
Subject: [PATCH 04/10] Gate the batched_mm sentinel clamp and extend the
 EP-off skip to finegrained-fp8 and deepgemm

---
 src/transformers/integrations/deepgemm.py      | 18 ++++++++++++------
 .../integrations/finegrained_fp8.py            | 10 ++++++----
 src/transformers/integrations/moe.py           | 15 ++++++++-------
 3 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py
index ab266873f821..fce8fb7fd90f 100644
--- a/src/transformers/integrations/deepgemm.py
+++ b/src/transformers/integrations/deepgemm.py
@@ -504,6 +504,7 @@ def _dispatch_routed_input(
     num_experts: int,
     m_alignment: int,
     use_psum_layout: bool,
+    is_expert_parallel: bool = False,
 ) -> tuple:
     """Sort tokens by expert id and build the M-grouped padded layout.
 
@@ -533,8 +534,10 @@ def _dispatch_routed_input(
     # keeps any per-row gather (e.g. bias) in-bounds — bias added at sentinel positions falls
     # in rows the kernel skips, so harmless. Safe to mutate now: the layout was built from the
     # unclamped tensor and nothing downstream needs the sentinel info from `expert_ids_g` itself.
-    sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1)
-    expert_ids_g.clamp_(max=num_experts - 1)
+    sentinel_mask = None
+    if is_expert_parallel:
+        sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1)
+        expert_ids_g.clamp_(max=num_experts - 1)
     return (
         sorted_hidden_states_g,
         sample_weights_g,
@@ -550,7 +553,7 @@ def _dispatch_routed_input(
 def _combine_routed_output(
     out_padded: torch.Tensor,
     sorted_weights: torch.Tensor,
-    sentinel_mask: torch.Tensor,
+    sentinel_mask: torch.Tensor | None,
     perm: torch.Tensor,
     sorted_to_padded: torch.Tensor,
     num_tokens: int,
@@ -563,7 +566,8 @@ def _combine_routed_output(
     weighted = out * sorted_weights.to(out.dtype).unsqueeze(-1)
     # Sentinel rows past the valid expert blocks may carry NaN from allocator
     # reuse (`0 * NaN = NaN`); zero them so the top-k reduction stays finite.
-    weighted.masked_fill_(sentinel_mask, 0.0)
+    if sentinel_mask is not None:
+        weighted.masked_fill_(sentinel_mask, 0.0)
     inv_perm = torch.empty_like(perm)
     inv_perm[perm] = torch.arange(perm.size(0), device=out.device)
     # Deterministic reshape+sum (index_add_ with duplicates is non-deterministic on CUDA).
@@ -646,7 +650,8 @@ def deepgemm_bf16_experts_forward(
         grouped_layout,
         total_padded_rows,
     ) = _dispatch_routed_input(
-        hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100()
+        hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100(),
+        is_expert_parallel=self.is_expert_parallel,
     )
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
@@ -732,7 +737,8 @@ def deepgemm_fp8_fp4_experts_forward(
         grouped_layout,
         total_padded_rows,
     ) = _dispatch_routed_input(
-        hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100()
+        hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100(),
+        is_expert_parallel=self.is_expert_parallel,
     )
     sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None
 
diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py
index 342abd9c35af..f90a642be2a3 100644
--- a/src/transformers/integrations/finegrained_fp8.py
+++ b/src/transformers/integrations/finegrained_fp8.py
@@ -468,7 +468,7 @@ def fp8_batched_mm_experts_forward(
     # EP sentinel handling: leave `expert_ids` unclamped — the batched kernel early-returns on
     # `expert_id >= NUM_EXPERTS`, leaving sentinel output rows uninitialized. The post-mask below
     # zeroes them before the per-token reduction so `uninit * 0 = NaN` can't poison the sum.
-    sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1)
+    sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1) if self.is_expert_parallel else None
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
     weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv
@@ -506,7 +506,8 @@ def fp8_batched_mm_experts_forward(
 
     # Post-mask sentinel rows: kernel left them uninitialized, so zero them out
     # before the reduction below (uninit may be NaN; NaN * 0 = NaN).
-    weighted_out.masked_fill_(sentinel_mask, 0.0)
+    if self.is_expert_parallel:
+        weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Accumulate results using deterministic reshape+sum instead of index_add_
     # (index_add_ with duplicate indices is non-deterministic on CUDA due to atomicAdd)
@@ -556,7 +557,7 @@ def fp8_grouped_mm_experts_forward(
     # valid rows, so sentinel-tail `proj_out` rows are uninit; without the post-mask below,
     # `proj_out[sentinel] * 0 = NaN * 0 = NaN` would poison the per-token reduction. FP8
     # quantized weights are inference-only, so no bwd pre-mask is needed.
-    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
+    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) if self.is_expert_parallel else None
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
     weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv
@@ -595,7 +596,8 @@ def fp8_grouped_mm_experts_forward(
     weighted_out = proj_out * sample_weights_g.to(proj_out.dtype).unsqueeze(-1)  # (S, hidden_dim)
 
     # Post-mask (fwd path).
-    weighted_out.masked_fill_(sentinel_mask, 0.0)
+    if self.is_expert_parallel:
+        weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Restore original order
     inv_perm = torch.empty_like(perm)
diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index 3676f8415cb6..42fdeb39ca6b 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -125,11 +125,12 @@ def batched_mm_experts_forward(
     sample_weights = top_k_weights.reshape(-1)  # (S,)
     expert_ids = top_k_index.reshape(-1)  # (S,)
 
-    # Clamp EP sentinels so `gate_up_proj[expert_ids]` stays in-bounds. Routing weights are already
-    # zero at sentinel slots (RouterParallel masks them at dispatch), so the weighted mul drops
-    # those contributions — we pay the wasted GEMM compute because batched_mm has no offset to skip.
-    # Out-of-place to avoid mutating the caller's routing tensor (a contiguous `reshape(-1)` aliases it).
-    expert_ids = expert_ids.clamp(0, self.num_experts - 1)
+    if self.is_expert_parallel:
+        # Clamp EP sentinels so `gate_up_proj[expert_ids]` stays in-bounds. Routing weights are already
+        # zero at sentinel slots (RouterParallel masks them at dispatch), so the weighted mul drops
+        # those contributions — we pay the wasted GEMM compute because batched_mm has no offset to skip.
+        # Out-of-place to avoid mutating the caller's routing tensor (a contiguous `reshape(-1)` aliases it).
+        expert_ids = expert_ids.clamp(0, self.num_experts - 1)
 
     # Select gate_up or just up projection weights and biases
     if self.has_gate:
@@ -436,7 +437,7 @@ def grouped_mm_experts_forward(
         selected_biases = self.up_proj_bias[expert_ids_g] if self.has_bias else None
 
     # Pre-mask (bwd path).
-    if sentinel_mask is not None:
+    if self.is_expert_parallel:
         selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
 
     # --- Up projection per expert (grouped) ---
@@ -465,7 +466,7 @@ def grouped_mm_experts_forward(
     weighted_out = proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
 
     # Post-mask (fwd path).
-    if sentinel_mask is not None:
+    if self.is_expert_parallel:
         weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Restore original order

From 8f6dcf5e4b2b7a86a2c0d09bb2b777330474c474 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Thu, 3 Sep 2026 18:12:33 +0000
Subject: [PATCH 05/10] Apply ruff format to the deepgemm call sites

---
 src/transformers/integrations/deepgemm.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py
index fce8fb7fd90f..f7334f20b0f2 100644
--- a/src/transformers/integrations/deepgemm.py
+++ b/src/transformers/integrations/deepgemm.py
@@ -650,7 +650,12 @@ def deepgemm_bf16_experts_forward(
         grouped_layout,
         total_padded_rows,
     ) = _dispatch_routed_input(
-        hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100(),
+        hidden_states,
+        top_k_index,
+        top_k_weights,
+        self.num_experts,
+        deepgemm.m_alignment,
+        is_sm100(),
         is_expert_parallel=self.is_expert_parallel,
     )
 
@@ -737,7 +742,12 @@ def deepgemm_fp8_fp4_experts_forward(
         grouped_layout,
         total_padded_rows,
     ) = _dispatch_routed_input(
-        hidden_states, top_k_index, top_k_weights, self.num_experts, deepgemm.m_alignment, is_sm100(),
+        hidden_states,
+        top_k_index,
+        top_k_weights,
+        self.num_experts,
+        deepgemm.m_alignment,
+        is_sm100(),
         is_expert_parallel=self.is_expert_parallel,
     )
     sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None

From eb9dafc63fedd82d402f2a3640b80f4be232bcd8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Mon, 7 Sep 2026 21:47:14 +0000
Subject: [PATCH 06/10] 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.
---
 tests/kernels/test_finegrained_fp8.py |  4 ++--
 tests/kernels/test_utils.py           | 18 +++++++++++++++++-
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/tests/kernels/test_finegrained_fp8.py b/tests/kernels/test_finegrained_fp8.py
index c2e59a29731e..2e16b16bcead 100644
--- a/tests/kernels/test_finegrained_fp8.py
+++ b/tests/kernels/test_finegrained_fp8.py
@@ -264,7 +264,7 @@ def test_batched_mm_non_gated_uses_up_proj_and_activation(self):
     def test_batched_mm_passes_sentinel_expert_ids_unclamped(self):
         # EP sentinels (expert_ids >= num_experts) reach the kernel unclamped; the post-mask zeroes the
         # matching output rows before the per-token reduction (the kernel leaves them uninitialized).
-        experts = make_fp8_experts(num_experts=4, hidden=8, inter=16)
+        experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, is_expert_parallel=True)
         hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device)
         top_k_index = torch.tensor([[0, 4], [1, 4], [2, 4]], device=torch_device)  # 4 == num_experts -> sentinel
         top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device)
@@ -326,7 +326,7 @@ def test_grouped_mm_kernel_inputs_and_output(self):
     def test_grouped_mm_sentinels_dropped_from_histogram(self):
         # Sentinels are left unclamped so the sort pushes them to the tail and histc(max=num_experts-1)
         # drops them from tokens_per_expert -> no wasted GEMM rows; the post-mask zeroes their output.
-        experts = make_fp8_experts(num_experts=4, hidden=8, inter=16)
+        experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, is_expert_parallel=True)
         hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device)
         top_k_index = torch.tensor([[0, 4], [1, 4], [2, 4]], device=torch_device)  # three sentinels (== num_experts)
         top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device)
diff --git a/tests/kernels/test_utils.py b/tests/kernels/test_utils.py
index 1a3d110d9162..b751f00eb6dd 100644
--- a/tests/kernels/test_utils.py
+++ b/tests/kernels/test_utils.py
@@ -32,7 +32,18 @@
 
 
 def _build_experts(
-    *, num_experts, hidden, inter, has_gate, has_bias, is_transposed, weight_dtype, scale_dtype, hidden_act, **extra
+    *,
+    num_experts,
+    hidden,
+    inter,
+    has_gate,
+    has_bias,
+    is_transposed,
+    weight_dtype,
+    scale_dtype,
+    hidden_act,
+    is_expert_parallel,
+    **extra,
 ):
     def weight(out_dim, in_dim):
         # Non-transposed weights are (E, out, in); transposed are (E, in, out).
@@ -56,6 +67,7 @@ def apply_gate(gate_up):
         has_gate=has_gate,
         has_bias=has_bias,
         is_transposed=is_transposed,
+        is_expert_parallel=is_expert_parallel,
         act_fn=act_fn,
         _apply_gate=apply_gate,
         down_proj=weight(hidden, inter),
@@ -81,6 +93,7 @@ def make_experts(
     hidden_act="silu",
     is_concatenated=True,
     weight_dtype=torch.bfloat16,
+    is_expert_parallel=False,
 ):
     """BF16 experts stand-in (no scales) for the sonic-moe and DeepGEMM BF16 forwards. Carries
     `config.hidden_act` / `is_concatenated` (read by sonic-moe)."""
@@ -94,6 +107,7 @@ def make_experts(
         weight_dtype=weight_dtype,
         scale_dtype=None,
         hidden_act=hidden_act,
+        is_expert_parallel=is_expert_parallel,
         config=types.SimpleNamespace(hidden_act=hidden_act),
         is_concatenated=is_concatenated,
     )
@@ -111,6 +125,7 @@ def make_fp8_experts(
     scale_dtype=torch.float32,
     activation_scheme="dynamic",
     block_size=(128, 128),
+    is_expert_parallel=False,
 ):
     """FP8/FP4 experts stand-in (per-projection `_scale_inv`) for the DeepGEMM FP8 and finegrained-fp8
     forwards, plus the `_deepgemm_disabled` multi-device flag."""
@@ -124,6 +139,7 @@ def make_fp8_experts(
         weight_dtype=weight_dtype,
         scale_dtype=scale_dtype,
         hidden_act=hidden_act,
+        is_expert_parallel=is_expert_parallel,
         activation_scheme=activation_scheme,
         block_size=block_size,
         _deepgemm_disabled=False,

From a66940a0ea0a02c7ff354d2aae027dc7a559198c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Tue, 8 Sep 2026 16:38:25 +0000
Subject: [PATCH 07/10] Set is_expert_parallel at the call site instead of
 threading a parameter

---
 src/transformers/distributed/tensor_parallel.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 3a8b2ee6cab1..682beb0d0bf6 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -144,7 +144,7 @@ def context_around_forward(self, module, mesh):
     def transform_output_post_forward(self, module, output, mesh):
         return output
 
-    def install_forward(self, module, mesh, *, is_expert_parallel=False):
+    def install_forward(self, module, mesh):
         """Install pre / around / post transforms by replacing module.forward."""
         original_forward = module.forward
 
@@ -344,7 +344,7 @@ class ReplicatedWithGradAllReduce(TensorParallelLayer):
     summed across the mesh.
     """
 
-    def install_forward(self, module, mesh, *, is_expert_parallel=False):
+    def install_forward(self, module, mesh):
         # A module hook rather than `param.register_hook`: params are replaced during weight
         # loading, which happens after TP is applied, and would drop a param-level hook.
         def _all_reduce_grads(mod, grad_input, grad_output):
@@ -415,7 +415,7 @@ def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = True):
         self.sequence_dim = sequence_dim
         self.use_local_output = use_local_output
 
-    def install_forward(self, module, mesh, *, is_expert_parallel=False):
+    def install_forward(self, module, mesh):
         # Replicate the module's params (LayerNorm/RMSNorm ones-init → from_local is safe).
         for p_name, p in list(module.named_parameters(recurse=False)):
             module.register_parameter(
@@ -604,9 +604,8 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_
 
         return (hidden_states, *routing_args), kwargs
 
-    def install_forward(self, module, mesh, *, is_expert_parallel=False):
+    def install_forward(self, module, mesh):
         """Install the transforms but pass `is_expert_parallel` in the forward call."""
-        module.is_expert_parallel = is_expert_parallel
         original_forward = module.forward
         output_source = (
             Partial()
@@ -619,7 +618,7 @@ def install_forward(self, module, mesh, *, is_expert_parallel=False):
 
         def tp_forward(*args, **kwargs):
             args, kwargs = self.transform_inputs_pre_forward(
-                module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel
+                module, args, kwargs, mesh, is_expert_parallel=module.is_expert_parallel
             )
             with self.context_around_forward(module, mesh):
                 output = original_forward(*args, **kwargs)
@@ -817,9 +816,8 @@ def apply_tensor_parallelism(model, tp_mesh):
                 # MLA needs to know the qk_rope_head_dim to split the projection output into KV and RoPE parts.
                 # TODO: Store qk_rope_head_dim on MLA projection modules when the models initialize them.
                 module.config = model.config.get_text_config()
-            ALL_PARALLEL_STYLES[style_name].install_forward(
-                module, tp_mesh, is_expert_parallel=model.config.distributed_config.enable_expert_parallel
-            )
+            module.is_expert_parallel = model.config.distributed_config.enable_expert_parallel
+            ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh)
         module._is_hooked = True
 
     return model

From e9f3a20c74f62b6856625e5d86167f5bf99f6dbc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Tue, 8 Sep 2026 16:53:52 +0000
Subject: [PATCH 08/10] 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.
---
 src/transformers/distributed/tensor_parallel.py  | 7 ++++---
 src/transformers/integrations/finegrained_fp8.py | 4 ++--
 src/transformers/integrations/moe.py             | 4 ++--
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 682beb0d0bf6..405884c75885 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -541,6 +541,8 @@ def shard_param(self, module, param, mesh):
                     f"{expert_parallel_size} expert-parallel ranks."
                 )
             module.num_experts = global_num_experts // expert_parallel_size
+            # The experts forward masks sentinel rows only when its experts are actually split.
+            module.is_expert_parallel = True
         module._parameters[param] = torch.nn.Parameter(
             distribute_tensor(meta, mesh, [self.placement], src_data_rank=None),
             requires_grad=meta.requires_grad,
@@ -604,7 +606,7 @@ def transform_inputs_pre_forward(self, module, args, kwargs, mesh, *, is_expert_
 
         return (hidden_states, *routing_args), kwargs
 
-    def install_forward(self, module, mesh):
+    def install_forward(self, module, mesh, *, is_expert_parallel=False):
         """Install the transforms but pass `is_expert_parallel` in the forward call."""
         original_forward = module.forward
         output_source = (
@@ -618,7 +620,7 @@ def install_forward(self, module, mesh):
 
         def tp_forward(*args, **kwargs):
             args, kwargs = self.transform_inputs_pre_forward(
-                module, args, kwargs, mesh, is_expert_parallel=module.is_expert_parallel
+                module, args, kwargs, mesh, is_expert_parallel=is_expert_parallel
             )
             with self.context_around_forward(module, mesh):
                 output = original_forward(*args, **kwargs)
@@ -816,7 +818,6 @@ def apply_tensor_parallelism(model, tp_mesh):
                 # MLA needs to know the qk_rope_head_dim to split the projection output into KV and RoPE parts.
                 # TODO: Store qk_rope_head_dim on MLA projection modules when the models initialize them.
                 module.config = model.config.get_text_config()
-            module.is_expert_parallel = model.config.distributed_config.enable_expert_parallel
             ALL_PARALLEL_STYLES[style_name].install_forward(module, tp_mesh)
         module._is_hooked = True
 
diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py
index 9089241fdf75..f8f2775ed431 100644
--- a/src/transformers/integrations/finegrained_fp8.py
+++ b/src/transformers/integrations/finegrained_fp8.py
@@ -501,7 +501,7 @@ def fp8_batched_mm_experts_forward(
 
     # Post-mask sentinel rows: kernel left them uninitialized, so zero them out
     # before the reduction below (uninit may be NaN; NaN * 0 = NaN).
-    if self.is_expert_parallel:
+    if sentinel_mask is not None:
         weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Accumulate results using deterministic reshape+sum instead of index_add_
@@ -591,7 +591,7 @@ def fp8_grouped_mm_experts_forward(
     weighted_out = proj_out * sample_weights_g.to(proj_out.dtype).unsqueeze(-1)  # (S, hidden_dim)
 
     # Post-mask (fwd path).
-    if self.is_expert_parallel:
+    if sentinel_mask is not None:
         weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Restore original order
diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index 42fdeb39ca6b..386b6c0d1693 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -437,7 +437,7 @@ def grouped_mm_experts_forward(
         selected_biases = self.up_proj_bias[expert_ids_g] if self.has_bias else None
 
     # Pre-mask (bwd path).
-    if self.is_expert_parallel:
+    if sentinel_mask is not None:
         selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
 
     # --- Up projection per expert (grouped) ---
@@ -466,7 +466,7 @@ def grouped_mm_experts_forward(
     weighted_out = proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
 
     # Post-mask (fwd path).
-    if self.is_expert_parallel:
+    if sentinel_mask is not None:
         weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Restore original order

From b22c3025d746c53c20c560266dd6f7e57e72691f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Tue, 8 Sep 2026 16:57:16 +0000
Subject: [PATCH 09/10] Drop the dead default on _dispatch_routed_input

Both call sites pass it.
---
 src/transformers/integrations/deepgemm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py
index f7334f20b0f2..a23981d1bffe 100644
--- a/src/transformers/integrations/deepgemm.py
+++ b/src/transformers/integrations/deepgemm.py
@@ -504,7 +504,7 @@ def _dispatch_routed_input(
     num_experts: int,
     m_alignment: int,
     use_psum_layout: bool,
-    is_expert_parallel: bool = False,
+    is_expert_parallel: bool,
 ) -> tuple:
     """Sort tokens by expert id and build the M-grouped padded layout.
 

From ca7eeca34a926c149ba07da35610af4de2d9a915 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Thu, 10 Sep 2026 17:23:48 +0000
Subject: [PATCH 10/10] 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.
---
 src/transformers/distributed/tensor_parallel.py  | 2 +-
 src/transformers/integrations/deepgemm.py        | 4 ++--
 src/transformers/integrations/finegrained_fp8.py | 4 ++--
 src/transformers/integrations/moe.py             | 6 +++---
 tests/kernels/test_utils.py                      | 2 +-
 5 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py
index 13f93bc216b7..3151f6264a12 100644
--- a/src/transformers/distributed/tensor_parallel.py
+++ b/src/transformers/distributed/tensor_parallel.py
@@ -542,7 +542,7 @@ def shard_param(self, module, param, mesh):
                 )
             module.num_experts = global_num_experts // expert_parallel_size
             # The experts forward masks sentinel rows only when its experts are actually split.
-            module.is_expert_parallel = True
+            module._is_expert_parallel = True
         module._parameters[param] = torch.nn.Parameter(
             distribute_tensor(meta, mesh, [self.placement], src_data_rank=None),
             requires_grad=meta.requires_grad,
diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py
index a23981d1bffe..880526719c7e 100644
--- a/src/transformers/integrations/deepgemm.py
+++ b/src/transformers/integrations/deepgemm.py
@@ -656,7 +656,7 @@ def deepgemm_bf16_experts_forward(
         self.num_experts,
         deepgemm.m_alignment,
         is_sm100(),
-        is_expert_parallel=self.is_expert_parallel,
+        is_expert_parallel=self._is_expert_parallel,
     )
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
@@ -748,7 +748,7 @@ def deepgemm_fp8_fp4_experts_forward(
         self.num_experts,
         deepgemm.m_alignment,
         is_sm100(),
-        is_expert_parallel=self.is_expert_parallel,
+        is_expert_parallel=self._is_expert_parallel,
     )
     sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None
 
diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py
index f8f2775ed431..084b17733b16 100644
--- a/src/transformers/integrations/finegrained_fp8.py
+++ b/src/transformers/integrations/finegrained_fp8.py
@@ -463,7 +463,7 @@ def fp8_batched_mm_experts_forward(
     # EP sentinel handling: leave `expert_ids` unclamped — the batched kernel early-returns on
     # `expert_id >= NUM_EXPERTS`, leaving sentinel output rows uninitialized. The post-mask below
     # zeroes them before the per-token reduction so `uninit * 0 = NaN` can't poison the sum.
-    sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1) if self.is_expert_parallel else None
+    sentinel_mask = (expert_ids >= self.num_experts).unsqueeze(-1) if self._is_expert_parallel else None
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
     weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv
@@ -552,7 +552,7 @@ def fp8_grouped_mm_experts_forward(
     # valid rows, so sentinel-tail `proj_out` rows are uninit; without the post-mask below,
     # `proj_out[sentinel] * 0 = NaN * 0 = NaN` would poison the per-token reduction. FP8
     # quantized weights are inference-only, so no bwd pre-mask is needed.
-    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) if self.is_expert_parallel else None
+    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1) if self._is_expert_parallel else None
 
     weight_up = self.gate_up_proj if self.has_gate else self.up_proj
     weight_scale_up = self.gate_up_proj_scale_inv if self.has_gate else self.up_proj_scale_inv
diff --git a/src/transformers/integrations/moe.py b/src/transformers/integrations/moe.py
index e19ad6d98e26..ba61ef85bf7f 100644
--- a/src/transformers/integrations/moe.py
+++ b/src/transformers/integrations/moe.py
@@ -125,7 +125,7 @@ def batched_mm_experts_forward(
     sample_weights = top_k_weights.reshape(-1)  # (S,)
     expert_ids = top_k_index.reshape(-1)  # (S,)
 
-    if self.is_expert_parallel:
+    if self._is_expert_parallel:
         # Clamp EP sentinels so `gate_up_proj[expert_ids]` stays in-bounds. Routing weights are already
         # zero at sentinel slots (RouterParallel masks them at dispatch), so the weighted mul drops
         # those contributions — we pay the wasted GEMM compute because batched_mm has no offset to skip.
@@ -419,7 +419,7 @@ def grouped_mm_experts_forward(
     # sentinel positions falls in rows the kernel skips, so harmless). Safe to mutate now —
     # nothing downstream needs the sentinel info from `expert_ids_g` itself.
     sentinel_mask = None
-    if self.is_expert_parallel:
+    if self._is_expert_parallel:
         sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
         expert_ids_g.clamp_(max=self.num_experts - 1)
 
@@ -572,7 +572,7 @@ def __init__(self, config, *args, **kwargs):
             self.has_bias = has_bias
             self.is_transposed = is_transposed
             self.is_concatenated = is_concatenated
-            self.is_expert_parallel = False
+            self._is_expert_parallel = False
 
         @wraps(original_forward)
         def forward(self, *args, **kwargs):
diff --git a/tests/kernels/test_utils.py b/tests/kernels/test_utils.py
index b751f00eb6dd..dc467c451389 100644
--- a/tests/kernels/test_utils.py
+++ b/tests/kernels/test_utils.py
@@ -67,7 +67,7 @@ def apply_gate(gate_up):
         has_gate=has_gate,
         has_bias=has_bias,
         is_transposed=is_transposed,
-        is_expert_parallel=is_expert_parallel,
+        _is_expert_parallel=is_expert_parallel,
         act_fn=act_fn,
         _apply_gate=apply_gate,
         down_proj=weight(hidden, inter),