diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index f18cb119b..2f45c3a65 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -69,6 +69,7 @@ _DP: GroupCoordinator | None = None _VAE_DECODE: GroupCoordinator | None = None _DIT: ProcessGroup | None = None _VAE: ProcessGroup | None = None +_VAE_DECODE_PARALLEL_AXES = "tp-sp-pp-cfg" TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) @@ -201,6 +202,14 @@ def init_parallel_group_coordinator( ) +def _get_vae_decode_group_ranks( + rank_generator: RankGenerator, +) -> list[list[int]]: + # VAE decode happens after each DP replica owns a different request result. + # Decode can shard one request across TP/SP/PP/CFG ranks, but must not cross DP. + return rank_generator.get_ranks(_VAE_DECODE_PARALLEL_AXES) + + def get_tp_group() -> GroupCoordinator: assert _TP is not None, "tensor model parallel group is not initialized" return _TP @@ -461,7 +470,7 @@ def initialize_model_parallel( global _VAE_DECODE assert _VAE_DECODE is None, "VAE decode parallel group is already initialized" _VAE_DECODE = init_parallel_group_coordinator( - group_ranks=rank_generator.get_ranks("tp-sp-pp-cfg"), + group_ranks=_get_vae_decode_group_ranks(rank_generator), local_rank=get_world_group().local_rank, backend=backend, parallel_mode="vae_decode", diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_spatial_parallel_decode.py b/python/sglang/multimodal_gen/test/unit/test_vae_spatial_parallel_decode.py index bac9ff204..cffd746e8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_spatial_parallel_decode.py +++ b/python/sglang/multimodal_gen/test/unit/test_vae_spatial_parallel_decode.py @@ -255,7 +255,7 @@ class TestVAESpatialParallelDecode(unittest.TestCase): ) self.assertEqual( - rank_generator.get_ranks("tp-sp-pp-cfg"), + parallel_state._get_vae_decode_group_ranks(rank_generator), [list(range(0, 8)), list(range(8, 16))], )