Update Wall-X to 1.1.0 (#104)

wall_x/model/core/ops/moe.py

@@ -0,0 +1,97 @@
+"""MoE routing operators (permute/unpermute).
+
+These operators reorder tokens by expert assignment for efficient
+Mixture-of-Experts processing.
+"""
+
+import logging
+
+import torch
+
+from wall_x.model.core.ops.base import OpsProxy
+
+logger = logging.getLogger(__name__)
+
+
+class PermuteOp(OpsProxy):
+    """Reorder tokens by expert assignment for MoE processing.
+
+    Signature: permute(tokens, indices, num_out_tokens=None, max_token_num=0) -> (permuted_tokens, sorted_indices)
+    """
+
+    @property
+    def _external_accel_name(self):
+        return "permute"
+
+    def _get_cuda_kernel(self):
+        try:
+            from wall_x.model.core.ops._cuda_wrappers import permute_kernel
+
+            return permute_kernel
+        except ImportError:
+            return None
+        except Exception as e:
+            logger.warning("PermuteOp: CUDA kernel load failed: %s", e)
+            return None
+
+    def _pytorch_fallback(self, tokens, indices, num_out_tokens=None, max_token_num=0):
+        """PyTorch fallback for permute.
+
+        Args:
+            num_out_tokens: If set, truncate output to this many tokens (matches
+                external_accel behavior of discarding padding expert assignments).
+            max_token_num: Unused, kept for external_accel API compatibility.
+        """
+        del max_token_num  # unused, external_accel API compat
+        if indices.dim() == 1:
+            indices = indices.view(-1, 1)
+        expand_factor = indices.size(1)
+        flatten_indices = indices.view(-1)
+        # Keep int64 throughout to avoid precision loss on large token counts
+        sorted_indices = torch.argsort(flatten_indices, stable=True)
+        permuted_tokens = tokens.index_select(0, sorted_indices // expand_factor)
+        if num_out_tokens is not None:
+            permuted_tokens = permuted_tokens[:num_out_tokens]
+            sorted_indices = sorted_indices[:num_out_tokens]
+        return permuted_tokens, sorted_indices
+
+
+class UnpermuteOp(OpsProxy):
+    """Restore tokens to original order after MoE processing.
+
+    Signature: unpermute(permuted_tokens, sorted_indices, probs=None) -> restored_tokens
+    """
+
+    @property
+    def _external_accel_name(self):
+        return "unpermute"
+
+    def _get_cuda_kernel(self):
+        try:
+            from wall_x.model.core.ops._cuda_wrappers import unpermute_kernel
+
+            return unpermute_kernel
+        except ImportError:
+            return None
+        except Exception as e:
+            logger.warning("UnpermuteOp: CUDA kernel load failed: %s", e)
+            return None
+
+    def _pytorch_fallback(self, permuted_tokens, sorted_indices, probs=None):
+        if probs is not None:
+            merge_factor = probs.size(1)
+        else:
+            merge_factor = 1
+        unpermuted_tokens = torch.zeros_like(permuted_tokens)
+        unpermuted_tokens.index_copy_(0, sorted_indices.long(), permuted_tokens)
+        unpermuted_tokens = unpermuted_tokens.reshape(
+            -1, merge_factor, permuted_tokens.size(-1)
+        )
+        if probs is not None:
+            unpermuted_tokens = unpermuted_tokens * probs.unsqueeze(-1)
+        unpermuted_tokens = unpermuted_tokens.sum(dim=1)
+        return unpermuted_tokens
+
+
+permute = PermuteOp()
+unpermute = UnpermuteOp()