feat: bump build for shared experts

build/torch26-cxx11-cu118-x86_64-linux/megablocks/{_megablocks_76c7de7.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3c5605ba50f2661b9dc4c5609572323fb4f52787181109c5900c261c5e2bf602
+oid sha256:070067fec0e735e865610caf4fc33b384fe8c9c47a002c365f740c82c5af1bab
 size 10517576

build/torch26-cxx11-cu118-x86_64-linux/megablocks/_megablocks_9a1816c.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4e4c48e189572141f6a140dd83f9eca19eaebbc20c5cd686aa0263aafec14533
-size 10517576

build/torch26-cxx11-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch26-cxx11-cu118-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out

build/torch26-cxx11-cu124-x86_64-linux/megablocks/{_megablocks_76c7de7.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:12c21d6f72f90950adbda156534691dd753476a18719b416541e8d6920a173b4
+oid sha256:02dffd561ef226c1ec17c99e462c3c771879f078dde9b1e5cd8bd5992be5b3da
 size 11869392

build/torch26-cxx11-cu124-x86_64-linux/megablocks/_megablocks_9a1816c.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3d958a0c77589a5ede72336d1cab80ea9d6324ef6f8a9a187af2da4db74e1894
-size 11869392

build/torch26-cxx11-cu124-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out

build/torch26-cxx11-cu126-x86_64-linux/megablocks/{_megablocks_76c7de7.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:db7e3b7c3c15af78fe9ef0ba50c33cb2cb988bdf5dfb1f46807b7871e7c8e70e
+oid sha256:b5aa4e066ddbd863693ca8a5ec37fba34996226442dfa407e4a49b779497001d
 size 11931048

build/torch26-cxx11-cu126-x86_64-linux/megablocks/_megablocks_9a1816c.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d41a4f5bbc160f51b058d3ba36e9087e9f15d35ae4782f36c984dd7199ee8ede
-size 11931048

build/torch26-cxx11-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out

build/torch26-cxx98-cu118-x86_64-linux/megablocks/{_megablocks_76c7de7.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d9971a30d397598ee0a58118b8cca337d142de1ca34404532dfda6328122ab11
+oid sha256:fababa7e0d2c20c98afaebef6165a8145b33d80cdadba28f895c14dd2a7b2823
 size 10510040

build/torch26-cxx98-cu118-x86_64-linux/megablocks/_megablocks_9a1816c.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:01f0c774e900380d3c0721dfe15591c67be5d5eb5ad687af6c89a88ecdff4f2a
-size 10510040

build/torch26-cxx98-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out

build/torch26-cxx98-cu124-x86_64-linux/megablocks/{_megablocks_76c7de7.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c136e90b35e7fd43fcc4d987588f68b3f4cfea295a00f1fda343acc9c8848577
+oid sha256:9e3663f46030f07e030efe94c26495d17b2703551a46c0ca3acf8b25ecb2a238
 size 11857920

build/torch26-cxx98-cu124-x86_64-linux/megablocks/_megablocks_9a1816c.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:09a5f57ae37af9f5b14c4a0f21d1679e32f5b7424973c36dac9bbbecbfbf7374
-size 11857920

build/torch26-cxx98-cu124-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_megablocks_76c7de7.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a3f893773ec7b8157a4531a57821807f5f27ac48ceaa695c342cc7a39ad318dc
-size 11927768

build/torch26-cxx98-cu126-x86_64-linux/megablocks/{_megablocks_9a1816c.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d155b22a3a413d23e1d1b6f65fd3700b2e004e45daf1cca1b397b8e0b4d68616
+oid sha256:d1571732c5954914d5ddf0f12ebc4074d88d907130d71d898de43958e3b9a5d1
 size 11923672

build/torch26-cxx98-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out

build/torch27-cxx11-cu118-x86_64-linux/megablocks/{_megablocks_76c7de7.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:aff7108245384777d22e9023ae3fd4cf2bcb0015a0938e314d556dbd3e59fe00
+oid sha256:a39b315c5359b79a67282160b5b344853aa06b5a5c9d8efafb903eb4f249b645
 size 10517816

build/torch27-cxx11-cu118-x86_64-linux/megablocks/_megablocks_9a1816c.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:002c2687dbc5693308fe32eaebe2f45ed3c85454fd45bc06d7b30e9c1a6d8949
-size 10517816

build/torch27-cxx11-cu118-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch27-cxx11-cu118-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out

build/torch27-cxx11-cu126-x86_64-linux/megablocks/{_megablocks_76c7de7.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:aa1eeccba0a3a26435538a2aa87bc22a40c0201a79979872f6296af984e7bf1e
+oid sha256:4870e4a9a831c30c7177b9b23b2b20d64f47242f16d818be1884b4e130e063c1
 size 11931080

build/torch27-cxx11-cu126-x86_64-linux/megablocks/_megablocks_9a1816c.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ef9197ea269734d4e0528887ab3c353fa8ba10ccf9a82c9abe85b72bc0ea3553
-size 11931080

build/torch27-cxx11-cu126-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch27-cxx11-cu126-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out

build/torch27-cxx11-cu128-x86_64-linux/megablocks/{_megablocks_76c7de7.abi3.so → _megablocks_89e2950.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:64e2fd33ed4a5e9497ad304763c3c174ade26702a8e43fe8e7b3d3e79eb1e021
+oid sha256:37844f7b2972aae75a1eeb8cda3b573a93ef27dd5a73b2cfb95fca1f41da07d9
 size 17892624

build/torch27-cxx11-cu128-x86_64-linux/megablocks/_megablocks_9a1816c.abi3.so

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b071dec56af72c9e6b8408106b97fb42355b08e94cc1200bb6f4d3f42ba0e97e
-size 17892624

build/torch27-cxx11-cu128-x86_64-linux/megablocks/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _megablocks_76c7de7
-ops = torch.ops._megablocks_76c7de7
+from . import _megablocks_89e2950
+ops = torch.ops._megablocks_89e2950
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_megablocks_76c7de7::{op_name}"
+    return f"_megablocks_89e2950::{op_name}"

build/torch27-cxx11-cu128-x86_64-linux/megablocks/layers.py

@@ -152,6 +152,66 @@ def mlp_forward(
     return torch.bmm(x, w2) + w2_bias[..., None, :]
 
 
+# Shared expert MLP forward pass
+def shared_mlp_forward(
+    x: torch.Tensor,
+    up_proj_weight: torch.Tensor,
+    down_proj_weight: torch.Tensor,
+    up_proj_bias: Optional[torch.Tensor] = None,
+    down_proj_bias: Optional[torch.Tensor] = None,
+    activation_fn: Optional[Any] = None,
+    gradient_scale: Optional[float] = None,
+) -> torch.Tensor:
+    # Default activation function
+    if activation_fn is None:
+        activation_fn = torch.nn.functional.gelu
+
+    # Scale weights
+    up_proj_weight = scale_grad(up_proj_weight, gradient_scale)
+    down_proj_weight = scale_grad(down_proj_weight, gradient_scale)
+    if up_proj_bias is not None:
+        up_proj_bias = scale_grad(up_proj_bias, gradient_scale)
+    if down_proj_bias is not None:
+        down_proj_bias = scale_grad(down_proj_bias, gradient_scale)
+
+    # Resolve dtensors
+    up_proj_weight = _layers.mlp.resolve_dtensor(up_proj_weight)
+    down_proj_weight = _layers.mlp.resolve_dtensor(down_proj_weight)
+    if up_proj_bias is not None:
+        up_proj_bias = _layers.mlp.resolve_dtensor(up_proj_bias)
+    if down_proj_bias is not None:
+        down_proj_bias = _layers.mlp.resolve_dtensor(down_proj_bias)
+
+    # Up projection
+    x = torch.nn.functional.linear(x, up_proj_weight, up_proj_bias)
+    
+    # Activation
+    x = activation_fn(x)
+    
+    # Down projection
+    x = torch.nn.functional.linear(x, down_proj_weight, down_proj_bias)
+    
+    return x
+
+
+# Combine outputs from shared expert and regular experts
+def combine_expert_shared_outputs(
+    shared_expert_out: torch.Tensor,
+    expert_out: torch.Tensor,
+    shared_expert_weighted_sum: bool = False,
+    moe_top_k: int = 1,
+) -> torch.Tensor:
+    if shared_expert_weighted_sum:
+        # Weighted sum based on number of experts used
+        total_experts = moe_top_k + 1
+        shared_weight = 1.0 / total_experts
+        expert_weight = moe_top_k / total_experts
+        return shared_expert_out * shared_weight + expert_out * expert_weight
+    else:
+        # Simple addition
+        return shared_expert_out + expert_out
+
+
 # Global variable to store load balancing loss
 _LOAD_BALANCING_LOSS = []
 
@@ -680,6 +740,136 @@ def moe_forward(
     return x, expert_weights, router_scores
 
 
+def moe_forward_with_shared_expert(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    moe_top_k: int,
+    moe_num_experts: int,
+    moe_jitter_eps: float = None,
+    moe_normalize_expert_weights: int = None,
+    uniform_expert_assignment: bool = False,
+    training: bool = False,
+    w1: torch.Tensor = None,
+    w2: torch.Tensor = None,
+    w1_bias: torch.Tensor = None,
+    w2_bias: torch.Tensor = None,
+    gradient_scale: Optional[float] = None,
+    alpha: float = 1.702,
+    sort_end_bit: int = 0,
+    expert_parallel_group: torch.distributed.ProcessGroup = None,
+    moe_capacity_factor: float = 1.0,
+    moe_expert_model_parallelism: bool = False,
+    forward_fn: Any = None,
+    hidden_size: int = None,
+    mlp_impl: str = "grouped",
+    # Shared expert parameters
+    shared_up_proj_weight: Optional[torch.Tensor] = None,
+    shared_down_proj_weight: Optional[torch.Tensor] = None,
+    shared_up_proj_bias: Optional[torch.Tensor] = None,
+    shared_down_proj_bias: Optional[torch.Tensor] = None,
+    shared_expert_weighted_sum: bool = False,
+    shared_activation_fn: Optional[Any] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    # First, compute regular MoE forward pass
+    expert_out, expert_weights, router_scores = moe_forward(
+        x=x,
+        router_weight=router_weight,
+        moe_top_k=moe_top_k,
+        moe_num_experts=moe_num_experts,
+        moe_jitter_eps=moe_jitter_eps,
+        moe_normalize_expert_weights=moe_normalize_expert_weights,
+        uniform_expert_assignment=uniform_expert_assignment,
+        training=training,
+        w1=w1,
+        w2=w2,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        gradient_scale=gradient_scale,
+        alpha=alpha,
+        sort_end_bit=sort_end_bit,
+        expert_parallel_group=expert_parallel_group,
+        moe_capacity_factor=moe_capacity_factor,
+        moe_expert_model_parallelism=moe_expert_model_parallelism,
+        forward_fn=forward_fn,
+        hidden_size=hidden_size,
+        mlp_impl=mlp_impl,
+    )
+    
+    # If shared expert weights provided, compute shared expert output
+    if shared_up_proj_weight is not None and shared_down_proj_weight is not None:
+        shared_expert_out = shared_mlp_forward(
+            x=x,
+            up_proj_weight=shared_up_proj_weight,
+            down_proj_weight=shared_down_proj_weight,
+            up_proj_bias=shared_up_proj_bias,
+            down_proj_bias=shared_down_proj_bias,
+            activation_fn=shared_activation_fn,
+            gradient_scale=gradient_scale,
+        )
+        
+        # Combine expert outputs
+        combined_out = combine_expert_shared_outputs(
+            shared_expert_out=shared_expert_out,
+            expert_out=expert_out,
+            shared_expert_weighted_sum=shared_expert_weighted_sum,
+            moe_top_k=moe_top_k,
+        )
+        
+        return combined_out, expert_weights, router_scores
+    
+    # Return regular MoE output if no shared expert
+    return expert_out, expert_weights, router_scores
+
+
+def create_shared_expert_weights(
+    hidden_size: int,
+    shared_expert_hidden_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    init_method: Any,
+    output_layer_init_method: Any = None,
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+    if output_layer_init_method is None:
+        output_layer_init_method = init_method
+        
+    # Create weight tensors
+    up_proj_weight = torch.empty(
+        shared_expert_hidden_size,
+        hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    down_proj_weight = torch.empty(
+        hidden_size,
+        shared_expert_hidden_size,
+        device=device,
+        dtype=dtype,
+    )
+    
+    # Initialize weights
+    init_method(up_proj_weight)
+    output_layer_init_method(down_proj_weight)
+    
+    # No bias by default
+    return up_proj_weight, down_proj_weight, None, None
+
+# HACK: Extract device_mesh from pre-hook closure - required for transformers integration
+# This exists because device_mesh is trapped in hook closures with no model attribute
+# Fragile - breaks if hook structure changes or Python internals change
+# TODO: Replace with a more robust solution when available
+def get_device_mesh(model):
+    # Extract device_mesh from child's unused pre_hook closure
+    try:
+        # Find the pre-hook that contains 'device_mesh' in its closure
+        hook = next(h for h in model.experts._forward_pre_hooks.values() if 'device_mesh' in h.__code__.co_freevars)
+        # Extract the device_mesh from the closure
+        return hook.__closure__[hook.__code__.co_freevars.index('device_mesh')].cell_contents
+    except Exception:
+        return None
+
+
 class MegaBlocksMoeMLP(torch.nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -691,8 +881,12 @@ class MegaBlocksMoeMLP(torch.nn.Module):
         moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
         moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
         uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
-        
+
         expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
         has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
         forward_fn = parallel_forward_once if has_parallel else forward_once
         
@@ -722,4 +916,86 @@ class MegaBlocksMoeMLP(torch.nn.Module):
             hidden_size=self.experts.hidden_size,
             mlp_impl=mlp_impl,
         )
+        return output, expert_weights_out
+
+
+class MegaBlocksMoeMLPWithSharedExpert(MegaBlocksMoeMLP):
+    
+    def __init__(self):
+        super().__init__()
+        # Shared expert weights will be set by the user
+        self.shared_up_proj_weight = None
+        self.shared_down_proj_weight = None
+        self.shared_up_proj_bias = None
+        self.shared_down_proj_bias = None
+        self.shared_expert_weighted_sum = False
+        self.shared_activation_fn = None
+        
+    def set_shared_expert_weights(
+        self,
+        up_proj_weight: torch.Tensor,
+        down_proj_weight: torch.Tensor,
+        up_proj_bias: Optional[torch.Tensor] = None,
+        down_proj_bias: Optional[torch.Tensor] = None,
+        weighted_sum: bool = False,
+        activation_fn: Optional[Any] = None,
+    ):
+        self.shared_up_proj_weight = up_proj_weight
+        self.shared_down_proj_weight = down_proj_weight
+        self.shared_up_proj_bias = up_proj_bias
+        self.shared_down_proj_bias = down_proj_bias
+        self.shared_expert_weighted_sum = weighted_sum
+        self.shared_activation_fn = activation_fn
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        moe_top_k = getattr(self.router, "top_k", 4)
+        moe_num_experts = getattr(self.experts, "num_experts", 128)
+        gradient_scale = getattr(self.experts, "gradient_scale", None)
+        alpha = getattr(self.experts, "alpha", 1.0)
+        moe_capacity_factor = getattr(self.experts, "capacity_factor", 1.0)
+        moe_jitter_eps = getattr(self.experts, "jitter_eps", None)
+        moe_normalize_expert_weights = getattr(self.experts, "normalize_expert_weights", None)
+        uniform_expert_assignment = getattr(self, "uniform_expert_assignment", False)
+
+        expert_parallel_group = getattr(self, "expert_parallel_group", None)
+        if expert_parallel_group is None:
+            device_mesh = get_device_mesh(self)
+            expert_parallel_group = device_mesh.get_group() if device_mesh else None
+
+        has_parallel = expert_parallel_group is not None and dist.is_initialized() and dist.get_world_size(expert_parallel_group) > 1
+        forward_fn = parallel_forward_once if has_parallel else forward_once
+        
+        sort_end_bit = max(int(torch.ceil(torch.log2(torch.tensor(moe_num_experts)))), 1)
+        mlp_impl = getattr(self, "mlp_impl", "grouped")
+        
+        output, expert_weights_out, *_ = moe_forward_with_shared_expert(
+            x=x,
+            router_weight=self.router.weight,
+            moe_top_k=moe_top_k,
+            moe_num_experts=moe_num_experts,
+            moe_jitter_eps=moe_jitter_eps,
+            moe_normalize_expert_weights=moe_normalize_expert_weights,
+            uniform_expert_assignment=uniform_expert_assignment,
+            training=self.training,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            w1_bias=self.experts.gate_up_proj_bias,
+            w2_bias=self.experts.down_proj_bias,
+            gradient_scale=gradient_scale,
+            alpha=alpha,
+            sort_end_bit=sort_end_bit,
+            expert_parallel_group=expert_parallel_group,
+            moe_capacity_factor=moe_capacity_factor,
+            moe_expert_model_parallelism=has_parallel,
+            forward_fn=forward_fn,
+            hidden_size=self.experts.hidden_size,
+            mlp_impl=mlp_impl,
+            # Shared expert parameters
+            shared_up_proj_weight=self.shared_up_proj_weight,
+            shared_down_proj_weight=self.shared_down_proj_weight,
+            shared_up_proj_bias=self.shared_up_proj_bias,
+            shared_down_proj_bias=self.shared_down_proj_bias,
+            shared_expert_weighted_sum=self.shared_expert_weighted_sum,
+            shared_activation_fn=self.shared_activation_fn,
+        )
         return output, expert_weights_out