DeepGEMM BF16 + mixed FP8/FP4 + MegaMoE + refactor

src/transformers/conversion_mapping.py

@@ -329,6 +329,10 @@ def _build_checkpoint_conversion_mapping():
                 source_patterns=r"^layers\.(\d+)\.self_attn\.indexer\.",
                 target_patterns=r"layers.\1.self_attn.compressor.indexer.",
             ),
+            WeightRenaming(
+                source_patterns=r"^layers\.(\d+)\.self_attn\.compressor\.indexer\.weights_proj\.",
+                target_patterns=r"layers.\1.self_attn.compressor.indexer.scorer.weights_proj.",
+            ),
             WeightRenaming(
                 source_patterns=r"^layers\.(\d+)\.self_attn\.compressor\.norm\.",
                 target_patterns=r"layers.\1.self_attn.compressor.kv_norm.",

src/transformers/integrations/hub_kernels.py

@@ -288,7 +288,7 @@ def register_kernel_mapping_transformers(*args, **kwargs):
     "mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
     "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "version": 1},
     "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 1},
-    "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 1},
+    "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2},
     "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"},
 }
 
@@ -376,7 +376,7 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _
             repo_id = _HUB_KERNEL_MAPPING[kernel_name]["repo_id"]
             revision = _HUB_KERNEL_MAPPING[kernel_name].get("revision", None)
             version = _HUB_KERNEL_MAPPING[kernel_name].get("version", None)
-            kernel = get_kernel(repo_id, revision=revision, version=version)
+            kernel = get_kernel(repo_id, revision=revision, version=version, allow_all_kernels=ALLOW_ALL_KERNELS)
             mapping[kernel_name] = kernel
         except FileNotFoundError as e:
             mapping[kernel_name] = None

src/transformers/integrations/moe.py

@@ -24,6 +24,7 @@
     is_torch_less_or_equal,
     is_torchdynamo_compiling,
 )
+from .deepgemm import deepgemm_bf16_experts_forward
 from .sonicmoe import sonicmoe_experts_forward
 
 
@@ -133,7 +134,8 @@ def batched_mm_experts_forward(
     # 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.
-    expert_ids.clamp_(0, self.num_experts - 1)
+    # 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:
@@ -483,6 +485,7 @@ class ExpertsInterface(GeneralInterface):
     """Interface for registering custom experts forward functions."""
 
     _global_mapping = {
+        "deepgemm": deepgemm_bf16_experts_forward,
         "batched_mm": batched_mm_experts_forward,
         "grouped_mm": grouped_mm_experts_forward,
         "sonicmoe": sonicmoe_experts_forward,

src/transformers/integrations/sonicmoe.py

@@ -28,6 +28,7 @@
 
 from ..utils import logging
 from .hub_kernels import lazy_load_kernel
+from .tensor_parallel import to_local
 
 
 logger = logging.get_logger(__name__)
@@ -169,18 +170,10 @@ def sonicmoe_experts_forward(
     # already zero (RouterParallel masks them at dispatch), so the per-token reduction
     # contributes nothing for sentinel slots.
 
-    # FSDP2 / EP wraps weights as DTensors but the kernel takes raw CUTLASS / CuteDSL pointers,
-    # so unwrap to local shards before reshaping. `to_local()` is autograd-aware — backward
-    # will rewrap the gradient as a DTensor matching each parameter's placements.
-    w1 = self.gate_up_proj
-    w2 = self.down_proj
-    b1 = self.gate_up_proj_bias if self.has_bias else None
-    b2 = self.down_proj_bias if self.has_bias else None
-    if isinstance(w1, torch.distributed.tensor.DTensor):
-        w1 = w1.to_local()
-        w2 = w2.to_local()
-        b1 = b1.to_local() if b1 is not None else None
-        b2 = b2.to_local() if b2 is not None else None
+    w1 = to_local(self.gate_up_proj)
+    w2 = to_local(self.down_proj)
+    b1 = to_local(self.gate_up_proj_bias) if self.has_bias else None
+    b2 = to_local(self.down_proj_bias) if self.has_bias else None
 
     # Map activation function
     act_name = getattr(self.config, "hidden_act", "silu").lower()

src/transformers/integrations/tensor_parallel.py

@@ -37,6 +37,21 @@
 logger = logging.get_logger(__name__)
 
 
+def to_local(t):
+    """Unwrap a `DTensor` to its local shard if needed; pass through otherwise.
+
+    Custom kernels (CUTLASS, CuteDSL, Triton) take raw tensor pointers and don't
+    understand `DTensor`, so weights wrapped by FSDP2 / EP need this unwrap before
+    they can be fed to the kernel. ``to_local()`` is autograd-aware on the train
+    path: backward rewraps the gradient as a DTensor matching each parameter's
+    placements.
+    """
+    if hasattr(torch.distributed, "tensor") and hasattr(torch.distributed.tensor, "DTensor"):
+        if isinstance(t, torch.distributed.tensor.DTensor):
+            return t.to_local()
+    return t
+
+
 def initialize_tensor_parallelism(
     tp_plan: str | dict[str, str] | None, tp_size: int | None = None, device_mesh=None, device_map=None
 ):
@@ -773,6 +788,25 @@ def _backward_hook(mod, grad_input, grad_output, mesh=device_mesh):
         module.register_full_backward_hook(_backward_hook)
 
 
+class AllReduceParallel(TensorParallelLayer):
+    """All-reduce a module's forward output across the TP mesh. Use as a declarative
+    sync point at the boundary of a multi-arg module whose compute ends in a partial
+    sum (e.g. the lightning indexer's score sum before its top-k).
+    """
+
+    def _prepare_input_fn(self, mod, inputs, device_mesh):
+        return inputs
+
+    def _prepare_output_fn(self, mod, outputs, device_mesh):
+        return all_reduce_forward(outputs, device_mesh)
+
+    def shard_tensor(self, param, tensor_idx=None, device=None, dtype=None):
+        return param[...].to(device=device, dtype=dtype)
+
+    def prepare_module_tp(self, module, device_mesh, **kwargs):
+        distribute_module(module, device_mesh, output_fn=self._prepare_output_fn)
+
+
 class MlaKvAProjParallel(TensorParallelLayer):
     """
     For MLA attention used in DeepSeek-V2 style models (deepseek_v2, longcat_flash, glm_moe_dsa, glm4_moe_lite):
@@ -1095,7 +1129,7 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     def _prepare_input_fn(self, mod, inputs, device_mesh):
-        return inputs[0] if inputs else inputs
+        return inputs
 
     def _prepare_output_fn(self, mod, outputs, device_mesh):
         """
@@ -1173,6 +1207,19 @@ def shard_tensor(
         return param[...].to(device=device, dtype=dtype)
 
 
+class RouterParallelMegaMoe(RouterParallel):
+    """Router TP plan used with DeepGEMM Mega MoE.
+
+    Mega MoE handles EP dispatch inside the kernel and wants raw global expert ids
+    with unmasked routing weights, so the router doesn't pre-shard per EP rank like
+    `RouterParallel._prepare_output_fn` does. The quantizer's `update_tp_plan` swaps
+    `"ep_router"` → `"megamoe_router"` when `experts_implementation == "deepgemm_megamoe"`.
+    """
+
+    def _prepare_output_fn(self, mod, outputs, device_mesh):
+        return outputs
+
+
 class MoeTensorParalellExperts(TensorParallelLayer):
     """
     Note: For tensor parallel, the MoEExpertsParallel TP layer handles gradient sync:
@@ -1198,7 +1245,7 @@ def _prepare_input_fn(self, mod, inputs, device_mesh):
         # and partial_expert_output is different on each GPU before all-reduce
         top_k_weights = all_reduce_backward(top_k_weights, device_mesh)
 
-        return (hidden_states, top_k_index, top_k_weights)
+        return hidden_states, top_k_index, top_k_weights
 
     def _prepare_output_fn(self, mod, outputs, device_mesh):
         # all_reduce_forward to sum partial expert outputs across GPUs
@@ -1212,6 +1259,32 @@ def shard_tensor(
         return param[...].to(device=device, dtype=dtype)
 
 
+class MoeTensorParalellMegaMoeExperts(TensorParallelLayer):
+    """TP layer for DeepGEMM Mega MoE experts.
+
+    Mega MoE is inference-only (the kernel has no backward) and handles EP dispatch +
+    combine + per-rank token sharding internally — so we skip the gradient-sync hooks
+    that the regular `MoeTensorParalellExperts` would apply, and we forward the EP
+    `process_group` into the module so the symm-buffer rendezvous can run on first
+    forward. The quantizer's `update_tp_plan` swaps the experts plan key from
+    `"moe_tp_experts"` to `"megamoe_experts"` when
+    `from_pretrained(..., experts_implementation="deepgemm_megamoe")`.
+    """
+
+    def _prepare_input_fn(self, mod, inputs, device_mesh):
+        hidden_states, top_k_index, top_k_weights = inputs[0], inputs[1], inputs[2]
+        return hidden_states, top_k_index, top_k_weights, device_mesh.get_group()
+
+    def _prepare_output_fn(self, mod, outputs, device_mesh):
+        # Kernel returned the fully-combined gathered output; no further reduction.
+        return outputs
+
+    def shard_tensor(
+        self, param: torch.Tensor, tensor_idx: int | None = None, device=None, dtype=None
+    ) -> torch.Tensor:
+        return param[...].to(device=device, dtype=dtype)
+
+
 class MoeIdentityExpertParallel(TensorParallelLayer):
     """
     TP class for zero/identity experts in MoE layers.
@@ -1250,10 +1323,13 @@ class ParallelInterface(GeneralInterface):
             "sequence_parallel": SequenceParallel(),
             "grouped_gemm": GroupedGemmParallel(),
             "ep_router": RouterParallel(),
+            "megamoe_router": RouterParallelMegaMoe(),
             "moe_tp_experts": MoeTensorParalellExperts(),
+            "megamoe_experts": MoeTensorParalellMegaMoeExperts(),
             "moe_identity_expert": MoeIdentityExpertParallel(),
             "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(),
             "mla_kv_a_proj": MlaKvAProjParallel(),
+            "all_reduce": AllReduceParallel(),
         }
         if is_torch_available() and _torch_distributed_available
         else {}
@@ -1274,6 +1350,7 @@ class ParallelInterface(GeneralInterface):
         "sequence_parallel": None,
         "replicated_with_grad_allreduce": None,
         "mla_kv_a_proj": None,
+        "all_reduce": None,
     }
 
     # Bias sharding: colwise shards bias, rowwise doesn't (bias is replicated and all-reduced)
@@ -1289,6 +1366,7 @@ class ParallelInterface(GeneralInterface):
         "sequence_parallel": None,
         "replicated_with_grad_allreduce": None,
         "mla_kv_a_proj": None,
+        "all_reduce": None,
     }
 
     @classmethod

src/transformers/modeling_utils.py

@@ -2212,6 +2212,18 @@ def set_experts_implementation(self, experts_implementation: str | dict):
             else experts_implementation.get("", self.config._experts_implementation)
         )
 
+        # MegaMoE is locked at load time: its TP plan is baked into `base_model_tp_plan`
+        # by `update_tp_plan` (and isn't re-evaluated) and `setup_megamoe_weights`
+        # mutates the expert weights into UTCCP layout on first forward. Either side of
+        # a switch would silently produce garbage, so reject it with a clear pointer.
+        current = self.config._experts_implementation
+        if "deepgemm_megamoe" in (current, requested_implementation) and current != requested_implementation:
+            raise RuntimeError(
+                f"Cannot switch experts implementation from {current!r} to {requested_implementation!r} "
+                "at runtime: `deepgemm_megamoe` is a load-time choice. Reload via "
+                "`from_pretrained(..., experts_implementation=...)` to switch."
+            )
+
         if requested_implementation != self.config._experts_implementation:
             requested_implementation = self._check_and_adjust_experts_implementation(requested_implementation)
             # Apply the change (on the internal attr, to avoid setting it recursively)

src/transformers/models/deepseek_v4/configuration_deepseek_v4.py

@@ -112,20 +112,28 @@ class DeepseekV4Config(PreTrainedConfig):
         "norm": (["hidden_states"], ["hidden_states"]),
     }
     base_model_ep_plan = {
-        # EP-only by default, same shape as gpt-oss: route on the gate, run the
-        # routed experts as a grouped-GEMM kernel sharded along the expert axis,
-        # and wrap the experts module with `moe_tp_experts` so its output gets
-        # all-reduced across ranks. Attention stays replicated (V4 is shared-KV
-        # MQA + a CSA / HCA compressor branch — both broadcast a single KV head
-        # across all attention heads via `repeat_kv`, so colwise-sharding
-        # `q_b_proj` would leave KV replicated and `repeat_kv` would no longer
-        # match the rank-local query head count). The shared MLP also stays
-        # replicated — it's small and not worth TP-ing. There's deliberately
-        # no `base_model_tp_plan` for V4: we don't ship a pure-TP plan, only EP.
+        # V4 ships EP only (no `base_model_tp_plan` — the runtime picks one plan or
+        # the other, never both, and V4 is MoE so EP is the only sensible config).
+        # MoE parallelism: route on the gate, run the routed experts as a grouped-GEMM
+        # kernel sharded along the expert axis, and wrap the experts module with
+        # `moe_tp_experts` so its output gets all-reduced across ranks. Same shape as
+        # gpt-oss. Main attention stays replicated: V4 is shared-KV MQA + a CSA / HCA
+        # compressor branch — both broadcast a single KV head across all attention
+        # heads via `repeat_kv`, so colwise-sharding `q_b_proj` would leave KV
+        # replicated and `repeat_kv` would no longer match the rank-local query head
+        # count. The shared MLP also stays replicated — it's small and not worth
+        # sharding. The Lightning Indexer is the one carve-out: its keys are
+        # replicated (own compressor at index_head_dim fed by replicated
+        # hidden_states), so head-sharding is well-formed; `q_b_proj` and the
+        # `scorer.weights_proj` go colwise, and the `scorer` output is all-reduced
+        # so every rank sees the same `index_scores` and picks the same top-k.
         "layers.*.mlp.gate": "ep_router",
         "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
         "layers.*.mlp.experts.down_proj": "grouped_gemm",
         "layers.*.mlp.experts": "moe_tp_experts",
+        "layers.*.self_attn.compressor.indexer.q_b_proj": "colwise",
+        "layers.*.self_attn.compressor.indexer.scorer.weights_proj": "colwise",
+        "layers.*.self_attn.compressor.indexer.scorer": "all_reduce",
     }
 
     vocab_size: int = 129280
@@ -283,12 +291,14 @@ def __post_init__(self, **kwargs):
             )
         self.qk_rope_head_dim = int(self.head_dim * self.partial_rotary_factor)
 
-        # yarn is applied ONLY to layers with a
-        # compressor (CSA/HCA); pure sliding-window layers use plain RoPE with
-        # `theta=rope_theta` (10000) and no scaling. Compress layers use
-        # `theta=compress_rope_theta` (160000) with yarn factor=16, and the reference
-        # does NOT multiply cos/sin by the yarn mscale — force `attention_factor=1.0`
-        # so transformers' `_compute_yarn_parameters` doesn't apply `0.1·log(16)+1`.
+        # yarn is applied ONLY to layers with a compressor (CSA/HCA); pure
+        # sliding-window layers use plain RoPE with `theta=rope_theta` (10000) and no
+        # scaling. Compress layers use `theta=compress_rope_theta` (160000) with yarn
+        # factor=16. For a YaRN compress layer we force `attention_factor=1.0`: the V4
+        # reference does NOT multiply cos/sin by YaRN's computed mscale, and leaving the
+        # key unset lets `_compute_yarn_parameters` derive `0.1·log(factor)+1 ≈ 1.277`.
+        # We only inject the key when `rope_type="yarn"` — the default-rope validator
+        # rejects it (and `compute_default_rope_parameters` hardcodes 1.0 anyway).
         rp = self.rope_parameters or {}
         if isinstance(rp.get("main"), dict) and isinstance(rp.get("compress"), dict):
             # Already nested — drop any leftover top-level keys.