try to support etp

Prayer3th · Prayer3th · commit 0f61aa039351 · 2025-10-30T14:12:15.000+08:00
diff --git a/python/sgl_jax/srt/layers/moe.py b/python/sgl_jax/srt/layers/moe.py
@@ -443,7 +443,7 @@ def _expert_all_to_all_collect(self, data, global_group_sizes, expert_shard_id,
             send_sizes,
             output_offsets,
             recv_sizes,
-            axis_name=("data", "tensor"),
+            axis_name=("tensor"),
         )
 
         return result

Original file line number	Diff line number	Diff line change
`@@ -443,7 +443,7 @@ def _expert_all_to_all_collect(self, data, global_group_sizes, expert_shard_id,`
`443`	`443`	`send_sizes,`
`444`	`444`	`output_offsets,`
`445`	`445`	`recv_sizes,`
`446`		`- axis_name=("data", "tensor"),`
	`446`	`+ axis_name=("tensor"),`
`447`	`447`	`)`
`448`	`448`
`449`	`449`	`return result`