Add Wall-X serving and Turtle2 TCP WebSocket bridge

wall_x/model/core/ops/moe.py

@@ -5,6 +5,7 @@ Mixture-of-Experts processing.
 """
 
 import logging
+import os
 
 import torch
 
@@ -13,6 +14,14 @@ from wall_x.model.core.ops.base import OpsProxy
 logger = logging.getLogger(__name__)
 
 
+def _force_pytorch_backend() -> bool:
+    return os.environ.get("WALL_X_FORCE_TORCH_MOE_OPS", "").lower() in {
+        "1",
+        "true",
+        "yes",
+    }
+
+
 class PermuteOp(OpsProxy):
     """Reorder tokens by expert assignment for MoE processing.
 
@@ -21,9 +30,11 @@ class PermuteOp(OpsProxy):
 
     @property
     def _external_accel_name(self):
-        return "permute"
+        return None if _force_pytorch_backend() else "permute"
 
     def _get_cuda_kernel(self):
+        if _force_pytorch_backend():
+            return None
         try:
             from wall_x.model.core.ops._cuda_wrappers import permute_kernel
 
@@ -64,9 +75,11 @@ class UnpermuteOp(OpsProxy):
 
     @property
     def _external_accel_name(self):
-        return "unpermute"
+        return None if _force_pytorch_backend() else "unpermute"
 
     def _get_cuda_kernel(self):
+        if _force_pytorch_backend():
+            return None
         try:
             from wall_x.model.core.ops._cuda_wrappers import unpermute_kernel
 

wall_x/model/core/vla_mixin.py

@@ -321,6 +321,20 @@ class ActionModelMixMin:
             noisy_action_emb = noisy_action_emb.to(
                 inputs_embeds.device, inputs_embeds.dtype
             )
+            action_token_count = int(mask.sum().item())
+            if noisy_action_emb.ndim != 3 or (
+                action_token_count * inputs_embeds.shape[-1]
+                != noisy_action_emb.numel()
+            ):
+                raise ValueError(
+                    "Flow action placeholder/embedding mismatch: "
+                    f"input_ids={tuple(input_ids.shape)}, "
+                    f"action_tokens={action_token_count}, "
+                    f"action_chunk={tuple(action_chunk.shape)}, "
+                    f"inputs_embeds={tuple(inputs_embeds.shape)}, "
+                    f"noisy_action_emb={tuple(noisy_action_emb.shape)}, "
+                    f"configured_horizon={self.config.action_horizon_flow}"
+                )
             inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb)
 
         return inputs_embeds, flow, adarms_cond
@@ -530,7 +544,10 @@ class ActionGenerationMixin(GenerationMixin):
             target_modules=target_modules,
             lora_dropout=lora_dropout,
             bias="none",
-            task_type="CAUSAL_LM",
+            # ``self.model`` is Wall-X's decoder, not the outer generation model.
+            # Leave task_type unset so PEFT uses its generic wrapper and forwards
+            # Wall-X's custom MoE arguments unchanged.
+            task_type=None,
         )
         self.model = get_peft_model(self.model, config)
         # Log trainable parameter information