Build

build/torch25-cxx11-cu118-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch25-cxx11-cu118-x86_64-linux/moe/{_moe_2f2wzwk42r5t2.abi3.so → _moe_3c64348.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:020e963e8c6209e6afaaa5e8d22a8aeeabda5695daeddc50601a4079197cd267
-size 85827728
+oid sha256:8c3f7e77fec716ceef6ecbc561aca6c77357f00fd24b17b4f285249fda0422e8
+size 85827648

build/torch25-cxx11-cu118-x86_64-linux/moe/_ops.py

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

build/torch25-cxx11-cu118-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

build/torch25-cxx11-cu118-x86_64-linux/moe/layers.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()

build/torch25-cxx11-cu121-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch25-cxx11-cu121-x86_64-linux/moe/{_moe_qx7m4hiw6tx7s.abi3.so → _moe_3c64348.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:92ef82e1e2f8ca4513fdde1864278d1fb69a3ef4111741c34dee8d386f1825a7
-size 86030432
+oid sha256:6a8d550f099a536818532aa332a83a9d9eb106cb7431ed29f19bb29b9c2d269e
+size 86034480

build/torch25-cxx11-cu121-x86_64-linux/moe/_ops.py

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

build/torch25-cxx11-cu121-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

build/torch25-cxx11-cu121-x86_64-linux/moe/layers.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()

build/torch25-cxx11-cu124-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch25-cxx11-cu124-x86_64-linux/moe/{_moe_ctvji3e7dq64w.abi3.so → _moe_3c64348.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1d0460b6a9520a8b77008965a7c732ad448fefcec2d10f89d5bd888316cc6057
-size 85737096
+oid sha256:37f4cdf7ef050df0b438f7af05046fa52512ef38f5dcde712d60c22d41d6ea41
+size 85732952

build/torch25-cxx11-cu124-x86_64-linux/moe/_ops.py

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

build/torch25-cxx11-cu124-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

build/torch25-cxx11-cu124-x86_64-linux/moe/layers.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()

build/torch25-cxx98-cu118-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch25-cxx98-cu118-x86_64-linux/moe/{_moe_2dhx2xkm6c5wu.abi3.so → _moe_3c64348.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7db27a5f590418a070defd28bb017103c63e3de18ab52e0dc080ad9f22c68233
-size 85823632
+oid sha256:d9d6b4f6ae4102cc11e5c9bdb467967f8ca4a62ab700383d2e76d0457871e991
+size 85819488

build/torch25-cxx98-cu118-x86_64-linux/moe/_ops.py

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

build/torch25-cxx98-cu118-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

build/torch25-cxx98-cu118-x86_64-linux/moe/layers.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()

build/torch25-cxx98-cu121-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch25-cxx98-cu121-x86_64-linux/moe/_moe_3c64348.abi3.so

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

build/torch25-cxx98-cu121-x86_64-linux/moe/_moe_wfwoejktckaue.abi3.so

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

build/torch25-cxx98-cu121-x86_64-linux/moe/_ops.py

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

build/torch25-cxx98-cu121-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

build/torch25-cxx98-cu121-x86_64-linux/moe/layers.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()

build/torch25-cxx98-cu124-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch25-cxx98-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so

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

build/torch25-cxx98-cu124-x86_64-linux/moe/_moe_plhnk6yxrdq3c.abi3.so

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

build/torch25-cxx98-cu124-x86_64-linux/moe/_ops.py

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

build/torch25-cxx98-cu124-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

build/torch25-cxx98-cu124-x86_64-linux/moe/layers.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()

build/torch26-cxx11-cu118-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch26-cxx11-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so

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

build/torch26-cxx11-cu118-x86_64-linux/moe/_moe_mrhezivmofzdg.abi3.so

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

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

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

build/torch26-cxx11-cu118-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

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

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()

build/torch26-cxx11-cu124-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch26-cxx11-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so

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

build/torch26-cxx11-cu124-x86_64-linux/moe/_moe_lvqy7x44edhqo.abi3.so

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

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

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

build/torch26-cxx11-cu124-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

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

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()

build/torch26-cxx11-cu126-x86_64-linux/moe/__init__.py

@@ -1,5 +1,6 @@
 import torch
 
+from . import layers
 from ._ops import ops
 from .fp8_utils import per_token_group_quant_fp8, w8a8_block_fp8_matmul
 from .fused_marlin_moe import fused_marlin_moe
@@ -87,6 +88,7 @@ __all__ = [
     "fused_topk",
     "gptq_marlin_moe_repack",
     "grouped_topk",
+    "layers",
     "moe_align_block_size",
     "moe_sum",
     "per_token_group_quant_fp8",

build/torch26-cxx11-cu126-x86_64-linux/moe/_moe_3c64348.abi3.so

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

build/torch26-cxx11-cu126-x86_64-linux/moe/_moe_gasoel7noy6kw.abi3.so

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

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

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

build/torch26-cxx11-cu126-x86_64-linux/moe/fused_moe.py

@@ -1022,6 +1022,7 @@ def inplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1034,22 +1035,23 @@ def inplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> None:
     fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        True,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1059,6 +1061,7 @@ def outplace_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1071,22 +1074,23 @@ def outplace_fused_experts(
     block_shape: Optional[List[int]] = None,
 ) -> torch.Tensor:
     return fused_experts_impl(
-        hidden_states,
-        w1,
-        w2,
-        topk_weights,
-        topk_ids,
-        False,
-        use_fp8_w8a8,
-        use_int8_w8a16,
-        use_int4_w4a16,
-        w1_scale,
-        w2_scale,
-        w1_zp,
-        w2_zp,
-        a1_scale,
-        a2_scale,
-        block_shape,
+        hidden_states=hidden_states,
+        w1=w1,
+        w2=w2,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=False,
+        activation=activation,
+        use_fp8_w8a8=use_fp8_w8a8,
+        use_int8_w8a16=use_int8_w8a16,
+        use_int4_w4a16=use_int4_w4a16,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_zp=w1_zp,
+        w2_zp=w2_zp,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        block_shape=block_shape,
     )
 
 
@@ -1097,6 +1101,7 @@ def fused_experts(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1110,40 +1115,42 @@ def fused_experts(
 ):
     if inplace:
         inplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
         return hidden_states
     else:
         return outplace_fused_experts(
-            hidden_states,
-            w1,
-            w2,
-            topk_weights,
-            topk_ids,
-            use_fp8_w8a8,
-            use_int8_w8a16,
-            use_int4_w4a16,
-            w1_scale,
-            w2_scale,
-            w1_zp,
-            w2_zp,
-            a1_scale,
-            a2_scale,
-            block_shape,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            use_fp8_w8a8=use_fp8_w8a8,
+            use_int8_w8a16=use_int8_w8a16,
+            use_int4_w4a16=use_int4_w4a16,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
         )
 
 
@@ -1154,6 +1161,7 @@ def fused_experts_impl(
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
     inplace: bool = False,
+    activation: str = "silu",
     use_fp8_w8a8: bool = False,
     use_int8_w8a16: bool = False,
     use_int4_w4a16: bool = False,
@@ -1281,7 +1289,12 @@ def fused_experts_impl(
             block_shape=block_shape,
         )
 
-        ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        if activation == "silu":
+            ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        elif activation == "gelu":
+            ops.gelu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
         invoke_fused_moe_kernel(
             intermediate_cache2,
@@ -1320,6 +1333,7 @@ def fused_moe(
     topk: int,
     renormalize: bool,
     inplace: bool = False,
+    activation: str = "silu",
     use_grouped_topk: bool = False,
     num_expert_group: Optional[int] = None,
     topk_group: Optional[int] = None,
@@ -1404,6 +1418,7 @@ def fused_moe(
         topk_weights,
         topk_ids,
         inplace=inplace,
+        activation=activation,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a16=use_int8_w8a16,
         use_int4_w4a16=use_int4_w4a16,

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

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+from .fused_moe import fused_moe
+
+
+def _llama4_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+    topk_weights, topk_ids = torch.topk(gating_output, topk, dim=1)
+    topk_weights = torch.sigmoid(topk_weights.float()).to(hidden_states.dtype)
+    return topk_weights, topk_ids
+
+
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
+class Llama4TextMoe(nn.Module):
+    experts: nn.Module
+    router: nn.Linear
+    shared_expert: nn.Module
+    top_k: int
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
+
+        _fix_llama4_experts(hidden_states, self.experts)
+
+        router_logits = self.router(hidden_states)
+        out = fused_moe(
+            hidden_states,
+            w1=self.experts.gate_up_proj,
+            w2=self.experts.down_proj,
+            gating_output=router_logits,
+            topk=self.top_k,
+            renormalize=False,
+            custom_routing_function=_llama4_topk,
+        )
+
+        out += self.shared_expert(hidden_states)
+
+        return out, router_logits.t()