Build uploaded using `kernels`.

build/torch210-cxx11-cpu-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3a81c0cc23130a95d05263f0509e8de560183f6472f458f4316c97e6e8d8f533
+oid sha256:1bb9607d2d00b6eb3f3fe58da8dd972deb37b0658b8682807fc2863129f7aa8d
 size 2219056

build/torch210-cxx11-cpu-x86_64-linux/_ops.py

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

build/torch210-cxx11-cpu-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch210-cxx11-cu126-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d482577c55ffe1abd34983ce45eeeb280a817e55f92d6585b5e92173b2860749
+oid sha256:321e1bb305fd100b1abc99234f480634d05a901ee3a758628d94615d535e2caf
 size 15061032

build/torch210-cxx11-cu126-x86_64-linux/_ops.py

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

build/torch210-cxx11-cu126-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch210-cxx11-cu128-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c0876dbd4267e12fa67f24fac60cedbee8e6dd41b85104c4c241b173729bee9a
+oid sha256:83c64c2e54082d931c9fc3027ef6522bf3f3acd4c49d4c5c14dbfcb5ab038b12
 size 21009952

build/torch210-cxx11-cu128-x86_64-linux/_ops.py

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

build/torch210-cxx11-cu128-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch210-cxx11-cu130-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4c7bc97e0aadcd94b0f6d3d7198269823d894fd5a36f6af9744864211ae0fd71
+oid sha256:f48c4762cbfdf923c9547acd7d792dd7edec4bcfe5a857b605ce370f807be23a
 size 12041568

build/torch210-cxx11-cu130-x86_64-linux/_ops.py

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

build/torch210-cxx11-cu130-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch210-cxx11-xpu20253-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dbf6091a3c2622e19367385fb8c82b507f841749bc9c4177421884232856c021
+oid sha256:e840b67c3d3ee92b1150b7c0e4eaab1eda0998347131838eea3bc1bd44049093
 size 4227888

build/torch210-cxx11-xpu20253-x86_64-linux/_ops.py

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

build/torch210-cxx11-xpu20253-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch29-cxx11-cpu-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8b3f1c2f3058c4c5c08291c7a51be003046657e7567454a779911c7cebfdc3d9
+oid sha256:24c19663574a3afb94a458ee318e8b63d47d24f6b1f457a605c115a567810a08
 size 2201176

build/torch29-cxx11-cpu-x86_64-linux/_ops.py

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

build/torch29-cxx11-cpu-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch29-cxx11-cu126-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4d58bdd86403eaa524fac1db9361b0025a175f4b10dcddd8fa0bf99892172e54
+oid sha256:bc4e092bd6f32001e850abf73dd6ee609e9a25800d87fd9e19a0e4a6c30f8e9c
 size 15046808

build/torch29-cxx11-cu126-x86_64-linux/_ops.py

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

build/torch29-cxx11-cu126-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch29-cxx11-cu128-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a5c3c17f0fa54822f12b05fe5c22f8b61ad1a9711a02de13a706e1e8f63e141b
+oid sha256:9018001f72f4a1b7f364d1ca582d8a756cbe452ed798efc4c42e74c49ca1839c
 size 20995680

build/torch29-cxx11-cu128-x86_64-linux/_ops.py

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

build/torch29-cxx11-cu128-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch29-cxx11-cu130-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:609492272ed9672ab824abf87b08f078f409696c8db453ccc5f46dff39d84f98
+oid sha256:49caf38e644493142784e8ad8fac70c1ec9f249c798399950f4228570a570c04
 size 12031392

build/torch29-cxx11-cu130-x86_64-linux/_ops.py

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

build/torch29-cxx11-cu130-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(

build/torch29-cxx11-xpu20252-x86_64-linux/{_megablocks_099ac3c.abi3.so → _megablocks_9be3a32.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:82d4807a02abe216da87ac6d4fbbf4870fdefa64ef182d09ab3408528107f08b
+oid sha256:bb6f2e895e92997f9d93107066513438e413bdba0012d0ee59737105b7ff6f1c
 size 4075712

build/torch29-cxx11-xpu20252-x86_64-linux/_ops.py

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

build/torch29-cxx11-xpu20252-x86_64-linux/xpu_fused_moe.py

@@ -3,7 +3,9 @@
 import os
 import torch
 
-from ._ops import ops
+from ._ops import ops, add_op_namespace_prefix
+
+from torch.library import register_fake
 
 
 def resolve_dtensor(weight: torch.Tensor):
@@ -14,74 +16,65 @@ def resolve_dtensor(weight: torch.Tensor):
     return weight
 
 
-# Install meta kernels for torch.compile compatibility
-def _install_xpu_meta_kernels():
-    """Install meta kernels for XPU MoE operations to support torch.compile"""
-    
-    # Patch cutlass_grouped_gemm_interface
-    if hasattr(ops, "cutlass_grouped_gemm_interface"):
-        original_gemm = ops.cutlass_grouped_gemm_interface
-        
-        def gemm_with_meta(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, 
-                          expert_first_token_offset, N, K, num_experts,
-                          is_B_int4, is_B_mxfp4):
-            if torch.compiler.is_compiling():
-                # Meta implementation - ptr_D is the output, return it
-                return ptr_D
-            return original_gemm(ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D,
-                               expert_first_token_offset, N, K, num_experts,
-                               is_B_int4, is_B_mxfp4)
-        
-        ops.cutlass_grouped_gemm_interface = gemm_with_meta
-    
-    # Patch fused_moe_prologue
-    if hasattr(ops, "fused_moe_prologue"):
-        original_prologue = ops.fused_moe_prologue
-        
-        def prologue_with_meta(input, token_selected_experts, token_final_scales,
-                              workspace, hidden_size, inter_size, num_experts_on_rank):
-            if torch.compiler.is_compiling():
-                # Meta implementation - this op modifies workspace in-place
-                return None
-            return original_prologue(input, token_selected_experts, token_final_scales,
-                                    workspace, hidden_size, inter_size, num_experts_on_rank)
-        
-        ops.fused_moe_prologue = prologue_with_meta
-    
-    # Patch moe_gather
-    if hasattr(ops, "moe_gather"):
-        original_gather = ops.moe_gather
-        
-        def gather_with_meta(output, moe_output, topk_weights, 
-                            unpermuted_row_to_permuted_row, num_experts):
-            if torch.compiler.is_compiling():
-                # Meta implementation - output is modified in-place
-                return None
-            return original_gather(output, moe_output, topk_weights,
-                                  unpermuted_row_to_permuted_row, num_experts)
-        
-        ops.moe_gather = gather_with_meta
-    
-    # Patch activation ops
-    for act_name in ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul", 
-                     "gelu_fast", "gelu_new", "gelu_quick", "mul_and_silu",
-                     "swigluoai_and_mul"]:
-        if hasattr(ops, act_name):
-            original_act = getattr(ops, act_name)
-            
-            def make_act_wrapper(orig_fn):
-                def act_with_meta(*args, **kwargs):
-                    if torch.compiler.is_compiling():
-                        # Meta implementation - in-place ops, return None
-                        return None
-                    return orig_fn(*args, **kwargs)
-                return act_with_meta
-            
-            setattr(ops, act_name, make_act_wrapper(original_act))
-
-
-# Install meta kernels on module load
-_install_xpu_meta_kernels()
+# Register fake/meta kernels for torch.compile compatibility
+def _register_xpu_fake_kernels():
+    """Register fake kernels for XPU MoE operations to support torch.compile."""
+
+    def _register_if_available(op_name, fn):
+        if hasattr(ops, op_name):
+            register_fake(add_op_namespace_prefix(op_name))(fn)
+
+    _register_if_available(
+        "cutlass_grouped_gemm_interface",
+        lambda ptr_A, ptr_B, ptr_scales, ptr_bias, ptr_D, expert_first_token_offset, N, K, num_experts, is_B_int4, is_B_mxfp4: ptr_D,
+    )
+
+    _register_if_available(
+        "fused_moe_prologue",
+        lambda input, token_selected_experts, token_final_scales, workspace, hidden_size, inter_size, num_experts_on_rank: None,
+    )
+
+    _register_if_available(
+        "moe_gather",
+        lambda output, moe_output, topk_weights, unpermuted_row_to_permuted_row, num_experts: None,
+    )
+
+    _register_if_available(
+        "silu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "mul_and_silu",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_tanh_and_mul",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_fast",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_new",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "gelu_quick",
+        lambda out, input: None,
+    )
+    _register_if_available(
+        "swigluoai_and_mul",
+        lambda out, input, alpha=1.702, limit=7.0: None,
+    )
+
+
+# Register fake kernels on module load
+_register_xpu_fake_kernels()
 
 
 # default
@@ -151,6 +144,21 @@ def compute_num_tokens_per_block(num_tokens, num_experts_per_node):
     return 1024
 
 
+def _bytes_to_typed_tensor(byte_tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """Reinterpret a uint8 buffer as a typed tensor by copying bytes.
+
+    This avoids `Tensor.view(dtype)` which can fail under torch.compile
+    constant folding when shape divisibility is not proven.
+    """
+    if byte_tensor.dtype != torch.uint8:
+        raise ValueError("byte_tensor must be uint8")
+    itemsize = torch.empty((), dtype=dtype).element_size()
+    numel = byte_tensor.numel() // itemsize
+    out = torch.empty((numel,), dtype=dtype, device=byte_tensor.device)
+    out.view(torch.uint8).copy_(byte_tensor.contiguous())
+    return out
+
+
 def implement_zp(qweight):
     # change u4 to s4 to avoid zero point in gemm kernel
     # only support default zero point now
@@ -321,7 +329,7 @@ def xpu_fused_moe(hidden_states,
     config_ws("permuted_token_final_scales", permuted_token_final_scales_size)
     config_ws("overlapped_gemm1_gemm2_inputs", permuted_data_size)
 
-    workspace = torch.zeros(map_offset,
+    workspace = torch.empty(map_offset,
                             dtype=torch.uint8,
                             device=hidden_states.device)
     if topk_ids.dtype == torch.int32:
@@ -335,14 +343,25 @@ def xpu_fused_moe(hidden_states,
         inter_size=inter_size,
         num_experts_on_rank=num_experts_per_node)
 
-    expert_first_token_offset = workspace[
+    expert_first_token_offset_bytes = workspace[
         ws_map["expert_first_token_offset"][1]:
         ws_map["expert_first_token_offset"][1] +
-        expert_first_token_offset_size].view(torch.int64)
-    unpermuted_row_to_permuted_row = workspace[
+        expert_first_token_offset_size]
+    unpermuted_row_to_permuted_row_bytes = workspace[
         ws_map["unpermuted_row_to_permuted_row"][1]:
         ws_map["unpermuted_row_to_permuted_row"][1] +
-        src_to_dest_map_size].view(torch.int32)
+        src_to_dest_map_size]
+
+    if torch.compiler.is_compiling():
+        expert_first_token_offset = _bytes_to_typed_tensor(
+            expert_first_token_offset_bytes, torch.int64
+        )
+        unpermuted_row_to_permuted_row = _bytes_to_typed_tensor(
+            unpermuted_row_to_permuted_row_bytes, torch.int32
+        )
+    else:
+        expert_first_token_offset = expert_first_token_offset_bytes.view(torch.int64)
+        unpermuted_row_to_permuted_row = unpermuted_row_to_permuted_row_bytes.view(torch.int32)
     gemm1_input = workspace[ws_map["overlapped_gemm1_gemm2_inputs"][1]:
                             ws_map["overlapped_gemm1_gemm2_inputs"][1] +
                             permuted_data_size].view(hidden_states.dtype).view(