Skip the expert-parallel sentinel masking when expert parallelism is off

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,

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,
 ) -> 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,13 @@ 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 +742,13 @@ 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
 

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)
+    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
@@ -501,7 +501,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 sentinel_mask is not None:
+        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)
@@ -551,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)
+    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
@@ -590,7 +591,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 sentinel_mask is not None:
+        weighted_out.masked_fill_(sentinel_mask, 0.0)
 
     # Restore original order
     inv_perm = torch.empty_like(perm)

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:
@@ -417,8 +418,10 @@ 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)
+    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,15 +437,17 @@ 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(
         selected_hidden_states_g, selected_weights, offsets, bias=selected_biases, is_transposed=self.is_transposed
     )  # (S, 2 * intermediate_dim) or  (S, intermediate_dim) depending on whether we have gating
 
     # Zero the sentinel-tail rows the kernel left uninitialized (fwd output and bwd `d_input`).
-    proj_out = proj_out.masked_fill(sentinel_mask, 0.0)
+    if sentinel_mask is not None:
+        proj_out = proj_out.masked_fill(sentinel_mask, 0.0)
 
     # Apply gating or activation
     if self.has_gate:
@@ -462,7 +467,8 @@ def grouped_mm_experts_forward(
     )  # (S, hidden_dim)
 
     # Same: zero the uninitialized sentinel-tail rows.
-    proj_out = proj_out.masked_fill(sentinel_mask, 0.0)
+    if sentinel_mask is not None:
+        proj_out = proj_out.masked_fill(sentinel_mask, 0.0)
 
     # Apply routing weights
     weighted_out = proj_out * sample_weights_g.unsqueeze(-1)  # (S, hidden_dim)
@@ -566,6 +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
 
         @wraps(original_forward)
         def forward(self, *args, **kwargs):

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)

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,