try to support etp

Author: Prayer3th

python/sgl_jax/srt/layers/moe.py

@@ -263,9 +263,9 @@ def __call__(self, hidden_states, topk_weights, topk_ids):
                 P(None),  # hidden_states
                 P(None),  # topk_weights
                 P(None),  # topk_ids
-                P("expert", None, "data"),  # w0_weights
-                P("expert", None, "data"),  # w1_weights
-                P("expert", "data", None),  # wo_weights
+                P("tensor", None, "data"),  # w0_weights
+                P("tensor", None, "data"),  # w1_weights
+                P("tensor", "data", None),  # wo_weights
             ),
             out_specs=P(None),
             check_rep=False,