Fix Llama expert layout

torch-ext/moe/layers.py

@@ -16,6 +16,14 @@ def _llama4_topk(
     return topk_weights, topk_ids
 
 
+def _fix_llama4_experts(hidden_states: torch.Tensor, experts: nn.Module):
+    if experts.gate_up_proj.shape[-1] != hidden_states.shape[-1]:
+        experts.gate_up_proj = nn.Parameter(
+            experts.gate_up_proj.transpose(1, 2).contiguous()
+        )
+        experts.down_proj = nn.Parameter(experts.down_proj.transpose(1, 2).contiguous())
+
+
 class Llama4TextMoe(nn.Module):
     experts: nn.Module
     router: nn.Linear
@@ -25,6 +33,8 @@ class Llama4TextMoe(nn.Module):
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1])
 
+        _fix_llama4_experts(hidden_states, self.experts)
+
         router_logits = self.router(hidden_states)
         out = fused_moe(
             hidden_states,