[diffusion] chore: document VAE decode parallel group axes (#29364)

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",

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))],
         )