try tp support etp

Prayer3th · Prayer3th · commit d5e817d51743 · 2025-11-12T20:12:52.000+08:00
diff --git a/python/sgl_jax/srt/layers/moe.py b/python/sgl_jax/srt/layers/moe.py
@@ -385,8 +385,7 @@ def _gmm_compute_with_sharded_weights(
         )
 
         # Need to reduce over the intermediate dimension which is sharded on "tensor" axis
-        mesh_shape_tensor = jax.lax.psum(1, "tensor")
-        if mesh_shape_tensor > 1:
+        if self.tp_size > 1:
             intermediate_output = jax.lax.psum(intermediate_output, "tensor")
 
         return intermediate_output

Original file line number	Diff line number	Diff line change
`@@ -385,8 +385,7 @@ def _gmm_compute_with_sharded_weights(`
`385`	`385`	`)`
`386`	`386`
`387`	`387`	`# Need to reduce over the intermediate dimension which is sharded on "tensor" axis`
`388`		`- mesh_shape_tensor = jax.lax.psum(1, "tensor")`
`389`		`- if mesh_shape_tensor > 1:`
	`388`	`+ if self.tp_size > 1:`
`390`	`389`	`intermediate_output = jax.lax.psum(intermediate_output, "tensor")`
`391`	`390`
`392`	`391`	`return intermediate_output`