Build

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

@@ -1 +1,135 @@
-import moe._custom_ops as ops
+from typing import TYPE_CHECKING
+
+import torch
+
+# neuron has torch version that doesn't even have impl_abstract
+if TYPE_CHECKING:
+
+    def register_fake(fn):
+        return lambda name: fn
+
+else:
+    try:
+        from torch.library import register_fake
+    except ImportError:
+        from torch.library import impl_abstract as register_fake
+
+from ._ops import add_op_namespace_prefix, ops
+from .fused_marlin_moe import fused_marlin_moe
+from .fused_moe import fused_moe, fused_topk, grouped_topk
+from .scalar_type import ScalarType, scalar_types
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.gptq_marlin_repack(
+            b_q_weight[e], perm[e], size_k, size_n, num_bits
+        )
+    return output
+
+
+def awq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
+    return output
+
+
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    ops.moe_sum(input, output)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indicies: torch.Tensor,
+    gating_output: float,
+) -> None:
+    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
+
+
+if hasattr(ops, "marlin_gemm_moe"):
+
+    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
+    def marlin_gemm_moe_fake(
+        a: torch.Tensor,
+        b_q_weights: torch.Tensor,
+        sorted_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_zero_points: torch.Tensor,
+        g_idx: torch.Tensor,
+        perm: torch.Tensor,
+        workspace: torch.Tensor,
+        b_q_type: ScalarType,
+        size_m: torch.SymInt,
+        size_n: torch.SymInt,
+        size_k: torch.SymInt,
+        is_k_full: bool,
+        num_experts: int,
+        topk: int,
+        moe_block_size: int,
+        replicate_input: bool,
+        apply_weights: bool,
+    ) -> torch.Tensor:
+        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(out, x)
+    return out
+
+
+__all__ = [
+    "gptq_marlin_moe_repack",
+    "awq_marlin_moe_repack",
+    "fused_marlin_moe",
+    "moe_sum",
+    "moe_align_block_size",
+    "topk_softmax",
+    "fused_moe",
+]

build/torch24-cxx11-cu118-x86_64-linux/moe/_custom_ops.py

@@ -1,135 +0,0 @@
-from typing import TYPE_CHECKING
-
-import torch
-
-# neuron has torch version that doesn't even have impl_abstract
-if TYPE_CHECKING:
-
-    def register_fake(fn):
-        return lambda name: fn
-
-else:
-    try:
-        from torch.library import register_fake
-    except ImportError:
-        from torch.library import impl_abstract as register_fake
-
-try:
-    from ._ops import ops, add_op_namespace_prefix
-except ImportError as e:
-    # Fallback for local development.
-    try:
-        import _moe
-
-        ops = torch._moe
-
-        def add_op_namespace_prefix(op_name: str):
-            return f"_quantization::{op_name}"
-
-    except ImportError:
-        raise e
-
-from .scalar_type import ScalarType
-
-def gptq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.gptq_marlin_repack(
-            b_q_weight[e], perm[e], size_k, size_n, num_bits
-        )
-    return output
-
-
-def awq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
-    return output
-
-
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
-    ops.moe_sum(input, output)
-
-
-def moe_align_block_size(
-    topk_ids: torch.Tensor,
-    num_experts: int,
-    block_size: int,
-    sorted_token_ids: torch.Tensor,
-    experts_ids: torch.Tensor,
-    num_tokens_post_pad: torch.Tensor,
-) -> None:
-    ops.moe_align_block_size(
-        topk_ids,
-        num_experts,
-        block_size,
-        sorted_token_ids,
-        experts_ids,
-        num_tokens_post_pad,
-    )
-
-
-def topk_softmax(
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    token_expert_indicies: torch.Tensor,
-    gating_output: float,
-) -> None:
-    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
-
-if hasattr(ops, "marlin_gemm_moe"):
-
-    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
-
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out

build/{torch24-cxx11-cu124-x86_64-linux/moe/_moe_0_0_1.abi3.so → torch24-cxx11-cu118-x86_64-linux/moe/_moe_wtjc356yopxde.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c1dd7f6fb98ad1ed39a402e1e42f3231645949dcc5cef28739f4e093883e0184
-size 84063064
+oid sha256:6330aa66b63067a8c9c031419773dc47e8853a717ef20b03c57df76660188831
+size 84165640

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

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

build/torch24-cxx11-cu118-x86_64-linux/moe/fused_marlin_moe.py

@@ -7,7 +7,7 @@ import torch
 
 from .fused_moe import fused_topk, moe_align_block_size, try_get_optimal_moe_config
 from .scalar_type import scalar_types
-import moe._custom_ops as ops
+import moe as ops
 
 
 def get_scalar_type(num_bits: int, has_zp: bool):

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

@@ -11,7 +11,7 @@ import triton.language as tl
 
 from .platforms import current_platform
 from .fp8 import scaled_fp8_quant
-import moe._custom_ops as ops
+import moe as ops
 
 VLLM_FUSED_MOE_CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768"))
 

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

@@ -1 +1,135 @@
-import moe._custom_ops as ops
+from typing import TYPE_CHECKING
+
+import torch
+
+# neuron has torch version that doesn't even have impl_abstract
+if TYPE_CHECKING:
+
+    def register_fake(fn):
+        return lambda name: fn
+
+else:
+    try:
+        from torch.library import register_fake
+    except ImportError:
+        from torch.library import impl_abstract as register_fake
+
+from ._ops import add_op_namespace_prefix, ops
+from .fused_marlin_moe import fused_marlin_moe
+from .fused_moe import fused_moe, fused_topk, grouped_topk
+from .scalar_type import ScalarType, scalar_types
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.gptq_marlin_repack(
+            b_q_weight[e], perm[e], size_k, size_n, num_bits
+        )
+    return output
+
+
+def awq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
+    return output
+
+
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    ops.moe_sum(input, output)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indicies: torch.Tensor,
+    gating_output: float,
+) -> None:
+    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
+
+
+if hasattr(ops, "marlin_gemm_moe"):
+
+    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
+    def marlin_gemm_moe_fake(
+        a: torch.Tensor,
+        b_q_weights: torch.Tensor,
+        sorted_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_zero_points: torch.Tensor,
+        g_idx: torch.Tensor,
+        perm: torch.Tensor,
+        workspace: torch.Tensor,
+        b_q_type: ScalarType,
+        size_m: torch.SymInt,
+        size_n: torch.SymInt,
+        size_k: torch.SymInt,
+        is_k_full: bool,
+        num_experts: int,
+        topk: int,
+        moe_block_size: int,
+        replicate_input: bool,
+        apply_weights: bool,
+    ) -> torch.Tensor:
+        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(out, x)
+    return out
+
+
+__all__ = [
+    "gptq_marlin_moe_repack",
+    "awq_marlin_moe_repack",
+    "fused_marlin_moe",
+    "moe_sum",
+    "moe_align_block_size",
+    "topk_softmax",
+    "fused_moe",
+]

build/torch24-cxx11-cu121-x86_64-linux/moe/_custom_ops.py

@@ -1,135 +0,0 @@
-from typing import TYPE_CHECKING
-
-import torch
-
-# neuron has torch version that doesn't even have impl_abstract
-if TYPE_CHECKING:
-
-    def register_fake(fn):
-        return lambda name: fn
-
-else:
-    try:
-        from torch.library import register_fake
-    except ImportError:
-        from torch.library import impl_abstract as register_fake
-
-try:
-    from ._ops import ops, add_op_namespace_prefix
-except ImportError as e:
-    # Fallback for local development.
-    try:
-        import _moe
-
-        ops = torch._moe
-
-        def add_op_namespace_prefix(op_name: str):
-            return f"_quantization::{op_name}"
-
-    except ImportError:
-        raise e
-
-from .scalar_type import ScalarType
-
-def gptq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.gptq_marlin_repack(
-            b_q_weight[e], perm[e], size_k, size_n, num_bits
-        )
-    return output
-
-
-def awq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
-    return output
-
-
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
-    ops.moe_sum(input, output)
-
-
-def moe_align_block_size(
-    topk_ids: torch.Tensor,
-    num_experts: int,
-    block_size: int,
-    sorted_token_ids: torch.Tensor,
-    experts_ids: torch.Tensor,
-    num_tokens_post_pad: torch.Tensor,
-) -> None:
-    ops.moe_align_block_size(
-        topk_ids,
-        num_experts,
-        block_size,
-        sorted_token_ids,
-        experts_ids,
-        num_tokens_post_pad,
-    )
-
-
-def topk_softmax(
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    token_expert_indicies: torch.Tensor,
-    gating_output: float,
-) -> None:
-    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
-
-if hasattr(ops, "marlin_gemm_moe"):
-
-    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
-
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out

build/{torch24-cxx11-cu118-x86_64-linux/moe/_moe_0_0_1.abi3.so → torch24-cxx11-cu121-x86_64-linux/moe/_moe_fidhfyl4jgbje.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3c1fc3d94e9a7d4b7c0cf13dd3a9633efef2ed265222b17b22d32282818b7bd1
-size 84165576
+oid sha256:b0ca4f733821a564c525a36bb13e35ae960dc1e20f6472b177f67b9b165597ff
+size 84364504

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

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

build/torch24-cxx11-cu121-x86_64-linux/moe/fused_marlin_moe.py

@@ -7,7 +7,7 @@ import torch
 
 from .fused_moe import fused_topk, moe_align_block_size, try_get_optimal_moe_config
 from .scalar_type import scalar_types
-import moe._custom_ops as ops
+import moe as ops
 
 
 def get_scalar_type(num_bits: int, has_zp: bool):

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

@@ -11,7 +11,7 @@ import triton.language as tl
 
 from .platforms import current_platform
 from .fp8 import scaled_fp8_quant
-import moe._custom_ops as ops
+import moe as ops
 
 VLLM_FUSED_MOE_CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768"))
 

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

@@ -1 +1,135 @@
-import moe._custom_ops as ops
+from typing import TYPE_CHECKING
+
+import torch
+
+# neuron has torch version that doesn't even have impl_abstract
+if TYPE_CHECKING:
+
+    def register_fake(fn):
+        return lambda name: fn
+
+else:
+    try:
+        from torch.library import register_fake
+    except ImportError:
+        from torch.library import impl_abstract as register_fake
+
+from ._ops import add_op_namespace_prefix, ops
+from .fused_marlin_moe import fused_marlin_moe
+from .fused_moe import fused_moe, fused_topk, grouped_topk
+from .scalar_type import ScalarType, scalar_types
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.gptq_marlin_repack(
+            b_q_weight[e], perm[e], size_k, size_n, num_bits
+        )
+    return output
+
+
+def awq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
+    return output
+
+
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    ops.moe_sum(input, output)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indicies: torch.Tensor,
+    gating_output: float,
+) -> None:
+    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
+
+
+if hasattr(ops, "marlin_gemm_moe"):
+
+    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
+    def marlin_gemm_moe_fake(
+        a: torch.Tensor,
+        b_q_weights: torch.Tensor,
+        sorted_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_zero_points: torch.Tensor,
+        g_idx: torch.Tensor,
+        perm: torch.Tensor,
+        workspace: torch.Tensor,
+        b_q_type: ScalarType,
+        size_m: torch.SymInt,
+        size_n: torch.SymInt,
+        size_k: torch.SymInt,
+        is_k_full: bool,
+        num_experts: int,
+        topk: int,
+        moe_block_size: int,
+        replicate_input: bool,
+        apply_weights: bool,
+    ) -> torch.Tensor:
+        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(out, x)
+    return out
+
+
+__all__ = [
+    "gptq_marlin_moe_repack",
+    "awq_marlin_moe_repack",
+    "fused_marlin_moe",
+    "moe_sum",
+    "moe_align_block_size",
+    "topk_softmax",
+    "fused_moe",
+]

build/torch24-cxx11-cu124-x86_64-linux/moe/_custom_ops.py

@@ -1,135 +0,0 @@
-from typing import TYPE_CHECKING
-
-import torch
-
-# neuron has torch version that doesn't even have impl_abstract
-if TYPE_CHECKING:
-
-    def register_fake(fn):
-        return lambda name: fn
-
-else:
-    try:
-        from torch.library import register_fake
-    except ImportError:
-        from torch.library import impl_abstract as register_fake
-
-try:
-    from ._ops import ops, add_op_namespace_prefix
-except ImportError as e:
-    # Fallback for local development.
-    try:
-        import _moe
-
-        ops = torch._moe
-
-        def add_op_namespace_prefix(op_name: str):
-            return f"_quantization::{op_name}"
-
-    except ImportError:
-        raise e
-
-from .scalar_type import ScalarType
-
-def gptq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.gptq_marlin_repack(
-            b_q_weight[e], perm[e], size_k, size_n, num_bits
-        )
-    return output
-
-
-def awq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
-    return output
-
-
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
-    ops.moe_sum(input, output)
-
-
-def moe_align_block_size(
-    topk_ids: torch.Tensor,
-    num_experts: int,
-    block_size: int,
-    sorted_token_ids: torch.Tensor,
-    experts_ids: torch.Tensor,
-    num_tokens_post_pad: torch.Tensor,
-) -> None:
-    ops.moe_align_block_size(
-        topk_ids,
-        num_experts,
-        block_size,
-        sorted_token_ids,
-        experts_ids,
-        num_tokens_post_pad,
-    )
-
-
-def topk_softmax(
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    token_expert_indicies: torch.Tensor,
-    gating_output: float,
-) -> None:
-    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
-
-if hasattr(ops, "marlin_gemm_moe"):
-
-    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
-
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out

build/{torch24-cxx98-cu118-x86_64-linux/moe/_moe_0_0_1.abi3.so → torch24-cxx11-cu124-x86_64-linux/moe/_moe_sg5gu4g3brle6.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1994e01d53c190da58a4a864b648421b515e2171abd320184164507e1aa4f1fe
-size 84157816
+oid sha256:0864e745883f687c46c9ce743f1e2887113734c57268b9bc0e290185be28cf65
+size 84063128

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

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

build/torch24-cxx11-cu124-x86_64-linux/moe/fused_marlin_moe.py

@@ -7,7 +7,7 @@ import torch
 
 from .fused_moe import fused_topk, moe_align_block_size, try_get_optimal_moe_config
 from .scalar_type import scalar_types
-import moe._custom_ops as ops
+import moe as ops
 
 
 def get_scalar_type(num_bits: int, has_zp: bool):

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

@@ -11,7 +11,7 @@ import triton.language as tl
 
 from .platforms import current_platform
 from .fp8 import scaled_fp8_quant
-import moe._custom_ops as ops
+import moe as ops
 
 VLLM_FUSED_MOE_CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768"))
 

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

@@ -1 +1,135 @@
-import moe._custom_ops as ops
+from typing import TYPE_CHECKING
+
+import torch
+
+# neuron has torch version that doesn't even have impl_abstract
+if TYPE_CHECKING:
+
+    def register_fake(fn):
+        return lambda name: fn
+
+else:
+    try:
+        from torch.library import register_fake
+    except ImportError:
+        from torch.library import impl_abstract as register_fake
+
+from ._ops import add_op_namespace_prefix, ops
+from .fused_marlin_moe import fused_marlin_moe
+from .fused_moe import fused_moe, fused_topk, grouped_topk
+from .scalar_type import ScalarType, scalar_types
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.gptq_marlin_repack(
+            b_q_weight[e], perm[e], size_k, size_n, num_bits
+        )
+    return output
+
+
+def awq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
+    return output
+
+
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    ops.moe_sum(input, output)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indicies: torch.Tensor,
+    gating_output: float,
+) -> None:
+    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
+
+
+if hasattr(ops, "marlin_gemm_moe"):
+
+    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
+    def marlin_gemm_moe_fake(
+        a: torch.Tensor,
+        b_q_weights: torch.Tensor,
+        sorted_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_zero_points: torch.Tensor,
+        g_idx: torch.Tensor,
+        perm: torch.Tensor,
+        workspace: torch.Tensor,
+        b_q_type: ScalarType,
+        size_m: torch.SymInt,
+        size_n: torch.SymInt,
+        size_k: torch.SymInt,
+        is_k_full: bool,
+        num_experts: int,
+        topk: int,
+        moe_block_size: int,
+        replicate_input: bool,
+        apply_weights: bool,
+    ) -> torch.Tensor:
+        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(out, x)
+    return out
+
+
+__all__ = [
+    "gptq_marlin_moe_repack",
+    "awq_marlin_moe_repack",
+    "fused_marlin_moe",
+    "moe_sum",
+    "moe_align_block_size",
+    "topk_softmax",
+    "fused_moe",
+]

build/torch24-cxx98-cu118-x86_64-linux/moe/_custom_ops.py

@@ -1,135 +0,0 @@
-from typing import TYPE_CHECKING
-
-import torch
-
-# neuron has torch version that doesn't even have impl_abstract
-if TYPE_CHECKING:
-
-    def register_fake(fn):
-        return lambda name: fn
-
-else:
-    try:
-        from torch.library import register_fake
-    except ImportError:
-        from torch.library import impl_abstract as register_fake
-
-try:
-    from ._ops import ops, add_op_namespace_prefix
-except ImportError as e:
-    # Fallback for local development.
-    try:
-        import _moe
-
-        ops = torch._moe
-
-        def add_op_namespace_prefix(op_name: str):
-            return f"_quantization::{op_name}"
-
-    except ImportError:
-        raise e
-
-from .scalar_type import ScalarType
-
-def gptq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.gptq_marlin_repack(
-            b_q_weight[e], perm[e], size_k, size_n, num_bits
-        )
-    return output
-
-
-def awq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
-    return output
-
-
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
-    ops.moe_sum(input, output)
-
-
-def moe_align_block_size(
-    topk_ids: torch.Tensor,
-    num_experts: int,
-    block_size: int,
-    sorted_token_ids: torch.Tensor,
-    experts_ids: torch.Tensor,
-    num_tokens_post_pad: torch.Tensor,
-) -> None:
-    ops.moe_align_block_size(
-        topk_ids,
-        num_experts,
-        block_size,
-        sorted_token_ids,
-        experts_ids,
-        num_tokens_post_pad,
-    )
-
-
-def topk_softmax(
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    token_expert_indicies: torch.Tensor,
-    gating_output: float,
-) -> None:
-    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
-
-if hasattr(ops, "marlin_gemm_moe"):
-
-    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
-
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out

build/{torch24-cxx11-cu121-x86_64-linux/moe/_moe_0_0_1.abi3.so → torch24-cxx98-cu118-x86_64-linux/moe/_moe_v3wdnwni3a5ce.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cd5492f9d9216ee88cfc40f373b19207c8e5f04ba8c55c58aec3ecc9f9ad3239
-size 84364440
+oid sha256:e83b7db92da1ee38a3a4e5a453d4279024e6af95efcf0ad4b34e275029e44729
+size 84157912

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

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

build/torch24-cxx98-cu118-x86_64-linux/moe/fused_marlin_moe.py

@@ -7,7 +7,7 @@ import torch
 
 from .fused_moe import fused_topk, moe_align_block_size, try_get_optimal_moe_config
 from .scalar_type import scalar_types
-import moe._custom_ops as ops
+import moe as ops
 
 
 def get_scalar_type(num_bits: int, has_zp: bool):

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

@@ -11,7 +11,7 @@ import triton.language as tl
 
 from .platforms import current_platform
 from .fp8 import scaled_fp8_quant
-import moe._custom_ops as ops
+import moe as ops
 
 VLLM_FUSED_MOE_CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768"))
 

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

@@ -1 +1,135 @@
-import moe._custom_ops as ops
+from typing import TYPE_CHECKING
+
+import torch
+
+# neuron has torch version that doesn't even have impl_abstract
+if TYPE_CHECKING:
+
+    def register_fake(fn):
+        return lambda name: fn
+
+else:
+    try:
+        from torch.library import register_fake
+    except ImportError:
+        from torch.library import impl_abstract as register_fake
+
+from ._ops import add_op_namespace_prefix, ops
+from .fused_marlin_moe import fused_marlin_moe
+from .fused_moe import fused_moe, fused_topk, grouped_topk
+from .scalar_type import ScalarType, scalar_types
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.gptq_marlin_repack(
+            b_q_weight[e], perm[e], size_k, size_n, num_bits
+        )
+    return output
+
+
+def awq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
+    return output
+
+
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    ops.moe_sum(input, output)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indicies: torch.Tensor,
+    gating_output: float,
+) -> None:
+    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
+
+
+if hasattr(ops, "marlin_gemm_moe"):
+
+    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
+    def marlin_gemm_moe_fake(
+        a: torch.Tensor,
+        b_q_weights: torch.Tensor,
+        sorted_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_zero_points: torch.Tensor,
+        g_idx: torch.Tensor,
+        perm: torch.Tensor,
+        workspace: torch.Tensor,
+        b_q_type: ScalarType,
+        size_m: torch.SymInt,
+        size_n: torch.SymInt,
+        size_k: torch.SymInt,
+        is_k_full: bool,
+        num_experts: int,
+        topk: int,
+        moe_block_size: int,
+        replicate_input: bool,
+        apply_weights: bool,
+    ) -> torch.Tensor:
+        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(out, x)
+    return out
+
+
+__all__ = [
+    "gptq_marlin_moe_repack",
+    "awq_marlin_moe_repack",
+    "fused_marlin_moe",
+    "moe_sum",
+    "moe_align_block_size",
+    "topk_softmax",
+    "fused_moe",
+]

build/torch24-cxx98-cu121-x86_64-linux/moe/_custom_ops.py

@@ -1,135 +0,0 @@
-from typing import TYPE_CHECKING
-
-import torch
-
-# neuron has torch version that doesn't even have impl_abstract
-if TYPE_CHECKING:
-
-    def register_fake(fn):
-        return lambda name: fn
-
-else:
-    try:
-        from torch.library import register_fake
-    except ImportError:
-        from torch.library import impl_abstract as register_fake
-
-try:
-    from ._ops import ops, add_op_namespace_prefix
-except ImportError as e:
-    # Fallback for local development.
-    try:
-        import _moe
-
-        ops = torch._moe
-
-        def add_op_namespace_prefix(op_name: str):
-            return f"_quantization::{op_name}"
-
-    except ImportError:
-        raise e
-
-from .scalar_type import ScalarType
-
-def gptq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.gptq_marlin_repack(
-            b_q_weight[e], perm[e], size_k, size_n, num_bits
-        )
-    return output
-
-
-def awq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
-    return output
-
-
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
-    ops.moe_sum(input, output)
-
-
-def moe_align_block_size(
-    topk_ids: torch.Tensor,
-    num_experts: int,
-    block_size: int,
-    sorted_token_ids: torch.Tensor,
-    experts_ids: torch.Tensor,
-    num_tokens_post_pad: torch.Tensor,
-) -> None:
-    ops.moe_align_block_size(
-        topk_ids,
-        num_experts,
-        block_size,
-        sorted_token_ids,
-        experts_ids,
-        num_tokens_post_pad,
-    )
-
-
-def topk_softmax(
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    token_expert_indicies: torch.Tensor,
-    gating_output: float,
-) -> None:
-    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
-
-if hasattr(ops, "marlin_gemm_moe"):
-
-    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
-
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out

build/torch24-cxx98-cu121-x86_64-linux/moe/_moe_0_0_1.abi3.so

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

build/torch24-cxx98-cu121-x86_64-linux/moe/_moe_hrq7opevcb4ug.abi3.so

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

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

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

build/torch24-cxx98-cu121-x86_64-linux/moe/fused_marlin_moe.py

@@ -7,7 +7,7 @@ import torch
 
 from .fused_moe import fused_topk, moe_align_block_size, try_get_optimal_moe_config
 from .scalar_type import scalar_types
-import moe._custom_ops as ops
+import moe as ops
 
 
 def get_scalar_type(num_bits: int, has_zp: bool):

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

@@ -11,7 +11,7 @@ import triton.language as tl
 
 from .platforms import current_platform
 from .fp8 import scaled_fp8_quant
-import moe._custom_ops as ops
+import moe as ops
 
 VLLM_FUSED_MOE_CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768"))
 

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

@@ -1 +1,135 @@
-import moe._custom_ops as ops
+from typing import TYPE_CHECKING
+
+import torch
+
+# neuron has torch version that doesn't even have impl_abstract
+if TYPE_CHECKING:
+
+    def register_fake(fn):
+        return lambda name: fn
+
+else:
+    try:
+        from torch.library import register_fake
+    except ImportError:
+        from torch.library import impl_abstract as register_fake
+
+from ._ops import add_op_namespace_prefix, ops
+from .fused_marlin_moe import fused_marlin_moe
+from .fused_moe import fused_moe, fused_topk, grouped_topk
+from .scalar_type import ScalarType, scalar_types
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.gptq_marlin_repack(
+            b_q_weight[e], perm[e], size_k, size_n, num_bits
+        )
+    return output
+
+
+def awq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
+    return output
+
+
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    ops.moe_sum(input, output)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indicies: torch.Tensor,
+    gating_output: float,
+) -> None:
+    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
+
+
+if hasattr(ops, "marlin_gemm_moe"):
+
+    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
+    def marlin_gemm_moe_fake(
+        a: torch.Tensor,
+        b_q_weights: torch.Tensor,
+        sorted_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_zero_points: torch.Tensor,
+        g_idx: torch.Tensor,
+        perm: torch.Tensor,
+        workspace: torch.Tensor,
+        b_q_type: ScalarType,
+        size_m: torch.SymInt,
+        size_n: torch.SymInt,
+        size_k: torch.SymInt,
+        is_k_full: bool,
+        num_experts: int,
+        topk: int,
+        moe_block_size: int,
+        replicate_input: bool,
+        apply_weights: bool,
+    ) -> torch.Tensor:
+        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(out, x)
+    return out
+
+
+__all__ = [
+    "gptq_marlin_moe_repack",
+    "awq_marlin_moe_repack",
+    "fused_marlin_moe",
+    "moe_sum",
+    "moe_align_block_size",
+    "topk_softmax",
+    "fused_moe",
+]

build/torch24-cxx98-cu124-x86_64-linux/moe/_custom_ops.py

@@ -1,135 +0,0 @@
-from typing import TYPE_CHECKING
-
-import torch
-
-# neuron has torch version that doesn't even have impl_abstract
-if TYPE_CHECKING:
-
-    def register_fake(fn):
-        return lambda name: fn
-
-else:
-    try:
-        from torch.library import register_fake
-    except ImportError:
-        from torch.library import impl_abstract as register_fake
-
-try:
-    from ._ops import ops, add_op_namespace_prefix
-except ImportError as e:
-    # Fallback for local development.
-    try:
-        import _moe
-
-        ops = torch._moe
-
-        def add_op_namespace_prefix(op_name: str):
-            return f"_quantization::{op_name}"
-
-    except ImportError:
-        raise e
-
-from .scalar_type import ScalarType
-
-def gptq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.gptq_marlin_repack(
-            b_q_weight[e], perm[e], size_k, size_n, num_bits
-        )
-    return output
-
-
-def awq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
-    return output
-
-
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
-    ops.moe_sum(input, output)
-
-
-def moe_align_block_size(
-    topk_ids: torch.Tensor,
-    num_experts: int,
-    block_size: int,
-    sorted_token_ids: torch.Tensor,
-    experts_ids: torch.Tensor,
-    num_tokens_post_pad: torch.Tensor,
-) -> None:
-    ops.moe_align_block_size(
-        topk_ids,
-        num_experts,
-        block_size,
-        sorted_token_ids,
-        experts_ids,
-        num_tokens_post_pad,
-    )
-
-
-def topk_softmax(
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    token_expert_indicies: torch.Tensor,
-    gating_output: float,
-) -> None:
-    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
-
-if hasattr(ops, "marlin_gemm_moe"):
-
-    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
-
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out

build/torch24-cxx98-cu124-x86_64-linux/moe/_moe_0_0_1.abi3.so

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

build/torch24-cxx98-cu124-x86_64-linux/moe/_moe_p3swbnotpexcc.abi3.so

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

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

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

build/torch24-cxx98-cu124-x86_64-linux/moe/fused_marlin_moe.py

@@ -7,7 +7,7 @@ import torch
 
 from .fused_moe import fused_topk, moe_align_block_size, try_get_optimal_moe_config
 from .scalar_type import scalar_types
-import moe._custom_ops as ops
+import moe as ops
 
 
 def get_scalar_type(num_bits: int, has_zp: bool):

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

@@ -11,7 +11,7 @@ import triton.language as tl
 
 from .platforms import current_platform
 from .fp8 import scaled_fp8_quant
-import moe._custom_ops as ops
+import moe as ops
 
 VLLM_FUSED_MOE_CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768"))
 

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

@@ -1 +1,135 @@
-import moe._custom_ops as ops
+from typing import TYPE_CHECKING
+
+import torch
+
+# neuron has torch version that doesn't even have impl_abstract
+if TYPE_CHECKING:
+
+    def register_fake(fn):
+        return lambda name: fn
+
+else:
+    try:
+        from torch.library import register_fake
+    except ImportError:
+        from torch.library import impl_abstract as register_fake
+
+from ._ops import add_op_namespace_prefix, ops
+from .fused_marlin_moe import fused_marlin_moe
+from .fused_moe import fused_moe, fused_topk, grouped_topk
+from .scalar_type import ScalarType, scalar_types
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.gptq_marlin_repack(
+            b_q_weight[e], perm[e], size_k, size_n, num_bits
+        )
+    return output
+
+
+def awq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
+    return output
+
+
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    ops.moe_sum(input, output)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indicies: torch.Tensor,
+    gating_output: float,
+) -> None:
+    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
+
+
+if hasattr(ops, "marlin_gemm_moe"):
+
+    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
+    def marlin_gemm_moe_fake(
+        a: torch.Tensor,
+        b_q_weights: torch.Tensor,
+        sorted_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_zero_points: torch.Tensor,
+        g_idx: torch.Tensor,
+        perm: torch.Tensor,
+        workspace: torch.Tensor,
+        b_q_type: ScalarType,
+        size_m: torch.SymInt,
+        size_n: torch.SymInt,
+        size_k: torch.SymInt,
+        is_k_full: bool,
+        num_experts: int,
+        topk: int,
+        moe_block_size: int,
+        replicate_input: bool,
+        apply_weights: bool,
+    ) -> torch.Tensor:
+        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(out, x)
+    return out
+
+
+__all__ = [
+    "gptq_marlin_moe_repack",
+    "awq_marlin_moe_repack",
+    "fused_marlin_moe",
+    "moe_sum",
+    "moe_align_block_size",
+    "topk_softmax",
+    "fused_moe",
+]

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

@@ -1,135 +0,0 @@
-from typing import TYPE_CHECKING
-
-import torch
-
-# neuron has torch version that doesn't even have impl_abstract
-if TYPE_CHECKING:
-
-    def register_fake(fn):
-        return lambda name: fn
-
-else:
-    try:
-        from torch.library import register_fake
-    except ImportError:
-        from torch.library import impl_abstract as register_fake
-
-try:
-    from ._ops import ops, add_op_namespace_prefix
-except ImportError as e:
-    # Fallback for local development.
-    try:
-        import _moe
-
-        ops = torch._moe
-
-        def add_op_namespace_prefix(op_name: str):
-            return f"_quantization::{op_name}"
-
-    except ImportError:
-        raise e
-
-from .scalar_type import ScalarType
-
-def gptq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.gptq_marlin_repack(
-            b_q_weight[e], perm[e], size_k, size_n, num_bits
-        )
-    return output
-
-
-def awq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
-    return output
-
-
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
-    ops.moe_sum(input, output)
-
-
-def moe_align_block_size(
-    topk_ids: torch.Tensor,
-    num_experts: int,
-    block_size: int,
-    sorted_token_ids: torch.Tensor,
-    experts_ids: torch.Tensor,
-    num_tokens_post_pad: torch.Tensor,
-) -> None:
-    ops.moe_align_block_size(
-        topk_ids,
-        num_experts,
-        block_size,
-        sorted_token_ids,
-        experts_ids,
-        num_tokens_post_pad,
-    )
-
-
-def topk_softmax(
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    token_expert_indicies: torch.Tensor,
-    gating_output: float,
-) -> None:
-    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
-
-if hasattr(ops, "marlin_gemm_moe"):
-
-    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
-
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out

build/torch25-cxx11-cu118-x86_64-linux/moe/_moe_0_0_1.abi3.so

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

build/torch25-cxx11-cu118-x86_64-linux/moe/_moe_nskz7v224zllw.abi3.so

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

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

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

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

@@ -7,7 +7,7 @@ import torch
 
 from .fused_moe import fused_topk, moe_align_block_size, try_get_optimal_moe_config
 from .scalar_type import scalar_types
-import moe._custom_ops as ops
+import moe as ops
 
 
 def get_scalar_type(num_bits: int, has_zp: bool):

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

@@ -11,7 +11,7 @@ import triton.language as tl
 
 from .platforms import current_platform
 from .fp8 import scaled_fp8_quant
-import moe._custom_ops as ops
+import moe as ops
 
 VLLM_FUSED_MOE_CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "32768"))
 

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

@@ -1 +1,135 @@
-import moe._custom_ops as ops
+from typing import TYPE_CHECKING
+
+import torch
+
+# neuron has torch version that doesn't even have impl_abstract
+if TYPE_CHECKING:
+
+    def register_fake(fn):
+        return lambda name: fn
+
+else:
+    try:
+        from torch.library import register_fake
+    except ImportError:
+        from torch.library import impl_abstract as register_fake
+
+from ._ops import add_op_namespace_prefix, ops
+from .fused_marlin_moe import fused_marlin_moe
+from .fused_moe import fused_moe, fused_topk, grouped_topk
+from .scalar_type import ScalarType, scalar_types
+
+
+def gptq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.gptq_marlin_repack(
+            b_q_weight[e], perm[e], size_k, size_n, num_bits
+        )
+    return output
+
+
+def awq_marlin_moe_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    num_experts = b_q_weight.shape[0]
+    assert size_k % 16 == 0
+    output = torch.empty(
+        (num_experts, size_k // 16, size_n * (num_bits // 2)),
+        device=b_q_weight.device,
+        dtype=b_q_weight.dtype,
+    )
+    for e in range(num_experts):
+        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
+    return output
+
+
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    ops.moe_sum(input, output)
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    experts_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_token_ids,
+        experts_ids,
+        num_tokens_post_pad,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    token_expert_indicies: torch.Tensor,
+    gating_output: float,
+) -> None:
+    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
+
+
+if hasattr(ops, "marlin_gemm_moe"):
+
+    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
+    def marlin_gemm_moe_fake(
+        a: torch.Tensor,
+        b_q_weights: torch.Tensor,
+        sorted_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        b_scales: torch.Tensor,
+        b_zero_points: torch.Tensor,
+        g_idx: torch.Tensor,
+        perm: torch.Tensor,
+        workspace: torch.Tensor,
+        b_q_type: ScalarType,
+        size_m: torch.SymInt,
+        size_n: torch.SymInt,
+        size_k: torch.SymInt,
+        is_k_full: bool,
+        num_experts: int,
+        topk: int,
+        moe_block_size: int,
+        replicate_input: bool,
+        apply_weights: bool,
+    ) -> torch.Tensor:
+        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
+
+
+def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
+    ops.silu_and_mul(out, x)
+    return out
+
+
+__all__ = [
+    "gptq_marlin_moe_repack",
+    "awq_marlin_moe_repack",
+    "fused_marlin_moe",
+    "moe_sum",
+    "moe_align_block_size",
+    "topk_softmax",
+    "fused_moe",
+]

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

@@ -1,135 +0,0 @@
-from typing import TYPE_CHECKING
-
-import torch
-
-# neuron has torch version that doesn't even have impl_abstract
-if TYPE_CHECKING:
-
-    def register_fake(fn):
-        return lambda name: fn
-
-else:
-    try:
-        from torch.library import register_fake
-    except ImportError:
-        from torch.library import impl_abstract as register_fake
-
-try:
-    from ._ops import ops, add_op_namespace_prefix
-except ImportError as e:
-    # Fallback for local development.
-    try:
-        import _moe
-
-        ops = torch._moe
-
-        def add_op_namespace_prefix(op_name: str):
-            return f"_quantization::{op_name}"
-
-    except ImportError:
-        raise e
-
-from .scalar_type import ScalarType
-
-def gptq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.gptq_marlin_repack(
-            b_q_weight[e], perm[e], size_k, size_n, num_bits
-        )
-    return output
-
-
-def awq_marlin_moe_repack(
-    b_q_weight: torch.Tensor,
-    perm: torch.Tensor,
-    size_k: int,
-    size_n: int,
-    num_bits: int,
-) -> torch.Tensor:
-    num_experts = b_q_weight.shape[0]
-    assert size_k % 16 == 0
-    output = torch.empty(
-        (num_experts, size_k // 16, size_n * (num_bits // 2)),
-        device=b_q_weight.device,
-        dtype=b_q_weight.dtype,
-    )
-    for e in range(num_experts):
-        output[e] = ops.awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits)
-    return output
-
-
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
-    ops.moe_sum(input, output)
-
-
-def moe_align_block_size(
-    topk_ids: torch.Tensor,
-    num_experts: int,
-    block_size: int,
-    sorted_token_ids: torch.Tensor,
-    experts_ids: torch.Tensor,
-    num_tokens_post_pad: torch.Tensor,
-) -> None:
-    ops.moe_align_block_size(
-        topk_ids,
-        num_experts,
-        block_size,
-        sorted_token_ids,
-        experts_ids,
-        num_tokens_post_pad,
-    )
-
-
-def topk_softmax(
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    token_expert_indicies: torch.Tensor,
-    gating_output: float,
-) -> None:
-    ops.topk_softmax(topk_weights, topk_ids, token_expert_indicies, gating_output)
-
-if hasattr(ops, "marlin_gemm_moe"):
-
-    @register_fake(add_op_namespace_prefix("marlin_gemm_moe"))
-    def marlin_gemm_moe_fake(
-        a: torch.Tensor,
-        b_q_weights: torch.Tensor,
-        sorted_ids: torch.Tensor,
-        topk_weights: torch.Tensor,
-        topk_ids: torch.Tensor,
-        b_scales: torch.Tensor,
-        b_zero_points: torch.Tensor,
-        g_idx: torch.Tensor,
-        perm: torch.Tensor,
-        workspace: torch.Tensor,
-        b_q_type: ScalarType,
-        size_m: torch.SymInt,
-        size_n: torch.SymInt,
-        size_k: torch.SymInt,
-        is_k_full: bool,
-        num_experts: int,
-        topk: int,
-        moe_block_size: int,
-        replicate_input: bool,
-        apply_weights: bool,
-    ) -> torch.Tensor:
-        return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
-
-
-
-def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
-    ops.silu_and_mul(out, x)
-    return out

build/torch25-cxx11-cu121-x86_64-linux/moe/_moe_0_0_1.abi3.so

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

build/torch25-cxx11-cu121-x86_64-linux/moe/_moe_t32bhzwhzero6.abi3.so

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

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

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