diff --git a/build/torch25-cxx11-cu118-x86_64-linux/moe/__init__.py b/build/torch25-cxx11-cu118-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch25-cxx11-cu118-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch25-cxx11-cu118-x86_64-linux/moe/_moe_2f2wzwk42r5t2.abi3.so b/build/torch25-cxx11-cu118-x86_64-linux/moe/_moe_2f2wzwk42r5t2.abi3.so deleted file mode 100755 index b1f0ac7d52d2cbb7b49dd4e3e23eaf0b6acd3364..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu118-x86_64-linux/moe/_moe_2f2wzwk42r5t2.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:020e963e8c6209e6afaaa5e8d22a8aeeabda5695daeddc50601a4079197cd267 -size 85827728 diff --git a/build/torch25-cxx11-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch25-cxx11-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..8fdb7912fde97381917330860d2c62c0d2b740f8 --- /dev/null +++ b/build/torch25-cxx11-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8c3f7e77fec716ceef6ecbc561aca6c77357f00fd24b17b4f285249fda0422e8 +size 85827648 diff --git a/build/torch25-cxx11-cu118-x86_64-linux/moe/_ops.py b/build/torch25-cxx11-cu118-x86_64-linux/moe/_ops.py index 83a6a6a42d633c9b40e263b40b028086d2609b80..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch25-cxx11-cu118-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx11-cu118-x86_64-linux/moe/fused_moe.py b/build/torch25-cxx11-cu118-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch25-cxx11-cu118-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch25-cxx11-cu118-x86_64-linux/moe/layers.py b/build/torch25-cxx11-cu118-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch25-cxx11-cu121-x86_64-linux/moe/__init__.py b/build/torch25-cxx11-cu121-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch25-cxx11-cu121-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch25-cxx11-cu121-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch25-cxx11-cu121-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..b0a35c37b686f9ab537bce3187a77a823762e4fd --- /dev/null +++ b/build/torch25-cxx11-cu121-x86_64-linux/moe/_moe_3c64348.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6a8d550f099a536818532aa332a83a9d9eb106cb7431ed29f19bb29b9c2d269e +size 86034480 diff --git a/build/torch25-cxx11-cu121-x86_64-linux/moe/_moe_qx7m4hiw6tx7s.abi3.so b/build/torch25-cxx11-cu121-x86_64-linux/moe/_moe_qx7m4hiw6tx7s.abi3.so deleted file mode 100755 index cfdb823fabc296c258f58ce8e03a347be7eb558f..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu121-x86_64-linux/moe/_moe_qx7m4hiw6tx7s.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:92ef82e1e2f8ca4513fdde1864278d1fb69a3ef4111741c34dee8d386f1825a7 -size 86030432 diff --git a/build/torch25-cxx11-cu121-x86_64-linux/moe/_ops.py b/build/torch25-cxx11-cu121-x86_64-linux/moe/_ops.py index 0d77c5d3e29106cf62a45153770fafbff59b2932..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch25-cxx11-cu121-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx11-cu121-x86_64-linux/moe/fused_moe.py b/build/torch25-cxx11-cu121-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch25-cxx11-cu121-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch25-cxx11-cu121-x86_64-linux/moe/layers.py b/build/torch25-cxx11-cu121-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch25-cxx11-cu124-x86_64-linux/moe/__init__.py b/build/torch25-cxx11-cu124-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch25-cxx11-cu124-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch25-cxx11-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch25-cxx11-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..cd3abe819249ac59a4275791b8a862f1b4fcfb51 --- /dev/null +++ b/build/torch25-cxx11-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:37f4cdf7ef050df0b438f7af05046fa52512ef38f5dcde712d60c22d41d6ea41 +size 85732952 diff --git a/build/torch25-cxx11-cu124-x86_64-linux/moe/_moe_ctvji3e7dq64w.abi3.so b/build/torch25-cxx11-cu124-x86_64-linux/moe/_moe_ctvji3e7dq64w.abi3.so deleted file mode 100755 index db5c3be15bc329bc0aa5b87d34223b747751484e..0000000000000000000000000000000000000000 --- a/build/torch25-cxx11-cu124-x86_64-linux/moe/_moe_ctvji3e7dq64w.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:1d0460b6a9520a8b77008965a7c732ad448fefcec2d10f89d5bd888316cc6057 -size 85737096 diff --git a/build/torch25-cxx11-cu124-x86_64-linux/moe/_ops.py b/build/torch25-cxx11-cu124-x86_64-linux/moe/_ops.py index b22bd5b27938464e6c7359b1974db9b472effa6b..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch25-cxx11-cu124-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx11-cu124-x86_64-linux/moe/fused_moe.py b/build/torch25-cxx11-cu124-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch25-cxx11-cu124-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch25-cxx11-cu124-x86_64-linux/moe/layers.py b/build/torch25-cxx11-cu124-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch25-cxx98-cu118-x86_64-linux/moe/__init__.py b/build/torch25-cxx98-cu118-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch25-cxx98-cu118-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch25-cxx98-cu118-x86_64-linux/moe/_moe_2dhx2xkm6c5wu.abi3.so b/build/torch25-cxx98-cu118-x86_64-linux/moe/_moe_2dhx2xkm6c5wu.abi3.so deleted file mode 100755 index 56dc9c91ddc02b281dcf7c996071bee341ef026c..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu118-x86_64-linux/moe/_moe_2dhx2xkm6c5wu.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:7db27a5f590418a070defd28bb017103c63e3de18ab52e0dc080ad9f22c68233 -size 85823632 diff --git a/build/torch25-cxx98-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch25-cxx98-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..3200e898823da36ba67ae4b83ad8f21822ce0af2 --- /dev/null +++ b/build/torch25-cxx98-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d9d6b4f6ae4102cc11e5c9bdb467967f8ca4a62ab700383d2e76d0457871e991 +size 85819488 diff --git a/build/torch25-cxx98-cu118-x86_64-linux/moe/_ops.py b/build/torch25-cxx98-cu118-x86_64-linux/moe/_ops.py index 97f4f6344f4a61a52c8077cdc7884400e56c558b..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch25-cxx98-cu118-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx98-cu118-x86_64-linux/moe/fused_moe.py b/build/torch25-cxx98-cu118-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch25-cxx98-cu118-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch25-cxx98-cu118-x86_64-linux/moe/layers.py b/build/torch25-cxx98-cu118-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch25-cxx98-cu121-x86_64-linux/moe/__init__.py b/build/torch25-cxx98-cu121-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch25-cxx98-cu121-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch25-cxx98-cu121-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch25-cxx98-cu121-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..b4b61b98277df27702770292107a256af5232f79 --- /dev/null +++ b/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 diff --git a/build/torch25-cxx98-cu121-x86_64-linux/moe/_moe_wfwoejktckaue.abi3.so b/build/torch25-cxx98-cu121-x86_64-linux/moe/_moe_wfwoejktckaue.abi3.so deleted file mode 100755 index a679d38667a74beffaca30bb9c6628c6b7d0b1c0..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu121-x86_64-linux/moe/_moe_wfwoejktckaue.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:e52b7dea31f1c9ee0c6ff81e5eb66f794d8a6475eb145e41f80974a73d80a331 -size 86026776 diff --git a/build/torch25-cxx98-cu121-x86_64-linux/moe/_ops.py b/build/torch25-cxx98-cu121-x86_64-linux/moe/_ops.py index 9aebf2be86d230b6de5510163c7b53dcc3aa7c51..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch25-cxx98-cu121-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx98-cu121-x86_64-linux/moe/fused_moe.py b/build/torch25-cxx98-cu121-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch25-cxx98-cu121-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch25-cxx98-cu121-x86_64-linux/moe/layers.py b/build/torch25-cxx98-cu121-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch25-cxx98-cu124-x86_64-linux/moe/__init__.py b/build/torch25-cxx98-cu124-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch25-cxx98-cu124-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch25-cxx98-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch25-cxx98-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..66b9a676d4a311788348e999cc1a5727b7545730 --- /dev/null +++ b/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 diff --git a/build/torch25-cxx98-cu124-x86_64-linux/moe/_moe_plhnk6yxrdq3c.abi3.so b/build/torch25-cxx98-cu124-x86_64-linux/moe/_moe_plhnk6yxrdq3c.abi3.so deleted file mode 100755 index a16cd30c7ff53b3d73fa081369b6443efa5fb184..0000000000000000000000000000000000000000 --- a/build/torch25-cxx98-cu124-x86_64-linux/moe/_moe_plhnk6yxrdq3c.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:2ea126bf0453d6f6fe7320362d147687ef554aacc37884ec5951ebc19eafe0c6 -size 85729320 diff --git a/build/torch25-cxx98-cu124-x86_64-linux/moe/_ops.py b/build/torch25-cxx98-cu124-x86_64-linux/moe/_ops.py index 5b5d43b13c586ead5f177bcb71ba17c078eb016d..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch25-cxx98-cu124-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch25-cxx98-cu124-x86_64-linux/moe/fused_moe.py b/build/torch25-cxx98-cu124-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch25-cxx98-cu124-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch25-cxx98-cu124-x86_64-linux/moe/layers.py b/build/torch25-cxx98-cu124-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch26-cxx11-cu118-x86_64-linux/moe/__init__.py b/build/torch26-cxx11-cu118-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch26-cxx11-cu118-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch26-cxx11-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch26-cxx11-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..af300e41192d97a3ad467cefb778381f7ecaa9d1 --- /dev/null +++ b/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 diff --git a/build/torch26-cxx11-cu118-x86_64-linux/moe/_moe_mrhezivmofzdg.abi3.so b/build/torch26-cxx11-cu118-x86_64-linux/moe/_moe_mrhezivmofzdg.abi3.so deleted file mode 100755 index 5423719c9ac75f96528fc0b7386a108aedc996b1..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu118-x86_64-linux/moe/_moe_mrhezivmofzdg.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:0ce59ec9d04ecd293feb8e7d4d2d637c441ecbdcec6127e2ab7fba9c868d1ba5 -size 85827872 diff --git a/build/torch26-cxx11-cu118-x86_64-linux/moe/_ops.py b/build/torch26-cxx11-cu118-x86_64-linux/moe/_ops.py index f78c4d78eceaa86a9d245eea4d5562167db8f59b..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch26-cxx11-cu118-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu118-x86_64-linux/moe/fused_moe.py b/build/torch26-cxx11-cu118-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch26-cxx11-cu118-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch26-cxx11-cu118-x86_64-linux/moe/layers.py b/build/torch26-cxx11-cu118-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch26-cxx11-cu124-x86_64-linux/moe/__init__.py b/build/torch26-cxx11-cu124-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch26-cxx11-cu124-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch26-cxx11-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch26-cxx11-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..da09d3efed1cefd76677b34e4b1a498d97510efb --- /dev/null +++ b/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 diff --git a/build/torch26-cxx11-cu124-x86_64-linux/moe/_moe_lvqy7x44edhqo.abi3.so b/build/torch26-cxx11-cu124-x86_64-linux/moe/_moe_lvqy7x44edhqo.abi3.so deleted file mode 100755 index 51f9c5792c1d7bcb03a8906e9bc60e779ba1b343..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu124-x86_64-linux/moe/_moe_lvqy7x44edhqo.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:cc1e383aa51ecc99476634640f580e7b5d89aadc33222c7f428c40a86149b3b5 -size 85733304 diff --git a/build/torch26-cxx11-cu124-x86_64-linux/moe/_ops.py b/build/torch26-cxx11-cu124-x86_64-linux/moe/_ops.py index 510c892dcc2479877c6c2fc5c20f6a534dc90d51..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch26-cxx11-cu124-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu124-x86_64-linux/moe/fused_moe.py b/build/torch26-cxx11-cu124-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch26-cxx11-cu124-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch26-cxx11-cu124-x86_64-linux/moe/layers.py b/build/torch26-cxx11-cu124-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch26-cxx11-cu126-x86_64-linux/moe/__init__.py b/build/torch26-cxx11-cu126-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch26-cxx11-cu126-x86_64-linux/moe/__init__.py +++ b/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", diff --git a/build/torch26-cxx11-cu126-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch26-cxx11-cu126-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..802e4068a69f1f779844baaaba5f42a4720d66e5 --- /dev/null +++ b/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 diff --git a/build/torch26-cxx11-cu126-x86_64-linux/moe/_moe_gasoel7noy6kw.abi3.so b/build/torch26-cxx11-cu126-x86_64-linux/moe/_moe_gasoel7noy6kw.abi3.so deleted file mode 100755 index cd914e2830fbe3fbdcf31c5fa2f37c384d2c36d5..0000000000000000000000000000000000000000 --- a/build/torch26-cxx11-cu126-x86_64-linux/moe/_moe_gasoel7noy6kw.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:71784baa3bf781a0d915f0086fee151817a463d1d96879a7dc2e832b21fb1d8c -size 86192352 diff --git a/build/torch26-cxx11-cu126-x86_64-linux/moe/_ops.py b/build/torch26-cxx11-cu126-x86_64-linux/moe/_ops.py index c69fb498baf329dda803ec0f90dc4b7756fb1ff0..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch26-cxx11-cu126-x86_64-linux/moe/_ops.py +++ b/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}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx11-cu126-x86_64-linux/moe/fused_moe.py b/build/torch26-cxx11-cu126-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch26-cxx11-cu126-x86_64-linux/moe/fused_moe.py +++ b/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, diff --git a/build/torch26-cxx11-cu126-x86_64-linux/moe/layers.py b/build/torch26-cxx11-cu126-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/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() diff --git a/build/torch26-cxx98-cu118-x86_64-linux/moe/__init__.py b/build/torch26-cxx98-cu118-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch26-cxx98-cu118-x86_64-linux/moe/__init__.py +++ b/build/torch26-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", diff --git a/build/torch26-cxx98-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch26-cxx98-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..c91a35bfee9580c12483d0bf267c375e5ab11273 --- /dev/null +++ b/build/torch26-cxx98-cu118-x86_64-linux/moe/_moe_3c64348.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:37b78a69ccbfcb7874d9962c6f930780b5a340a1f270203a0fddc06d474b35a7 +size 85819960 diff --git a/build/torch26-cxx98-cu118-x86_64-linux/moe/_moe_cobe53r755p6a.abi3.so b/build/torch26-cxx98-cu118-x86_64-linux/moe/_moe_cobe53r755p6a.abi3.so deleted file mode 100755 index 0082e984366e264ad72eb429f4e138d45f5cbcaf..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu118-x86_64-linux/moe/_moe_cobe53r755p6a.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:939d0a2b70b45aae887a0c9671707d04586c4e253e955008647a240e16fb6f79 -size 85819984 diff --git a/build/torch26-cxx98-cu118-x86_64-linux/moe/_ops.py b/build/torch26-cxx98-cu118-x86_64-linux/moe/_ops.py index 153d250d92d9f1bf4b6e318287370706aa5cd385..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch26-cxx98-cu118-x86_64-linux/moe/_ops.py +++ b/build/torch26-cxx98-cu118-x86_64-linux/moe/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _moe_cobe53r755p6a -ops = torch.ops._moe_cobe53r755p6a +from . import _moe_3c64348 +ops = torch.ops._moe_3c64348 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_moe_cobe53r755p6a::{op_name}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu118-x86_64-linux/moe/fused_moe.py b/build/torch26-cxx98-cu118-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch26-cxx98-cu118-x86_64-linux/moe/fused_moe.py +++ b/build/torch26-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, diff --git a/build/torch26-cxx98-cu118-x86_64-linux/moe/layers.py b/build/torch26-cxx98-cu118-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/build/torch26-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() diff --git a/build/torch26-cxx98-cu124-x86_64-linux/moe/__init__.py b/build/torch26-cxx98-cu124-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch26-cxx98-cu124-x86_64-linux/moe/__init__.py +++ b/build/torch26-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", diff --git a/build/torch26-cxx98-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch26-cxx98-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..3791faf61392873f82ba687a515e3685b0dc9a0b --- /dev/null +++ b/build/torch26-cxx98-cu124-x86_64-linux/moe/_moe_3c64348.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:21d9aa1eb4ff8e8aaba398dfb26230a666fe648c55cdf27a758f2e75bc51c6c8 +size 85729720 diff --git a/build/torch26-cxx98-cu124-x86_64-linux/moe/_moe_iqmfy23diekyw.abi3.so b/build/torch26-cxx98-cu124-x86_64-linux/moe/_moe_iqmfy23diekyw.abi3.so deleted file mode 100755 index 25f80aad80865365eae32a1609be0219f7e6582e..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu124-x86_64-linux/moe/_moe_iqmfy23diekyw.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:5d54dcb90a0f197e27f52f399f573807ee692c3d8e8dad74b6869bf8292d7c5a -size 85725672 diff --git a/build/torch26-cxx98-cu124-x86_64-linux/moe/_ops.py b/build/torch26-cxx98-cu124-x86_64-linux/moe/_ops.py index 73e0213d234c5717aea5d708cff8e0938f14bce9..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch26-cxx98-cu124-x86_64-linux/moe/_ops.py +++ b/build/torch26-cxx98-cu124-x86_64-linux/moe/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _moe_iqmfy23diekyw -ops = torch.ops._moe_iqmfy23diekyw +from . import _moe_3c64348 +ops = torch.ops._moe_3c64348 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_moe_iqmfy23diekyw::{op_name}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu124-x86_64-linux/moe/fused_moe.py b/build/torch26-cxx98-cu124-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch26-cxx98-cu124-x86_64-linux/moe/fused_moe.py +++ b/build/torch26-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, diff --git a/build/torch26-cxx98-cu124-x86_64-linux/moe/layers.py b/build/torch26-cxx98-cu124-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/build/torch26-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() diff --git a/build/torch26-cxx98-cu126-x86_64-linux/moe/__init__.py b/build/torch26-cxx98-cu126-x86_64-linux/moe/__init__.py index cc806778863c03ccb3157343cd6331c1c6ca332c..137c8d2155a9d7aa4d06e946b4d90cf23dc287a2 100644 --- a/build/torch26-cxx98-cu126-x86_64-linux/moe/__init__.py +++ b/build/torch26-cxx98-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", diff --git a/build/torch26-cxx98-cu126-x86_64-linux/moe/_moe_3c64348.abi3.so b/build/torch26-cxx98-cu126-x86_64-linux/moe/_moe_3c64348.abi3.so new file mode 100755 index 0000000000000000000000000000000000000000..d5a22cf430cd1ad3a91eae92ea0e3791ab2398b0 --- /dev/null +++ b/build/torch26-cxx98-cu126-x86_64-linux/moe/_moe_3c64348.abi3.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:527e9184f7073f2772d03951173b9732909a47fbd72ce70f9b9e96659afad8ff +size 86180544 diff --git a/build/torch26-cxx98-cu126-x86_64-linux/moe/_moe_6xluzhr5x6fw4.abi3.so b/build/torch26-cxx98-cu126-x86_64-linux/moe/_moe_6xluzhr5x6fw4.abi3.so deleted file mode 100755 index 63a2723397b7c031719534dd6c23e3eaa1a85a23..0000000000000000000000000000000000000000 --- a/build/torch26-cxx98-cu126-x86_64-linux/moe/_moe_6xluzhr5x6fw4.abi3.so +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:8e51a9bd8f93268c6cea1b97079588199e3671b6affbd56d60f9c9899a7684ee -size 86184688 diff --git a/build/torch26-cxx98-cu126-x86_64-linux/moe/_ops.py b/build/torch26-cxx98-cu126-x86_64-linux/moe/_ops.py index 82fd79cf07706341ad17f8ea5d841b6d018cd676..6e38efbbebc2e72467f94eae660e797400511686 100644 --- a/build/torch26-cxx98-cu126-x86_64-linux/moe/_ops.py +++ b/build/torch26-cxx98-cu126-x86_64-linux/moe/_ops.py @@ -1,9 +1,9 @@ import torch -from . import _moe_6xluzhr5x6fw4 -ops = torch.ops._moe_6xluzhr5x6fw4 +from . import _moe_3c64348 +ops = torch.ops._moe_3c64348 def add_op_namespace_prefix(op_name: str): """ Prefix op by namespace. """ - return f"_moe_6xluzhr5x6fw4::{op_name}" \ No newline at end of file + return f"_moe_3c64348::{op_name}" \ No newline at end of file diff --git a/build/torch26-cxx98-cu126-x86_64-linux/moe/fused_moe.py b/build/torch26-cxx98-cu126-x86_64-linux/moe/fused_moe.py index af2d798cbe5d7c3c1760ce79f717ab5f6d7700ba..587afcd4e5f57da5c561e13568042f6654d549b4 100644 --- a/build/torch26-cxx98-cu126-x86_64-linux/moe/fused_moe.py +++ b/build/torch26-cxx98-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, diff --git a/build/torch26-cxx98-cu126-x86_64-linux/moe/layers.py b/build/torch26-cxx98-cu126-x86_64-linux/moe/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..fd58a5b9e3699d557f63613a734a21d8068b0749 --- /dev/null +++ b/build/torch26-cxx98-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()