chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/distributed/__init__.py

@@ -0,0 +1,63 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from sglang.multimodal_gen.runtime.distributed.communication_op import *
+from sglang.multimodal_gen.runtime.distributed.group_coordinator import (
+    get_local_torch_device,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    cleanup_dist_env_and_memory,
+    get_decode_parallel_group_coordinator,
+    get_decode_parallel_rank,
+    get_decode_parallel_world_size,
+    get_dp_group,
+    get_dp_rank,
+    get_dp_world_size,
+    get_sp_group,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+    get_tp_group,
+    get_tp_rank,
+    get_tp_world_size,
+    get_world_group,
+    get_world_rank,
+    get_world_size,
+    init_distributed_environment,
+    initialize_model_parallel,
+    maybe_init_distributed_environment_and_model_parallel,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.distributed.utils import *
+
+# SPDX-License-Identifier: Apache-2.0
+
+
+__all__ = [
+    # Initialization
+    "init_distributed_environment",
+    "initialize_model_parallel",
+    "cleanup_dist_env_and_memory",
+    "model_parallel_is_initialized",
+    "maybe_init_distributed_environment_and_model_parallel",
+    # World group
+    "get_world_group",
+    "get_world_rank",
+    "get_world_size",
+    # Data parallel group
+    "get_dp_group",
+    "get_dp_rank",
+    "get_dp_world_size",
+    # Sequence parallel group
+    "get_sp_group",
+    "get_sp_parallel_rank",
+    "get_sp_world_size",
+    # Tensor parallel group
+    "get_tp_group",
+    "get_tp_rank",
+    "get_tp_world_size",
+    # Decode parallel group
+    "get_decode_parallel_group_coordinator",
+    "get_decode_parallel_rank",
+    "get_decode_parallel_world_size",
+    # Get torch device
+    "get_local_torch_device",
+]

python/sglang/multimodal_gen/runtime/distributed/cfg_parallel_utils.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
+    _apply_cfg_postprocess,
+    _unwrap,
+    _wrap,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    cfg_model_parallel_all_gather,
+    cfg_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_classifier_free_guidance_world_size,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
+        CFGBranch,
+        CFGPolicy,
+    )
+
+# Tracks (n_branches, cfg_world_size, cfg_rank) tuples already logged so the
+# dispatch table is printed once per unique configuration, not once per step.
+_logged_dispatch_keys: set[tuple[int, int, int]] = set()
+
+
+def _run(
+    predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
+    bid: int,
+    branches,
+) -> tuple[torch.Tensor, ...]:
+    branch = branches[bid]
+    device = get_local_torch_device()
+    local_branch = dataclasses.replace(
+        branch,
+        kwargs={
+            k: v.to(device) if isinstance(v, torch.Tensor) else v
+            for k, v in branch.kwargs.items()
+        },
+    )
+    raw = predict_fn(local_branch)
+    return _wrap(raw)
+
+
+def run_cfg_parallel(
+    policy: CFGPolicy,
+    predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
+) -> list[torch.Tensor | tuple[torch.Tensor, ...]]:
+    """Dispatch CFG branches across ranks, all-gather results, return in branch order.
+
+    ``predict_fn`` is a closure capturing all step-varying state
+    (latent_model_input, timestep, model, etc.).  It is called with each
+    assigned ``CFGBranch`` and must return the raw ``_predict_noise`` output.
+
+    Idle ranks (cfg_world_size > n_branches) run branch 0 as a dummy forward
+    to obtain tensor shapes for the all-gather.
+
+    Returns a list indexed to match ``policy.branches``, identical on every rank.
+    """
+
+    cfg_rank = get_classifier_free_guidance_rank()
+    cfg_world_size = get_classifier_free_guidance_world_size()
+    branches = policy.branches
+    n_branches = len(branches)
+    assignments = dispatch_branches(n_branches, cfg_world_size)
+    branches_assigned_to_local_rank = assignments[cfg_rank]
+    max_num_branches_per_rank = max(len(a) for a in assignments)
+
+    if cfg_world_size > n_branches:
+        logger.warning_once(
+            "cfg_parallel_size=%d > n_branches=%d; %d GPU(s) will be idle for CFG",
+            cfg_world_size,
+            n_branches,
+            cfg_world_size - n_branches,
+        )
+
+    dispatch_key = (n_branches, cfg_world_size, cfg_rank)
+    if dispatch_key not in _logged_dispatch_keys:
+        _logged_dispatch_keys.add(dispatch_key)
+        branch_names = (
+            [branches[i].name for i in branches_assigned_to_local_rank]
+            if branches_assigned_to_local_rank
+            else ["(idle)"]
+        )
+        logger.info(
+            "CFG parallel dispatch: rank %d/%d -> [%s]",
+            cfg_rank,
+            cfg_world_size,
+            ", ".join(branch_names),
+        )
+
+    # perform the forward for local branches
+    predicts_from_local_branches: list[tuple[torch.Tensor, ...]] = [
+        _run(predict_fn, bid, branches) for bid in branches_assigned_to_local_rank
+    ]
+
+    if not predicts_from_local_branches:  # idle rank: run branch 0 for tensor shapes
+        predicts_from_local_branches.append(_run(predict_fn, 0, branches))
+
+    # pad the predicts to the length of max_num_branches_per_rank, to prepare for the all-gather later
+    ref = predicts_from_local_branches[0]
+    while len(predicts_from_local_branches) < max_num_branches_per_rank:
+        # TODO: cache this zero
+        predicts_from_local_branches.append(tuple(torch.zeros_like(t) for t in ref))
+
+    # All-gather each slot and output element with separate_tensors=True.
+    # all_slots[slot][elem] = list[Tensor] indexed by CFG rank; no reshape.
+    all_slots: list[list[list[torch.Tensor]]] = [
+        [
+            cfg_model_parallel_all_gather(p, dim=0, separate_tensors=True)
+            for p in slot_pred
+        ]
+        for slot_pred in predicts_from_local_branches
+    ]
+
+    # reorder the results in branch order: branch bid -> owner rank, slot.
+    n_elems = len(ref)
+    final: list[torch.Tensor | tuple[torch.Tensor, ...]] = []
+    for bid in range(n_branches):
+        owner = bid % cfg_world_size
+        slot = bid // cfg_world_size
+        elems = tuple(all_slots[slot][ei][owner] for ei in range(n_elems))
+        final.append(_unwrap(elems))
+    return final
+
+
+def run_two_branch_cfg_parallel(
+    policy: CFGPolicy,
+    predict_fn: Callable[[CFGBranch], torch.Tensor | tuple[torch.Tensor, ...]],
+    cfg_scale: float,
+    batch,
+    pipeline_config,
+) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    """Run standard two-pass CFG with the old all-reduce combine.
+
+    This keeps the existing WAN baselines: it avoids gathering both branch
+    predictions, and it preserves the bf16 arithmetic order used before the
+    multi-branch CFG dispatcher was added.
+    """
+
+    cfg_rank = get_classifier_free_guidance_rank()
+    pred_t = _run(predict_fn, cfg_rank, policy.branches)
+
+    if cfg_rank == 0:
+        partial = tuple(cfg_scale * p for p in pred_t)
+        cond_t = pred_t
+    else:
+        partial = tuple((1 - cfg_scale) * p for p in pred_t)
+        cond_t = tuple(torch.empty_like(p) for p in pred_t)
+
+    results = [cfg_model_parallel_all_reduce(p) for p in partial]
+    cond_t = tuple(get_cfg_group().broadcast(p, src=0) for p in cond_t)
+    results[0] = _apply_cfg_postprocess(results[0], cond_t[0], batch, pipeline_config)
+    return _unwrap(tuple(results))
+
+
+def dispatch_branches(n_branches: int, n_ranks: int) -> list[list[int]]:
+    """Assign branches to ranks in Round-robin fashion
+
+    Returns a list of length ``n_ranks`` where element ``r`` contains the
+    branch indices assigned to rank ``r``.  Branch ``i`` goes to rank
+    ``i % n_ranks``.
+
+    Example: 4 passes, 2 GPUs:
+        rank 0 -> [0, 2],  rank 1 -> [1, 3]
+    """
+    assignments: list[list[int]] = [[] for _ in range(n_ranks)]
+    for i in range(n_branches):
+        assignments[i % n_ranks].append(i)
+    return assignments

python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+
+
+@dataclass
+class CFGBranch:
+    """Immutable specification of one CFG branch forward pass.
+
+    Built once before the denoising loop; read-only across all steps.
+    """
+
+    name: str
+    is_conditional: bool
+    kwargs: dict[str, Any]
+
+    def configure_batch(self, batch: Req) -> None:
+        """Set batch state before this branch's forward pass.
+
+        Override for richer per-branch context (e.g. a branch index instead of
+        a single boolean) when a model needs more than two guidance modes.
+        """
+        batch.is_cfg_negative = not self.is_conditional
+
+
+@dataclass
+class CFGPolicy:
+    """Owns the CFG branches for one generation run and combines their predictions.
+
+    Built once before the denoising loop via ``build()``, then used read-only
+    across all steps.  Subclass and override ``build()`` / ``combine()`` for
+    custom CFG schemes (N-branch, multi-output, etc.).
+
+    The default implementation handles standard 2-branch CFG.  With a single
+    branch (CFG disabled) ``combine()`` returns the prediction unchanged.
+    """
+
+    branches: list[CFGBranch] = field(default_factory=list)
+
+    def build(
+        self,
+        batch: Req,
+        image_kwargs: dict[str, Any],
+        pos_cond_kwargs: dict[str, Any],
+        neg_cond_kwargs: dict[str, Any],
+    ) -> CFGPolicy:
+        """Return a new policy with branches populated.
+
+        Called once before the denoising loop.  The returned policy is
+        immutable for the lifetime of the run.  Override to declare N branches.
+        """
+        branches = [CFGBranch("conditional", True, {**image_kwargs, **pos_cond_kwargs})]
+        if batch.do_classifier_free_guidance:
+            branches.append(
+                CFGBranch("unconditional", False, {**image_kwargs, **neg_cond_kwargs})
+            )
+        return dataclasses.replace(self, branches=branches)
+
+    def combine(
+        self,
+        predictions: list[torch.Tensor | tuple[torch.Tensor, ...]],
+        batch: Req,
+        cfg_scale: float,
+        pipeline_config: Any,
+        *,
+        cfg_parallel: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """Combine branch predictions into the final noise estimate.
+
+        Default: standard 2-branch CFG formula applied element-wise, followed
+        by normalization / rescale / model-specific postprocess.
+        Single-branch (CFG disabled): returns the prediction unchanged.
+        Override for N-branch or multi-output models.
+        """
+        if len(predictions) == 1:
+            return predictions[0]
+        pos_t = _wrap(predictions[0])
+        neg_t = _wrap(predictions[1])
+        if cfg_parallel:
+            # Match the old CFG-parallel calculation: multiply the positive
+            # prediction by cfg_scale and the negative prediction by
+            # (1 - cfg_scale) before adding them. The serial CFG formula is
+            # mathematically equivalent, but bf16 rounding changes WAN outputs.
+            results = [
+                cfg_scale * p + (1 - cfg_scale) * n for p, n in zip(pos_t, neg_t)
+            ]
+        else:
+            results = [n + cfg_scale * (p - n) for p, n in zip(pos_t, neg_t)]
+        results[0] = _apply_cfg_postprocess(
+            results[0], pos_t[0], batch, pipeline_config
+        )
+        return _unwrap(tuple(results))
+
+
+# Helpers used by CFGPolicy and run_cfg_parallel.
+
+
+def _wrap(
+    pred: torch.Tensor | tuple[torch.Tensor, ...],
+) -> tuple[torch.Tensor, ...]:
+    return pred if isinstance(pred, tuple) else (pred,)
+
+
+def _unwrap(
+    pred: tuple[torch.Tensor, ...],
+) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    return pred[0] if len(pred) == 1 else pred
+
+
+def _apply_cfg_postprocess(
+    noise_pred: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    batch: Req,
+    pipeline_config: Any,
+) -> torch.Tensor:
+    if batch.cfg_normalization and float(batch.cfg_normalization) > 0:
+        noise_pred = _apply_cfg_normalization(
+            noise_pred, noise_pred_cond, float(batch.cfg_normalization)
+        )
+    if batch.guidance_rescale > 0.0:
+        noise_pred = _rescale_noise_cfg(
+            noise_pred, noise_pred_cond, guidance_rescale=batch.guidance_rescale
+        )
+    return pipeline_config.postprocess_cfg_noise(batch, noise_pred, noise_pred_cond)
+
+
+def _apply_cfg_normalization(
+    noise_pred: torch.Tensor,
+    noise_pred_cond: torch.Tensor,
+    cfg_normalization: float,
+) -> torch.Tensor:
+    cond_f = noise_pred_cond.float()
+    pred_f = noise_pred.float()
+    ori_norm = torch.linalg.vector_norm(cond_f)
+    new_norm = torch.linalg.vector_norm(pred_f)
+    max_norm = ori_norm * cfg_normalization
+    if new_norm > max_norm:
+        noise_pred = noise_pred * (max_norm / new_norm)
+    return noise_pred
+
+
+def _rescale_noise_cfg(
+    noise_cfg: torch.Tensor,
+    noise_pred_text: torch.Tensor,
+    guidance_rescale: float = 0.0,
+) -> torch.Tensor:
+    std_text = noise_pred_text.std(
+        dim=list(range(1, noise_pred_text.ndim)), keepdim=True
+    )
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    return guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg

python/sglang/multimodal_gen/runtime/distributed/communication_op.py

@@ -0,0 +1,67 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py
+
+import torch
+import torch.distributed as dist
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_cfg_group,
+    get_sp_group,
+    get_tp_group,
+)
+
+
+def tensor_model_parallel_all_reduce(
+    input_: torch.Tensor, tp_group: dist.ProcessGroup = None
+) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    tp_group = tp_group or get_tp_group()
+    return tp_group.all_reduce(input_)
+
+
+def tensor_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1, tp_group: dist.ProcessGroup = None
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    tp_group = tp_group or get_tp_group()
+    return tp_group.all_gather(input_, dim)
+
+
+# TODO: remove model, make it sequence_parallel
+def sequence_model_parallel_all_to_all_4D(
+    input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1
+) -> torch.Tensor:
+    """All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group."""
+    return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim)
+
+
+def sequence_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_sp_group().all_gather(input_, dim)
+
+
+def sequence_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
+    """All-reduce the input tensor across model parallel group."""
+    return get_sp_group().all_reduce(input_)
+
+
+def cfg_model_parallel_all_gather(
+    input_: torch.Tensor, dim: int = -1, separate_tensors: bool = False
+) -> torch.Tensor:
+    """All-gather the input tensor across model parallel group."""
+    return get_cfg_group().all_gather(input_, dim, separate_tensors)
+
+
+def cfg_model_parallel_all_reduce(
+    input_: torch.Tensor,
+    op: torch._C._distributed_c10d.ReduceOp = torch._C._distributed_c10d.ReduceOp.SUM,
+) -> torch.Tensor:
+    """All-reduce the input tensor across CFG parallel group."""
+    if not input_.is_contiguous():
+        input_ = input_.contiguous()
+    return get_cfg_group().all_reduce(input_, op=op)

python/sglang/multimodal_gen/runtime/distributed/device_communicators/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/runtime/distributed/device_communicators/base_device_communicator.py

@@ -0,0 +1,306 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py
+
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+from torch.distributed import ProcessGroup, ReduceOp
+
+
+class DistributedAutograd:
+    """Collection of autograd functions for distributed operations.
+
+    This class provides custom autograd functions for distributed operations like all_reduce,
+    all_gather, and all_to_all. Each operation is implemented as a static inner class with
+    proper forward and backward implementations.
+    """
+
+    class AllReduce(torch.autograd.Function):
+        """Differentiable all_reduce operation.
+
+        The gradient of all_reduce is another all_reduce operation since the operation
+        combines values from all ranks equally.
+        """
+
+        @staticmethod
+        def forward(
+            ctx: Any,
+            group: ProcessGroup,
+            input_: Tensor,
+            op: dist.ReduceOp | None = None,
+        ) -> Tensor:
+            ctx.group = group
+            ctx.op = op
+            output = input_.clone()
+            dist.all_reduce(output, group=group, op=op)
+            return output
+
+        @staticmethod
+        def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None]:
+            grad_output = grad_output.clone()
+            dist.all_reduce(grad_output, group=ctx.group, op=ctx.op)
+            return None, grad_output, None
+
+    class AllGather(torch.autograd.Function):
+        """Differentiable all_gather operation.
+
+        The operation gathers tensors from all ranks and concatenates them along a specified dimension.
+        The backward pass uses reduce_scatter to efficiently distribute gradients back to source ranks.
+        """
+
+        @staticmethod
+        def forward(
+            ctx: Any, group: ProcessGroup, input_: Tensor, world_size: int, dim: int
+        ) -> Tensor:
+            ctx.group = group
+            ctx.world_size = world_size
+            ctx.dim = dim
+            ctx.input_shape = input_.shape
+
+            input_size = input_.size()
+            output_size = (input_size[0] * world_size,) + input_size[1:]
+            output_tensor = torch.empty(
+                output_size, dtype=input_.dtype, device=input_.device
+            )
+
+            dist.all_gather_into_tensor(output_tensor, input_, group=group)
+
+            output_tensor = output_tensor.reshape((world_size,) + input_size)
+            output_tensor = output_tensor.movedim(0, dim)
+            output_tensor = output_tensor.reshape(
+                input_size[:dim]
+                + (world_size * input_size[dim],)
+                + input_size[dim + 1 :]
+            )
+            return output_tensor
+
+        @staticmethod
+        def backward(ctx: Any, grad_output: Tensor) -> tuple[None, Tensor, None, None]:
+            # Split the gradient tensor along the gathered dimension
+            dim_size = grad_output.size(ctx.dim) // ctx.world_size
+            grad_chunks = grad_output.reshape(
+                grad_output.shape[: ctx.dim]
+                + (ctx.world_size, dim_size)
+                + grad_output.shape[ctx.dim + 1 :]
+            )
+            grad_chunks = grad_chunks.movedim(ctx.dim, 0)
+
+            # Each rank only needs its corresponding gradient
+            grad_input = torch.empty(
+                ctx.input_shape, dtype=grad_output.dtype, device=grad_output.device
+            )
+            dist.reduce_scatter_tensor(
+                grad_input, grad_chunks.contiguous(), group=ctx.group
+            )
+
+            return None, grad_input, None, None
+
+    class AllToAll4D(torch.autograd.Function):
+        """Differentiable all_to_all operation specialized for 4D tensors.
+
+        This operation is particularly useful for attention operations where we need to
+        redistribute data across ranks for efficient parallel processing.
+
+        The operation supports two modes:
+        1. scatter_dim=2, gather_dim=1: Used for redistributing attention heads
+        2. scatter_dim=1, gather_dim=2: Used for redistributing sequence dimensions
+        """
+
+        @staticmethod
+        def forward(
+            ctx: Any,
+            group: ProcessGroup,
+            input_: Tensor,
+            world_size: int,
+            scatter_dim: int,
+            gather_dim: int,
+        ) -> Tensor:
+            ctx.group = group
+            ctx.world_size = world_size
+            ctx.scatter_dim = scatter_dim
+            ctx.gather_dim = gather_dim
+
+            if world_size == 1:
+                return input_
+
+            assert (
+                input_.dim() == 4
+            ), f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
+
+            if scatter_dim == 2 and gather_dim == 1:
+                bs, shard_seqlen, hn, hd = input_.shape
+                assert hn % world_size == 0, (
+                    f"head dimension ({hn}) must be divisible by sequence "
+                    f"parallel world size ({world_size})"
+                )
+                seqlen = shard_seqlen * world_size
+                shard_hn = hn // world_size
+
+                input_ = input_.transpose(0, 2).contiguous()  # hn, shard_seqlen, bs, hd
+                output = torch.empty_like(input_)
+
+                dist.all_to_all_single(
+                    output, input_, group=group
+                )  # hn, shard_seqlen, bs, hd
+
+                output = torch.cat(
+                    output.split(shard_hn), dim=1
+                )  # sharded hn, seqlen, bs, hd
+
+                output = output.transpose(
+                    0, 2
+                ).contiguous()  # bs, seqlen, sharded_hn, hd
+
+                return output
+            elif scatter_dim == 1 and gather_dim == 2:
+                bs, seqlen, shard_hn, hd = input_.shape
+                assert seqlen % world_size == 0, (
+                    f"sequence dimension ({seqlen}) must be divisible by sequence "
+                    f"parallel world size ({world_size})"
+                )
+                hn = shard_hn * world_size
+                shard_seqlen = seqlen // world_size
+
+                input_ = input_.transpose(0, 2).contiguous()  # shard_hn, seqlen, bs, hd
+
+                input_ = (
+                    input_.reshape(shard_hn, world_size, shard_seqlen, bs, hd)
+                    .transpose(0, 1)
+                    .reshape(shard_hn * world_size, shard_seqlen, bs, hd)
+                    .contiguous()
+                )
+
+                output = torch.empty_like(input_)
+
+                dist.all_to_all_single(output, input_, group=group)
+
+                output = output.transpose(
+                    0, 2
+                ).contiguous()  # bs, seqlen, sharded_hn, hd
+
+                return output
+            else:
+                raise RuntimeError(
+                    f"Invalid scatter_dim={scatter_dim}, gather_dim={gather_dim}. "
+                    f"Only (scatter_dim=2, gather_dim=1) and (scatter_dim=1, gather_dim=2) are supported."
+                )
+
+        @staticmethod
+        def backward(
+            ctx: Any, grad_output: Tensor
+        ) -> tuple[None, Tensor, None, None, None]:
+            if ctx.world_size == 1:
+                return None, grad_output, None, None, None
+
+            # For backward pass, we swap scatter_dim and gather_dim
+            output = DistributedAutograd.AllToAll4D.apply(
+                ctx.group, grad_output, ctx.world_size, ctx.gather_dim, ctx.scatter_dim
+            )
+            return None, output, None, None, None
+
+
+class DeviceCommunicatorBase:
+    """
+    Base class for device-specific communicator with autograd support.
+    It can use the `cpu_group` to initialize the communicator.
+    If the device has PyTorch integration (PyTorch can recognize its
+    communication backend), the `device_group` will also be given.
+    """
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        self.device = device or torch.device("cpu")
+        self.cpu_group = cpu_group
+        self.device_group = device_group
+        self.unique_name = unique_name
+        self.rank = dist.get_rank(cpu_group)
+        self.world_size = dist.get_world_size(cpu_group)
+        self.ranks = dist.get_process_group_ranks(cpu_group)
+        self.global_rank = dist.get_rank()
+        self.global_world_size = dist.get_world_size()
+        self.rank_in_group = dist.get_group_rank(self.cpu_group, self.global_rank)
+
+    def all_reduce(
+        self, input_: torch.Tensor, op: dist.ReduceOp | None = ReduceOp.SUM
+    ) -> torch.Tensor:
+        """Performs an all_reduce operation with gradient support."""
+        return DistributedAutograd.AllReduce.apply(self.device_group, input_, op)
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        """Performs an all_gather operation with gradient support."""
+        if dim < 0:
+            dim += input_.dim()
+        return DistributedAutograd.AllGather.apply(
+            self.device_group, input_, self.world_size, dim
+        )
+
+    def all_to_all_4D(
+        self, input_: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1
+    ) -> torch.Tensor:
+        """Performs a 4D all-to-all operation with gradient support."""
+        return DistributedAutograd.AllToAll4D.apply(
+            self.device_group, input_, self.world_size, scatter_dim, gather_dim
+        )
+
+    def gather(
+        self, input_: torch.Tensor, dst: int = 0, dim: int = -1
+    ) -> torch.Tensor | None:
+        """
+        NOTE: We assume that the input tensor is on the same device across
+        all the ranks.
+        NOTE: `dst` is the local rank of the destination rank.
+        """
+        world_size = self.world_size
+        assert (
+            -input_.dim() <= dim < input_.dim()
+        ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Allocate output tensor.
+        if self.rank_in_group == dst:
+            gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+        else:
+            gather_list = None
+        # Gather.
+        torch.distributed.gather(
+            input_, gather_list, dst=self.ranks[dst], group=self.device_group
+        )
+        if self.rank_in_group == dst:
+            output_tensor = torch.cat(gather_list, dim=dim)
+        else:
+            output_tensor = None
+        return output_tensor
+
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        """Sends a tensor to the destination rank in a non-blocking way"""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+        torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(
+        self, size: torch.Size, dtype: torch.dtype, src: int | None = None
+    ) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def destroy(self) -> None:
+        pass

python/sglang/multimodal_gen/runtime/distributed/device_communicators/cpu_communicator.py

@@ -0,0 +1,162 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from: https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/cpu_communicator.py
+
+import os
+
+import torch
+from torch.distributed import ProcessGroup
+
+from .base_device_communicator import DeviceCommunicatorBase
+
+
+class CpuCommunicator(DeviceCommunicatorBase):
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        from sglang.multimodal_gen.runtime.platforms import current_platform
+        from sglang.multimodal_gen.runtime.platforms.interface import CpuArchEnum
+
+        super().__init__(cpu_group, device, device_group, unique_name)
+        self.dist_module = torch.distributed
+
+        if (
+            (current_platform.get_cpu_architecture() == CpuArchEnum.X86)
+            and hasattr(torch.ops._C, "init_shm_manager")
+            and unique_name.startswith("tp")
+        ):
+            self.dist_module = _CPUSHMDistributed(self)
+
+    def all_reduce(
+        self,
+        input_: torch.Tensor,
+        op: torch.distributed.ReduceOp | None = torch.distributed.ReduceOp.SUM,
+    ) -> torch.Tensor:
+        self.dist_module.all_reduce(input_, group=self.device_group, op=op)
+        return input_
+
+    def gather(
+        self, input_: torch.Tensor, dst: int = 0, dim: int = -1
+    ) -> torch.Tensor | None:
+        """
+        NOTE: We assume that the input tensor is on the same device across
+        all the ranks.
+        NOTE: `dst` is the local rank of the destination rank.
+        """
+        world_size = self.world_size
+        assert (
+            -input_.dim() <= dim < input_.dim()
+        ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+
+        # Allocate output tensor.
+        if self.rank_in_group == dst:
+            gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+        else:
+            gather_list = None
+
+        # Gather.
+        self.dist_module.gather(
+            input_, gather_list, dst=self.ranks[dst], group=self.device_group
+        )
+
+        if self.rank_in_group == dst:
+            output_tensor = torch.cat(gather_list, dim=dim)
+        else:
+            output_tensor = None
+        return output_tensor
+
+    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        if dim < 0:
+            # Convert negative dim to positive.
+            dim += input_.dim()
+        input_size = input_.size()
+        # NOTE: we have to use concat-style all-gather here,
+        # stack-style all-gather has compatibility issues with
+        # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
+        output_size = (input_size[0] * self.world_size,) + input_size[1:]
+        # Allocate output tensor.
+        output_tensor = torch.empty(
+            output_size, dtype=input_.dtype, device=input_.device
+        )
+        # All-gather.
+        self.dist_module.all_gather_into_tensor(
+            output_tensor, input_, group=self.device_group
+        )
+
+        # Reshape
+        output_tensor = output_tensor.reshape((self.world_size,) + input_size)
+        output_tensor = output_tensor.movedim(0, dim)
+        output_tensor = output_tensor.reshape(
+            input_size[:dim]
+            + (self.world_size * input_size[dim],)
+            + input_size[dim + 1 :]
+        )
+        return output_tensor
+
+
+class _CPUSHMDistributed:
+
+    def __init__(self, communicator: CpuCommunicator):
+        instance_identifier = os.environ["VLLM_DIST_IDENT"]
+        unique_name = communicator.unique_name
+        instance_identifier = f"{instance_identifier}-{unique_name}"
+        self.communicator = communicator
+
+        group_ranks = [str(rank) for rank in self.communicator.ranks]
+        shm_group_identifier = f"[{'-'.join(group_ranks)}]"
+        self.group_name = f"{instance_identifier}-{shm_group_identifier}-cpushm"
+
+        self.handle = self._init_cpu_shm()
+
+    def _init_cpu_shm(self) -> int:
+        handle = torch.ops._C.init_shm_manager(
+            self.group_name,
+            self.communicator.world_size,
+            self.communicator.rank,
+        )
+        torch.distributed.barrier(self.communicator.device_group)
+        torch.ops._C.join_shm_manager(
+            handle,
+            self.group_name,
+        )
+        torch.distributed.barrier(self.communicator.device_group)
+
+        return int(handle)
+
+    def all_reduce(
+        self, input: torch.Tensor, group: ProcessGroup | None = None
+    ) -> None:
+        torch.ops._C.shm_allreduce(self.handle, input)
+
+    def gather(
+        self,
+        input: torch.Tensor,
+        gather_list: list[torch.Tensor] | None,
+        dst: int = -1,
+        group: ProcessGroup | None = None,
+    ) -> None:
+        # Note: different from the torch gather, here we use local dst rank.
+        torch.ops._C.shm_gather(
+            self.handle,
+            input,
+            gather_list,
+            torch.distributed.get_group_rank(group, dst),
+        )
+
+    def all_gather_into_tensor(
+        self,
+        output: torch.Tensor,
+        input: torch.Tensor,
+        group: ProcessGroup | None = None,
+    ) -> None:
+        torch.ops._C.shm_all_gather(self.handle, input, output)

python/sglang/multimodal_gen/runtime/distributed/device_communicators/cuda_communicator.py

@@ -0,0 +1,80 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py
+
+import torch
+from torch.distributed import ProcessGroup
+
+from sglang.multimodal_gen.runtime.distributed.device_communicators.base_device_communicator import (
+    DeviceCommunicatorBase,
+)
+
+
+class CudaCommunicator(DeviceCommunicatorBase):
+
+    def __init__(
+        self,
+        cpu_group: ProcessGroup,
+        device: torch.device | None = None,
+        device_group: ProcessGroup | None = None,
+        unique_name: str = "",
+    ):
+        super().__init__(cpu_group, device, device_group, unique_name)
+
+        from sglang.multimodal_gen.runtime.distributed.device_communicators.pynccl import (
+            PyNcclCommunicator,
+        )
+
+        self.pynccl_comm: PyNcclCommunicator | None = None
+        if self.world_size > 1:
+            self.pynccl_comm = PyNcclCommunicator(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+    def all_reduce(self, input_, op: torch.distributed.ReduceOp | None = None):
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
+        out = pynccl_comm.all_reduce(input_, op=op)
+        if out is None:
+            # fall back to the default all-reduce using PyTorch.
+            # this usually happens during testing.
+            # when we run the model, allreduce only happens for the TP
+            # group, where we always have either custom allreduce or pynccl.
+            out = input_.clone()
+            torch.distributed.all_reduce(out, group=self.device_group, op=op)
+        return out
+
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        """Sends a tensor to the destination rank in a non-blocking way"""
+        """NOTE: `dst` is the local rank of the destination rank."""
+        if dst is None:
+            dst = (self.rank_in_group + 1) % self.world_size
+
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.send(tensor, dst)
+        else:
+            torch.distributed.send(tensor, self.ranks[dst], self.device_group)
+
+    def recv(
+        self, size: torch.Size, dtype: torch.dtype, src: int | None = None
+    ) -> torch.Tensor:
+        """Receives a tensor from the source rank."""
+        """NOTE: `src` is the local rank of the source rank."""
+        if src is None:
+            src = (self.rank_in_group - 1) % self.world_size
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        pynccl_comm = self.pynccl_comm
+        if pynccl_comm is not None and not pynccl_comm.disabled:
+            pynccl_comm.recv(tensor, src)
+        else:
+            torch.distributed.recv(tensor, self.ranks[src], self.device_group)
+        return tensor
+
+    def destroy(self) -> None:
+        if self.pynccl_comm is not None:
+            self.pynccl_comm = None

python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl.py

@@ -0,0 +1,259 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py
+
+# ===================== import region =====================
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup, ReduceOp
+
+from sglang.multimodal_gen.runtime.distributed.device_communicators.pynccl_wrapper import (
+    NCCLLibrary,
+    buffer_type,
+    cudaStream_t,
+    ncclComm_t,
+    ncclDataTypeEnum,
+    ncclRedOpTypeEnum,
+    ncclUniqueId,
+)
+from sglang.multimodal_gen.runtime.distributed.utils import StatelessProcessGroup
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import current_stream
+
+logger = init_logger(__name__)
+
+
+class PyNcclCommunicator:
+
+    def __init__(
+        self,
+        group: ProcessGroup | StatelessProcessGroup,
+        device: int | str | torch.device,
+        library_path: str | None = None,
+    ):
+        """
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the PyNcclCommunicator to. If None,
+                it will be bind to f"cuda:{local_rank}".
+            library_path: the path to the NCCL library. If None, it will
+                use the default library path.
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device.
+        """
+        if not isinstance(group, StatelessProcessGroup):
+            assert dist.is_initialized()
+            assert (
+                dist.get_backend(group) != dist.Backend.NCCL
+            ), "PyNcclCommunicator should be attached to a non-NCCL group."
+            # note: this rank is the rank in the group
+            self.rank = dist.get_rank(group)
+            self.world_size = dist.get_world_size(group)
+        else:
+            self.rank = group.rank
+            self.world_size = group.world_size
+
+        self.group = group
+
+        # if world_size == 1, no need to create communicator
+        if self.world_size == 1:
+            self.available = False
+            self.disabled = True
+            return
+        try:
+            self.nccl = NCCLLibrary(library_path)
+        except Exception:
+            # disable because of missing NCCL library
+            # e.g. in a non-GPU environment
+            self.available = False
+            self.disabled = True
+            return
+
+        self.available = True
+        self.disabled = False
+
+        logger.info("sglang-diffusion is using nccl==%s", self.nccl.ncclGetVersion())
+
+        if self.rank == 0:
+            # get the unique id from NCCL
+            self.unique_id = self.nccl.ncclGetUniqueId()
+        else:
+            # construct an empty unique id
+            self.unique_id = ncclUniqueId()
+
+        if not isinstance(group, StatelessProcessGroup):
+            tensor = torch.ByteTensor(list(self.unique_id.internal))
+            ranks = dist.get_process_group_ranks(group)
+            # arg `src` in `broadcast` is the global rank
+            dist.broadcast(tensor, src=ranks[0], group=group)
+            byte_list = tensor.tolist()
+            for i, byte in enumerate(byte_list):
+                self.unique_id.internal[i] = byte
+        else:
+            self.unique_id = group.broadcast_obj(self.unique_id, src=0)
+        if isinstance(device, int):
+            device = torch.device(f"cuda:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+        # nccl communicator and stream will use this device
+        # `torch.cuda.device` is a context manager that changes the
+        # current cuda device to the specified one
+        with torch.cuda.device(device):
+            self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
+                self.world_size, self.unique_id, self.rank
+            )
+
+            stream = current_stream()
+            # A small all_reduce for warmup.
+            data = torch.zeros(1, device=device)
+            self.all_reduce(data)
+            if stream is not None:
+                stream.synchronize()
+            del data
+
+    def all_reduce(
+        self, in_tensor: torch.Tensor, op: ReduceOp = ReduceOp.SUM, stream=None
+    ) -> torch.Tensor:
+        if self.disabled:
+            return None
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert in_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {in_tensor.device}"
+        )
+
+        out_tensor = torch.empty_like(in_tensor)
+
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclAllReduce(
+            buffer_type(in_tensor.data_ptr()),
+            buffer_type(out_tensor.data_ptr()),
+            in_tensor.numel(),
+            ncclDataTypeEnum.from_torch(in_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+        return out_tensor
+
+    def all_gather(
+        self, output_tensor: torch.Tensor, input_tensor: torch.Tensor, stream=None
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclAllGather(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()),
+            input_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def reduce_scatter(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        op: ReduceOp = ReduceOp.SUM,
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclReduceScatter(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()),
+            output_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def send(self, tensor: torch.Tensor, dst: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclSend(
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            dst,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def recv(self, tensor: torch.Tensor, src: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclRecv(
+            buffer_type(tensor.data_ptr()),
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            src,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
+        if self.disabled:
+            return
+        assert tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        if src == self.rank:
+            sendbuff = buffer_type(tensor.data_ptr())
+            # NCCL requires the sender also to have a receive buffer
+            recvbuff = buffer_type(tensor.data_ptr())
+        else:
+            sendbuff = buffer_type()
+            recvbuff = buffer_type(tensor.data_ptr())
+        self.nccl.ncclBroadcast(
+            sendbuff,
+            recvbuff,
+            tensor.numel(),
+            ncclDataTypeEnum.from_torch(tensor.dtype),
+            src,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )

python/sglang/multimodal_gen/runtime/distributed/device_communicators/pynccl_wrapper.py

@@ -0,0 +1,451 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py
+
+# This file is a pure Python wrapper for the NCCL library.
+# The main purpose is to use NCCL combined with CUDA graph.
+# Before writing this script, we tried the following approach:
+# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself
+#  often gets stuck when initializing the NCCL communicator.
+# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce`
+#  contains many other potential cuda APIs, that are not allowed during
+#  capturing the CUDA graph. For further details, please check
+# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ .
+#
+# Another rejected idea is to write a C/C++ binding for NCCL. It is usually
+# doable, but we often encounter issues related with nccl versions, and need
+# to switch between different versions of NCCL. See
+# https://github.com/NVIDIA/nccl/issues/1234 for more details.
+# A C/C++ binding is not flexible enough to handle this. It requires
+# recompilation of the code every time we want to switch between different
+# versions. This current implementation, with a **pure** Python wrapper, is
+# more flexible. We can easily switch between different versions of NCCL by
+# changing the environment variable `SGLANG_DIFFUSION_NCCL_SO_PATH`, or the `so_file`
+# variable in the code.
+
+# TODO(will): support SGLANG_DIFFUSION_NCCL_SO_PATH
+
+import ctypes
+import platform
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch.distributed import ReduceOp
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import find_nccl_library
+
+logger = init_logger(__name__)
+
+# === export types and functions from nccl to Python ===
+# for the original nccl definition, please check
+# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
+
+ncclResult_t = ctypes.c_int
+ncclComm_t = ctypes.c_void_p
+
+
+class ncclUniqueId(ctypes.Structure):
+    _fields_ = [("internal", ctypes.c_byte * 128)]
+
+
+cudaStream_t = ctypes.c_void_p
+buffer_type = ctypes.c_void_p
+
+ncclDataType_t = ctypes.c_int
+
+
+class ncclDataTypeEnum:
+    ncclInt8 = 0
+    ncclChar = 0
+    ncclUint8 = 1
+    ncclInt32 = 2
+    ncclInt = 2
+    ncclUint32 = 3
+    ncclInt64 = 4
+    ncclUint64 = 5
+    ncclFloat16 = 6
+    ncclHalf = 6
+    ncclFloat32 = 7
+    ncclFloat = 7
+    ncclFloat64 = 8
+    ncclDouble = 8
+    ncclBfloat16 = 9
+    ncclNumTypes = 10
+
+    @classmethod
+    def from_torch(cls, dtype: torch.dtype) -> int:
+        if dtype == torch.int8:
+            return cls.ncclInt8
+        if dtype == torch.uint8:
+            return cls.ncclUint8
+        if dtype == torch.int32:
+            return cls.ncclInt32
+        if dtype == torch.int64:
+            return cls.ncclInt64
+        if dtype == torch.float16:
+            return cls.ncclFloat16
+        if dtype == torch.float32:
+            return cls.ncclFloat32
+        if dtype == torch.float64:
+            return cls.ncclFloat64
+        if dtype == torch.bfloat16:
+            return cls.ncclBfloat16
+        raise ValueError(f"Unsupported dtype: {dtype}")
+
+
+ncclRedOp_t = ctypes.c_int
+
+
+class ncclRedOpTypeEnum:
+    ncclSum = 0
+    ncclProd = 1
+    ncclMax = 2
+    ncclMin = 3
+    ncclAvg = 4
+    ncclNumOps = 5
+
+    @classmethod
+    def from_torch(cls, op: ReduceOp) -> int:
+        if op == ReduceOp.SUM:
+            return cls.ncclSum
+        if op == ReduceOp.PRODUCT:
+            return cls.ncclProd
+        if op == ReduceOp.MAX:
+            return cls.ncclMax
+        if op == ReduceOp.MIN:
+            return cls.ncclMin
+        if op == ReduceOp.AVG:
+            return cls.ncclAvg
+        raise ValueError(f"Unsupported op: {op}")
+
+
+@dataclass
+class Function:
+    name: str
+    restype: Any
+    argtypes: list[Any]
+
+
+class NCCLLibrary:
+    exported_functions = [
+        # const char* ncclGetErrorString(ncclResult_t result)
+        Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
+        # ncclResult_t  ncclGetVersion(int *version);
+        Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]),
+        # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
+        Function("ncclGetUniqueId", ncclResult_t, [ctypes.POINTER(ncclUniqueId)]),
+        # ncclResult_t  ncclCommInitRank(
+        #   ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
+        # note that ncclComm_t is a pointer type, so the first argument
+        # is a pointer to a pointer
+        Function(
+            "ncclCommInitRank",
+            ncclResult_t,
+            [ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int],
+        ),
+        # ncclResult_t  ncclAllReduce(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclAllReduce",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclAllGather(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclAllGather",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclReduceScatter(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
+        #   cudaStream_t stream);
+        # note that cudaStream_t is a pointer type, so the last argument
+        # is a pointer
+        Function(
+            "ncclReduceScatter",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ncclRedOp_t,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclSend(
+        #   const void* sendbuff, size_t count, ncclDataType_t datatype,
+        #   int dest, ncclComm_t comm, cudaStream_t stream);
+        Function(
+            "ncclSend",
+            ncclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t  ncclRecv(
+        #   void* recvbuff, size_t count, ncclDataType_t datatype,
+        #   int src, ncclComm_t comm, cudaStream_t stream);
+        Function(
+            "ncclRecv",
+            ncclResult_t,
+            [
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # ncclResult_t ncclBroadcast(
+        #   const void* sendbuff, void* recvbuff, size_t count,
+        #   ncclDataType_t datatype, int root, ncclComm_t comm,
+        #   cudaStream_t stream);
+        Function(
+            "ncclBroadcast",
+            ncclResult_t,
+            [
+                buffer_type,
+                buffer_type,
+                ctypes.c_size_t,
+                ncclDataType_t,
+                ctypes.c_int,
+                ncclComm_t,
+                cudaStream_t,
+            ],
+        ),
+        # be cautious! this is a collective call, it will block until all
+        # processes in the communicator have called this function.
+        # because Python object destruction can happen in random order,
+        # it is better not to call it at all.
+        # ncclResult_t  ncclCommDestroy(ncclComm_t comm);
+        Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
+    ]
+
+    # class attribute to store the mapping from the path to the library
+    # to avoid loading the same library multiple times
+    path_to_library_cache: dict[str, Any] = {}
+
+    # class attribute to store the mapping from library path
+    #  to the corresponding dictionary
+    path_to_dict_mapping: dict[str, dict[str, Any]] = {}
+
+    def __init__(self, so_file: str | None = None):
+
+        so_file = so_file or find_nccl_library()
+
+        try:
+            if so_file not in NCCLLibrary.path_to_dict_mapping:
+                lib = ctypes.CDLL(so_file)
+                NCCLLibrary.path_to_library_cache[so_file] = lib
+            self.lib = NCCLLibrary.path_to_library_cache[so_file]
+        except Exception as e:
+            logger.error(
+                "Failed to load NCCL library from %s ."
+                "It is expected if you are not running on NVIDIA/AMD/MTHREADS GPUs."
+                "Otherwise, the nccl library might not exist, be corrupted "
+                "or it does not support the current platform %s."
+                "If you already have the library, please set the "
+                "environment variable SGLANG_DIFFUSION_NCCL_SO_PATH"
+                " to point to the correct nccl library path.",
+                so_file,
+                platform.platform(),
+            )
+            raise e
+
+        if so_file not in NCCLLibrary.path_to_dict_mapping:
+            _funcs: dict[str, Any] = {}
+            for func in NCCLLibrary.exported_functions:
+                f = getattr(self.lib, func.name)
+                f.restype = func.restype
+                f.argtypes = func.argtypes
+                _funcs[func.name] = f
+            NCCLLibrary.path_to_dict_mapping[so_file] = _funcs
+        self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
+
+    def ncclGetErrorString(self, result: ncclResult_t) -> str:
+        return str(self._funcs["ncclGetErrorString"](result).decode("utf-8"))
+
+    def NCCL_CHECK(self, result: ncclResult_t) -> None:
+        if result != 0:
+            error_str = self.ncclGetErrorString(result)
+            raise RuntimeError(f"NCCL error: {error_str}")
+
+    def ncclGetVersion(self) -> str:
+        version = ctypes.c_int()
+        self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
+        version_str = str(version.value)
+        # something like 21903 --> "2.19.3"
+        major = version_str[0].lstrip("0")
+        minor = version_str[1:3].lstrip("0")
+        patch = version_str[3:].lstrip("0")
+        return f"{major}.{minor}.{patch}"
+
+    def ncclGetUniqueId(self) -> ncclUniqueId:
+        unique_id = ncclUniqueId()
+        self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id)))
+        return unique_id
+
+    def ncclCommInitRank(
+        self, world_size: int, unique_id: ncclUniqueId, rank: int
+    ) -> ncclComm_t:
+        comm = ncclComm_t()
+        self.NCCL_CHECK(
+            self._funcs["ncclCommInitRank"](
+                ctypes.byref(comm), world_size, unique_id, rank
+            )
+        )
+        return comm
+
+    def ncclAllReduce(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclAllReduce"](
+                sendbuff, recvbuff, count, datatype, op, comm, stream
+            )
+        )
+
+    def ncclReduceScatter(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        op: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # and `op` should be `ncclRedOp_t`
+        # both are aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclReduceScatter"](
+                sendbuff, recvbuff, count, datatype, op, comm, stream
+            )
+        )
+
+    def ncclAllGather(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        # `datatype` actually should be `ncclDataType_t`
+        # which is an aliases of `ctypes.c_int`
+        # when we pass int to a function, it will be converted to `ctypes.c_int`
+        # by ctypes automatically
+        self.NCCL_CHECK(
+            self._funcs["ncclAllGather"](
+                sendbuff, recvbuff, count, datatype, comm, stream
+            )
+        )
+
+    def ncclSend(
+        self,
+        sendbuff: buffer_type,
+        count: int,
+        datatype: int,
+        dest: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclSend"](sendbuff, count, datatype, dest, comm, stream)
+        )
+
+    def ncclRecv(
+        self,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        src: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclRecv"](recvbuff, count, datatype, src, comm, stream)
+        )
+
+    def ncclBroadcast(
+        self,
+        sendbuff: buffer_type,
+        recvbuff: buffer_type,
+        count: int,
+        datatype: int,
+        root: int,
+        comm: ncclComm_t,
+        stream: cudaStream_t,
+    ) -> None:
+        self.NCCL_CHECK(
+            self._funcs["ncclBroadcast"](
+                sendbuff, recvbuff, count, datatype, root, comm, stream
+            )
+        )
+
+    def ncclCommDestroy(self, comm: ncclComm_t) -> None:
+        self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
+
+
+__all__ = [
+    "NCCLLibrary",
+    "ncclDataTypeEnum",
+    "ncclRedOpTypeEnum",
+    "ncclUniqueId",
+    "ncclComm_t",
+    "cudaStream_t",
+    "buffer_type",
+]

python/sglang/multimodal_gen/runtime/distributed/parallel_groups.py

@@ -0,0 +1,91 @@
+# Reference: https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
+
+
+import torch
+
+
+class Singleton:
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super(Singleton, cls).__new__(cls, *args, **kwargs)
+        return cls._instance
+
+
+class ProcessGroupSingleton(Singleton):
+    def __init__(self):
+        self.ULYSSES_PG = None
+        self.RING_PG = None
+
+
+PROCESS_GROUP = ProcessGroupSingleton()
+
+
+def set_seq_parallel_pg_by_sp_groups(
+    sp_ulysses_degree,
+    sp_ring_degree,
+    rank: int,
+    sp_groups: list[list[int]],
+    use_ulysses_low: bool = True,
+):
+    """Create Ulysses/Ring process groups inside each SP group.
+
+    This is required when TP>1, because SP groups are not necessarily made of
+    consecutive global ranks (e.g., tp-sp order makes SP ranks strided).
+
+    Args:
+        sp_ulysses_degree: ulysses degree inside SP.
+        sp_ring_degree: ring degree inside SP.
+        rank: global rank of current process.
+        sp_groups: list of global-rank lists for each SP group.
+        use_ulysses_low: keep the same semantics as the original function.
+    """
+    sp_degree = sp_ring_degree * sp_ulysses_degree
+    assert sp_degree > 0
+    assert all(
+        len(g) == sp_degree for g in sp_groups
+    ), f"Each SP group must have size {sp_degree}, got sizes {[len(g) for g in sp_groups]}"
+
+    ulyssess_pg = None
+    ring_pg = None
+
+    num_ulysses_pgs = sp_ring_degree
+    num_ring_pgs = sp_ulysses_degree
+
+    def _map_indices_to_ranks(ranks: list[int], indices: list[int]) -> list[int]:
+        return [ranks[i] for i in indices]
+
+    # Important: call torch.distributed.new_group in the same order on all ranks.
+    for sp_ranks in sp_groups:
+        if use_ulysses_low:
+            for i in range(num_ulysses_pgs):
+                idx = list(range(i * sp_ulysses_degree, (i + 1) * sp_ulysses_degree))
+                ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx)
+                group = torch.distributed.new_group(ulysses_ranks)
+                if rank in ulysses_ranks:
+                    ulyssess_pg = group
+
+            for i in range(num_ring_pgs):
+                idx = list(range(i, sp_degree, num_ring_pgs))
+                ring_ranks = _map_indices_to_ranks(sp_ranks, idx)
+                group = torch.distributed.new_group(ring_ranks)
+                if rank in ring_ranks:
+                    ring_pg = group
+        else:
+            for i in range(num_ring_pgs):
+                idx = list(range(i * sp_ring_degree, (i + 1) * sp_ring_degree))
+                ring_ranks = _map_indices_to_ranks(sp_ranks, idx)
+                group = torch.distributed.new_group(ring_ranks)
+                if rank in ring_ranks:
+                    ring_pg = group
+
+            for i in range(num_ulysses_pgs):
+                idx = list(range(i, sp_degree, num_ulysses_pgs))
+                ulysses_ranks = _map_indices_to_ranks(sp_ranks, idx)
+                group = torch.distributed.new_group(ulysses_ranks)
+                if rank in ulysses_ranks:
+                    ulyssess_pg = group
+
+    PROCESS_GROUP.ULYSSES_PG = ulyssess_pg
+    PROCESS_GROUP.RING_PG = ring_pg

python/sglang/multimodal_gen/runtime/distributed/parallel_state.py

@@ -0,0 +1,912 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py
+# Copyright 2023 The vLLM team.
+# Adapted from
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+# Adapted from
+# Copyright 2024 xDiT team.
+# Adapted from
+# https://github.com/vllm-project/vllm/blob/main/vllm/distributed/parallel_state.py
+# Copyright 2023 The vLLM team.
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""sglang-diffusion distributed state.
+
+It takes over the control of the distributed environment from PyTorch.
+The typical workflow is:
+
+- call `init_distributed_environment` to initialize the distributed environment.
+- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to
+ initialize the model parallel groups.
+
+- any code dealing with the distributed stuff
+
+- call `destroy_model_parallel` to destroy the model parallel groups.
+- call `destroy_distributed_environment` to destroy the distributed environment.
+
+If you only need to use the distributed environment without model parallelism,
+ you can skip the model parallel initialization and destruction steps.
+"""
+
+import contextlib
+import datetime
+import os
+import weakref
+from collections import namedtuple
+from collections.abc import Callable
+from contextlib import contextmanager
+from multiprocessing import shared_memory
+from typing import Any, List, Optional
+from unittest.mock import patch
+
+import torch
+import torch.distributed
+from torch.distributed import ProcessGroup
+
+import sglang.multimodal_gen.envs as envs
+from sglang.multimodal_gen.runtime.distributed.utils import StatelessProcessGroup
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+from ..utils.distributed import RankGenerator
+from .group_coordinator import (
+    GroupCoordinator,
+    PipelineGroupCoordinator,
+    SequenceParallelGroupCoordinator,
+    get_local_torch_device,
+)
+
+logger = init_logger(__name__)
+
+_WORLD: GroupCoordinator | None = None
+_TP: GroupCoordinator | None = None
+_SP: SequenceParallelGroupCoordinator | None = None
+_PP: PipelineGroupCoordinator | None = None
+_CFG: GroupCoordinator | None = None
+_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"])
+
+
+def _split_tensor_dict(
+    tensor_dict: dict[str, torch.Tensor | Any],
+) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:
+    """Split the tensor dictionary into two parts:
+    1. A list of (key, value) pairs. If the value is a tensor, it is replaced
+         by its metadata.
+    2. A list of tensors.
+    """
+    metadata_list: list[tuple[str, Any]] = []
+    tensor_list: list[torch.Tensor] = []
+    for key, value in tensor_dict.items():
+        if isinstance(value, torch.Tensor):
+            # Note: we cannot use `value.device` here,
+            # because it contains not only the device type but also the device
+            # index (e.g. "cuda:0"). We only need the device type.
+            # receiving side will set the device index.
+            device = value.device.type
+            metadata_list.append(
+                (key, TensorMetadata(device, value.dtype, value.size()))
+            )
+            tensor_list.append(value)
+        else:
+            metadata_list.append((key, value))
+    return metadata_list, tensor_list
+
+
+_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
+
+
+def _register_group(group: "GroupCoordinator") -> None:
+    _groups[group.unique_name] = weakref.ref(group)
+
+
+def all_reduce(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
+    assert group_name in _groups, f"Group {group_name} is not found."
+    group = _groups[group_name]()
+    if group is None:
+        raise ValueError(f"Group {group_name} is destroyed.")
+    return group._all_reduce_out_place(tensor)
+
+
+def all_reduce_fake(tensor: torch.Tensor, group_name: str) -> torch.Tensor:
+    return torch.empty_like(tensor)
+
+
+def get_world_group() -> GroupCoordinator:
+    assert _WORLD is not None, "world group is not initialized"
+    return _WORLD
+
+
+def world_group_is_initialized() -> bool:
+    return _WORLD is not None
+
+
+def init_world_group(
+    ranks: list[int], local_rank: int, backend: str
+) -> GroupCoordinator:
+    return GroupCoordinator(
+        group_ranks=[ranks],
+        local_rank=local_rank,
+        torch_distributed_backend=backend,
+        use_device_communicator=True,
+        group_name="world",
+    )
+
+
+def _sync_srt_world_group() -> None:
+    import sglang.srt.distributed.parallel_state as srt_parallel_state
+
+    if srt_parallel_state._WORLD is None:
+        srt_parallel_state._WORLD = _WORLD
+
+
+def _clear_srt_world_group() -> None:
+    import sglang.srt.distributed.parallel_state as srt_parallel_state
+
+    if srt_parallel_state._WORLD is _WORLD:
+        srt_parallel_state._WORLD = None
+
+
+def init_parallel_group_coordinator(
+    group_ranks: List[List[int]],
+    local_rank: int,
+    backend: str,
+    parallel_mode: str,
+    **kwargs,
+) -> GroupCoordinator:
+    """Return a group coordinator for the given parallel mode."""
+    assert parallel_mode in [
+        "data",
+        "pipeline",
+        "tensor",
+        "sequence",
+        "classifier_free_guidance",
+        "vae_decode",
+    ], f"parallel_mode {parallel_mode} is not supported"
+    if parallel_mode == "pipeline":
+        return PipelineGroupCoordinator(
+            group_ranks=group_ranks,
+            local_rank=local_rank,
+            torch_distributed_backend=backend,
+            group_name="pp_group",
+        )
+    elif parallel_mode == "sequence":
+        return SequenceParallelGroupCoordinator(
+            group_ranks=group_ranks,
+            local_rank=local_rank,
+            torch_distributed_backend=backend,
+            group_name="sp_group",
+            **kwargs,
+        )
+    else:
+        return GroupCoordinator(
+            group_ranks=group_ranks,
+            local_rank=local_rank,
+            torch_distributed_backend=backend,
+            use_device_communicator=parallel_mode != "tensor",
+            use_srt_custom_allreduce=parallel_mode == "tensor",
+            group_name=(
+                "tp_group"
+                if parallel_mode == "tensor"
+                else (
+                    "vae_decode_group" if parallel_mode == "vae_decode" else "cfg_group"
+                )
+            ),
+        )
+
+
+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
+
+
+def init_distributed_environment(
+    world_size: int = 1,
+    rank: int = 0,
+    distributed_init_method: str = "env://",
+    local_rank: int = 0,
+    backend: str | None = None,
+    device_id: torch.device | None = None,
+    timeout: int | None = None,
+):
+    # Determine the appropriate backend based on the platform
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    if backend is None:
+        backend = current_platform.get_torch_distributed_backend_str()
+        logger.info(
+            "Using %s backend for %s platform", backend, current_platform.device_name
+        )
+
+    logger.debug(
+        "world_size=%d rank=%d local_rank=%d "
+        "distributed_init_method=%s backend=%s timeout=%s",
+        world_size,
+        rank,
+        local_rank,
+        distributed_init_method,
+        backend,
+        timeout,
+    )
+    if not torch.distributed.is_initialized():
+        assert distributed_init_method is not None, (
+            "distributed_init_method must be provided when initializing "
+            "distributed environment"
+        )
+
+        # For MPS, MUSA, and XPU, don't pass device_id as it doesn't support device indices
+        extra_args = (
+            {}
+            if (
+                current_platform.is_mps()
+                or current_platform.is_musa()
+                or current_platform.is_npu()
+                or current_platform.is_cpu()
+                or current_platform.is_xpu()
+            )
+            else dict(device_id=device_id)
+        )
+
+        if timeout is not None:
+
+            extra_args["timeout"] = datetime.timedelta(seconds=timeout)
+            logger.info(f"Setting distributed timeout to {timeout} seconds")
+
+        torch.distributed.init_process_group(
+            backend=backend,
+            init_method=distributed_init_method,
+            world_size=world_size,
+            rank=rank,
+            **extra_args,
+        )
+
+    # set the local rank
+    # local_rank is not available in torch ProcessGroup,
+    # see https://github.com/pytorch/pytorch/issues/122816
+    if local_rank == -1:
+        # local rank not set, this usually happens in single-node
+        # setting, where we can use rank as local rank
+        if distributed_init_method == "env://":
+            local_rank = envs.LOCAL_RANK
+        else:
+            local_rank = rank
+    global _WORLD
+    if _WORLD is None:
+        ranks = list(range(torch.distributed.get_world_size()))
+        _WORLD = init_world_group(ranks, local_rank, backend)
+    else:
+        assert (
+            _WORLD.world_size == torch.distributed.get_world_size()
+        ), "world group already initialized with a different world size"
+    _sync_srt_world_group()
+
+
+def get_sp_group() -> SequenceParallelGroupCoordinator:
+    assert _SP is not None, "sequence parallel group is not initialized"
+    return _SP
+
+
+def get_dp_group() -> GroupCoordinator:
+    assert _DP is not None, "data parallel group is not initialized"
+    return _DP
+
+
+# xDiT
+def initialize_model_parallel(
+    data_parallel_size: int = 1,
+    classifier_free_guidance_degree: int = 1,
+    sequence_parallel_degree: Optional[int] = None,
+    ulysses_degree: int = 1,
+    ring_degree: int = 1,
+    tensor_parallel_degree: int = 1,
+    pipeline_parallel_degree: int = 1,
+    vae_parallel_size: int = 0,
+    backend: Optional[str] = None,
+) -> None:
+    """
+    Initialize model parallel groups.
+
+    Arguments:
+        data_parallel_size: number of data parallelism groups.
+        classifier_free_guidance_degree: number of GPUs used for Classifier Free Guidance (CFG)
+        sequence_parallel_degree: number of GPUs used for sequence parallelism. sequence_parallel_degree = ulysses_degree * ring_degree
+        ulysses_degree: number of GPUs used for ulysses sequence parallelism.
+        ring_degree: number of GPUs used for ring sequence parallelism.
+        tensor_parallel_degree: number of GPUs used for tensor parallelism.
+        pipeline_parallel_degree: number of GPUs used for pipeline parallelism.
+        backend: distributed backend of pytorch collective comm.
+
+    Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
+    use 2 groups to parallelize the batch dim(dp), 2 groups to parallelize
+    split batch caused by CFG, and 2 GPUs to parallelize sequence.
+
+    dp_degree (2) * cfg_degree (2) * sp_degree (2) * pp_degree (2) = 16.
+
+    The present function will create 8 data-parallel groups,
+    8 CFG group, 8 pipeline-parallel group, and
+    8 sequence-parallel groups:
+        8 data-parallel groups:
+            [g0, g8], [g1, g9], [g2, g10], [g3, g11],
+            [g4, g12], [g5, g13], [g6, g14], [g7, g15]
+        8 CFG-parallel groups:
+            [g0, g4], [g1, g5], [g2, g6], [g3, g7],
+            [g8, g12], [g9, g13], [g10, g14], [g11, g15]
+        8 sequence-parallel groups:
+            [g0, g1], [g2, g3], [g4, g5], [g6, g7],
+            [g8, g9], [g10, g11], [g12, g13], [g14, g15]
+        8 pipeline-parallel groups:
+            [g0, g2], [g4, g6], [g8, g10], [g12, g14],
+            [g1, g3], [g5, g7], [g9, g11], [g13, g15]
+    Note that for efficiency, the caller should make sure adjacent ranks
+    are on the same DGX box. For example if we are using 2 DGX-1 boxes
+    with a total of 16 GPUs, rank 0 to 7 belong to the first box and
+    ranks 8 to 15 belong to the second box.
+    """
+
+    if backend is None:
+        from sglang.multimodal_gen.runtime.platforms import current_platform
+
+        backend = current_platform.get_torch_distributed_backend_str()
+    # Get world size and rank. Ensure some consistencies.
+    assert torch.distributed.is_initialized()
+    world_size: int = torch.distributed.get_world_size()
+    backend = backend or torch.distributed.get_backend(get_world_group().device_group)
+
+    dit_parallel_size = (
+        data_parallel_size
+        * classifier_free_guidance_degree
+        * sequence_parallel_degree
+        * pipeline_parallel_degree
+        * tensor_parallel_degree
+    )
+
+    if world_size < dit_parallel_size:
+        raise RuntimeError(
+            f"world_size ({world_size}) is less than "
+            f"tensor_parallel_degree ({tensor_parallel_degree}) x "
+            f"pipeline_parallel_degree ({pipeline_parallel_degree}) x"
+            f"sequence_parallel_degree ({sequence_parallel_degree}) x"
+            f"classifier_free_guidance_degree "
+            f"({classifier_free_guidance_degree}) x"
+            f"data_parallel_degree ({data_parallel_size})"
+        )
+
+    rank_generator: RankGenerator = RankGenerator(
+        tensor_parallel_degree,
+        sequence_parallel_degree,
+        pipeline_parallel_degree,
+        classifier_free_guidance_degree,
+        data_parallel_size,
+        "tp-sp-pp-cfg-dp",
+    )
+    global _DP
+    assert _DP is None, "data parallel group is already initialized"
+    _DP = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("dp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="data",
+    )
+
+    global _CFG
+    assert _CFG is None, "classifier_free_guidance group is already initialized"
+    _CFG = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("cfg"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="classifier_free_guidance",
+    )
+    global _PP
+    assert _PP is None, "pipeline model parallel group is already initialized"
+    _PP = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("pp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="pipeline",
+    )
+
+    global _SP
+    assert _SP is None, "sequence parallel group is already initialized"
+
+    try:
+        from .parallel_groups import PROCESS_GROUP as _YC_PROCESS_GROUP
+        from .parallel_groups import (
+            set_seq_parallel_pg_by_sp_groups as _set_seq_parallel_pg_by_sp_groups,
+        )
+    except ImportError:
+        _set_seq_parallel_pg_by_sp_groups = None
+
+        class _DummyProcessGroup:
+            ULYSSES_PG = torch.distributed.group.WORLD
+            RING_PG = torch.distributed.group.WORLD
+
+        PROCESS_GROUP = _DummyProcessGroup()
+    else:
+        # Build SGLang Diffusion SP sub-groups based on the true SP groups. This is
+        # critical when TP>1, because SP groups may be strided in global ranks
+        # (e.g., tp-sp order).
+        sp_groups = rank_generator.get_ranks("sp")
+        _set_seq_parallel_pg_by_sp_groups(
+            sp_ulysses_degree=ulysses_degree,
+            sp_ring_degree=ring_degree,
+            rank=get_world_group().rank,
+            sp_groups=sp_groups,
+        )
+        PROCESS_GROUP = _YC_PROCESS_GROUP
+
+    _SP = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("sp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="sequence",
+        ulysses_group=PROCESS_GROUP.ULYSSES_PG,
+        ring_group=PROCESS_GROUP.RING_PG,
+    )
+
+    global _TP
+    assert _TP is None, "Tensor parallel group is already initialized"
+    _TP = init_parallel_group_coordinator(
+        group_ranks=rank_generator.get_ranks("tp"),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="tensor",
+    )
+
+    global _VAE_DECODE
+    assert _VAE_DECODE is None, "VAE decode parallel group is already initialized"
+    _VAE_DECODE = init_parallel_group_coordinator(
+        group_ranks=_get_vae_decode_group_ranks(rank_generator),
+        local_rank=get_world_group().local_rank,
+        backend=backend,
+        parallel_mode="vae_decode",
+    )
+
+    if vae_parallel_size > 0:
+        init_vae_group(dit_parallel_size, vae_parallel_size, backend)
+    init_dit_group(dit_parallel_size, backend)
+
+
+def get_sp_world_size() -> int:
+    """Return world size for the sequence model parallel group."""
+    return get_sp_group().world_size
+
+
+def get_sp_parallel_rank() -> int:
+    """Return my rank for the sequence model parallel group."""
+    return get_sp_group().rank_in_group
+
+
+def get_world_size() -> int:
+    """Return world size for the world group."""
+    return get_world_group().world_size
+
+
+def get_world_rank() -> int:
+    """Return my rank for the world group."""
+    return get_world_group().rank
+
+
+def get_dp_world_size() -> int:
+    """Return world size for the data parallel group."""
+    return get_dp_group().world_size
+
+
+def get_dp_rank() -> int:
+    """Return my rank for the data parallel group."""
+    return get_dp_group().rank_in_group
+
+
+def maybe_init_distributed_environment_and_model_parallel(
+    tp_size: int,
+    sp_size: int,
+    cfg_degree: int = 1,
+    ulysses_degree: int = 1,
+    ring_degree: int = 1,
+    dp_size: int = 1,
+    distributed_init_method: str = "env://",
+    dist_timeout: int | None = None,
+):
+    from sglang.multimodal_gen.runtime.platforms import current_platform
+
+    if _WORLD is not None and model_parallel_is_initialized():
+        # make sure the tp and sp sizes are correct
+        assert (
+            get_tp_world_size() == tp_size
+        ), f"You are trying to initialize model parallel groups with size {tp_size}, but they are already initialized with size {get_tp_world_size()}"
+        assert (
+            get_sp_world_size() == sp_size
+        ), f"You are trying to initialize model parallel groups with size {sp_size}, but they are already initialized with size {get_sp_world_size()}"
+        return
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    rank = int(os.environ.get("RANK", 0))
+    device = get_local_torch_device()
+    logger.info(
+        "Initializing distributed environment with world_size=%d, device=%s, timeout=%s",
+        world_size,
+        device,
+        dist_timeout,
+        main_process_only=False,
+    )
+
+    init_distributed_environment(
+        world_size=world_size,
+        rank=rank,
+        local_rank=local_rank,
+        distributed_init_method=distributed_init_method,
+        device_id=device,
+        backend=current_platform.get_torch_distributed_backend_str(),
+        timeout=dist_timeout,
+    )
+    initialize_model_parallel(
+        data_parallel_size=dp_size,
+        classifier_free_guidance_degree=cfg_degree,
+        tensor_parallel_degree=tp_size,
+        ulysses_degree=ulysses_degree,
+        ring_degree=ring_degree,
+        sequence_parallel_degree=sp_size,
+    )
+
+    # Only set CUDA device if we're on a CUDA platform
+    if current_platform.is_cuda_alike():
+        device = torch.device(f"cuda:{local_rank}")
+        torch.cuda.set_device(device)
+    elif current_platform.is_npu():
+        device = torch.device(f"npu:{local_rank}")
+        torch.npu.set_device(device)
+
+
+def model_parallel_is_initialized() -> bool:
+    """Check if model parallel groups are initialized."""
+    return (
+        _DP is not None
+        and _CFG is not None
+        and _SP is not None
+        and _PP is not None
+        and _TP is not None
+        and _VAE_DECODE is not None
+    )
+
+
+_TP_STATE_PATCHED = False
+
+
+@contextmanager
+def patch_tensor_parallel_group(tp_group: GroupCoordinator):
+    """Patch the tp group temporarily until this function ends.
+
+    This method is for draft workers of speculative decoding to run draft model
+    with different tp degree from that of target model workers.
+
+    """
+    global _TP_STATE_PATCHED
+    assert not _TP_STATE_PATCHED, "Should not call when it's already patched"
+
+    _TP_STATE_PATCHED = True
+    old_tp_group = get_tp_group()
+    global _TP
+    _TP = tp_group
+    try:
+        yield
+    finally:
+        # restore the original state
+        _TP_STATE_PATCHED = False
+        _TP = old_tp_group
+
+
+def get_tp_world_size() -> int:
+    """Return world size for the tensor model parallel group."""
+    return get_tp_group().world_size
+
+
+def get_tp_rank() -> int:
+    """Return my rank for the tensor model parallel group."""
+    return get_tp_group().rank_in_group
+
+
+def destroy_distributed_environment() -> None:
+    global _WORLD
+    _clear_srt_world_group()
+    if _WORLD:
+        _WORLD.destroy()
+    _WORLD = None
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
+
+
+def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
+    destroy_model_parallel()
+    destroy_distributed_environment()
+    with contextlib.suppress(AssertionError):
+        torch.distributed.destroy_process_group()
+    if shutdown_ray:
+        import ray  # Lazy import Ray
+
+        ray.shutdown()
+
+
+def is_the_same_node_as(
+    pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0
+) -> list[int]:
+    """
+    This is a collective operation that returns if each rank is in the same node
+    as the source rank. It tests if processes are attached to the same
+    memory system (shared access to shared memory).
+    """
+    if isinstance(pg, ProcessGroup):
+        assert (
+            torch.distributed.get_backend(pg) != torch.distributed.Backend.NCCL
+        ), "in_the_same_node_as should be tested with a non-NCCL group."
+        # local rank inside the group
+        rank = torch.distributed.get_rank(group=pg)
+        world_size = torch.distributed.get_world_size(group=pg)
+
+        # global ranks of the processes in the group
+        ranks = torch.distributed.get_process_group_ranks(pg)
+    else:
+        rank = pg.rank
+        world_size = pg.world_size
+        ranks = list(range(world_size))
+
+    # local tensor in each process to store the result
+    is_in_the_same_node = torch.tensor([0] * world_size, dtype=torch.int32)
+
+    magic_message = b"magic_message"
+    shm = None
+
+    try:
+        with contextlib.suppress(OSError):
+            if rank == source_rank:
+                # create a shared memory segment
+                shm = shared_memory.SharedMemory(create=True, size=128)
+                shm.buf[: len(magic_message)] = magic_message
+                if isinstance(pg, ProcessGroup):
+                    torch.distributed.broadcast_object_list(
+                        [shm.name], src=ranks[source_rank], group=pg
+                    )
+                else:
+                    pg.broadcast_obj(shm.name, src=source_rank)
+                is_in_the_same_node[rank] = 1
+            else:
+                # try to open the shared memory segment
+                if isinstance(pg, ProcessGroup):
+                    recv = [None]
+                    torch.distributed.broadcast_object_list(
+                        recv, src=ranks[source_rank], group=pg
+                    )
+                    name = recv[0]
+                else:
+                    name = pg.broadcast_obj(None, src=source_rank)
+                # fix to https://stackoverflow.com/q/62748654/9191338
+                # Python incorrectly tracks shared memory even if it is not
+                # created by the process. The following patch is a workaround.
+                with patch(
+                    "multiprocessing.resource_tracker.register",
+                    lambda *args, **kwargs: None,
+                ):
+                    shm = shared_memory.SharedMemory(name=name)
+                if shm.buf[: len(magic_message)] == magic_message:
+                    is_in_the_same_node[rank] = 1
+    except Exception as e:
+        logger.error("Error ignored in is_in_the_same_node: %s", e)
+    finally:
+        if shm:
+            shm.close()
+
+    if isinstance(pg, ProcessGroup):
+        torch.distributed.barrier(group=pg)
+    else:
+        pg.barrier()
+
+    # clean up the shared memory segment
+    with contextlib.suppress(OSError):
+        if rank == source_rank and shm:
+            shm.unlink()
+
+    if isinstance(pg, ProcessGroup):
+        torch.distributed.all_reduce(is_in_the_same_node, group=pg)
+        aggregated_data = is_in_the_same_node
+    else:
+        aggregated_data = torch.zeros_like(is_in_the_same_node)
+        for i in range(world_size):
+            rank_data = pg.broadcast_obj(is_in_the_same_node, src=i)
+            aggregated_data += rank_data
+
+    return [x == 1 for x in aggregated_data.tolist()]
+
+
+def get_tensor_model_parallel_world_size() -> int:
+    """Return world size for the tensor model parallel group."""
+    return get_tp_world_size()
+
+
+def get_tensor_model_parallel_rank() -> int:
+    """Return my rank for the tensor model parallel group."""
+    return get_tp_rank()
+
+
+def get_sequence_parallel_world_size() -> int:
+    """Return world size for the sequence parallel group."""
+    return get_sp_world_size()
+
+
+def get_sequence_parallel_rank() -> int:
+    """Return my rank for the sequence parallel group."""
+    return get_sp_parallel_rank()
+
+
+def get_ulysses_parallel_world_size() -> int:
+    return get_sp_group().ulysses_world_size
+
+
+def get_ulysses_parallel_rank() -> int:
+    return get_sp_group().ulysses_rank
+
+
+def get_ring_parallel_world_size() -> int:
+    return get_sp_group().ring_world_size
+
+
+def get_ring_parallel_rank() -> int:
+    return get_sp_group().ring_rank
+
+
+# PP
+def get_pp_group() -> PipelineGroupCoordinator:
+    assert _PP is not None, "pipeline model parallel group is not initialized"
+    return _PP
+
+
+def get_pipeline_parallel_world_size() -> int:
+    """Return world size for the pipeline model parallel group."""
+    return get_pp_group().world_size
+
+
+def get_pipeline_parallel_rank() -> int:
+    """Return my rank for the pipeline model parallel group."""
+    return get_pp_group().rank_in_group
+
+
+def is_pipeline_first_stage() -> bool:
+    """Return True if in the first pipeline model parallel stage, False otherwise."""
+    return get_pipeline_parallel_rank() == 0
+
+
+def is_pipeline_last_stage() -> bool:
+    """Return True if in the last pipeline model parallel stage, False otherwise."""
+    return get_pipeline_parallel_rank() == (get_pipeline_parallel_world_size() - 1)
+
+
+# CFG
+def get_cfg_group() -> GroupCoordinator:
+    assert (
+        _CFG is not None
+    ), "classifier_free_guidance parallel group is not initialized"
+    return _CFG
+
+
+def get_classifier_free_guidance_world_size() -> int:
+    """Return world size for the classifier_free_guidance parallel group."""
+    return get_cfg_group().world_size
+
+
+def get_classifier_free_guidance_rank() -> int:
+    """Return my rank for the classifier_free_guidance parallel group."""
+    return get_cfg_group().rank_in_group
+
+
+def get_data_parallel_world_size() -> int:
+    """Return world size for the data parallel group."""
+    return get_dp_world_size()
+
+
+def get_data_parallel_rank() -> int:
+    """Return my rank for the data parallel group."""
+    return get_dp_rank()
+
+
+def is_dp_last_group() -> bool:
+    """Return True if in the last data parallel group, False otherwise."""
+    return (
+        get_sequence_parallel_rank() == (get_sequence_parallel_world_size() - 1)
+        and get_classifier_free_guidance_rank()
+        == (get_classifier_free_guidance_world_size() - 1)
+        and get_pipeline_parallel_rank() == (get_pipeline_parallel_world_size() - 1)
+    )
+
+
+def get_dit_world_size() -> int:
+    """Return world size for the DiT model (excluding VAE)."""
+    return (
+        get_data_parallel_world_size()
+        * get_classifier_free_guidance_world_size()
+        * get_sequence_parallel_world_size()
+        * get_pipeline_parallel_world_size()
+        * get_tensor_model_parallel_world_size()
+    )
+
+
+def get_vae_parallel_group() -> ProcessGroup:
+    assert _VAE is not None, "VAE parallel group is not initialized"
+    return _VAE
+
+
+def get_vae_parallel_world_size() -> int:
+    """Return world size for the VAE parallel group."""
+    return torch.distributed.get_world_size(group=get_vae_parallel_group())
+
+
+def get_vae_parallel_rank() -> int:
+    """Return my rank for the VAE parallel group."""
+    return torch.distributed.get_rank(group=get_vae_parallel_group())
+
+
+def get_decode_parallel_group_coordinator() -> GroupCoordinator:
+    assert _VAE_DECODE is not None, "VAE decode parallel group is not initialized"
+    return _VAE_DECODE
+
+
+def get_decode_parallel_world_size() -> int:
+    return get_decode_parallel_group_coordinator().world_size
+
+
+def get_decode_parallel_rank() -> int:
+    return get_decode_parallel_group_coordinator().rank_in_group
+
+
+def init_dit_group(
+    dit_parallel_size: int,
+    backend: str,
+) -> None:
+    global _DIT
+    assert _DIT is None, "DIT group is already initialized"
+    _DIT = torch.distributed.new_group(
+        ranks=list(range(dit_parallel_size)), backend=backend
+    )
+
+
+def get_dit_group() -> ProcessGroup:
+    assert _DIT is not None, "DIT group is not initialized"
+    return _DIT
+
+
+def init_vae_group(
+    dit_parallel_size: int,
+    vae_parallel_size: int,
+    backend: str,
+):
+    # Initialize VAE group first
+    global _VAE
+    assert _VAE is None, "VAE parallel group is already initialized"
+    vae_ranks = list(range(dit_parallel_size, dit_parallel_size + vae_parallel_size))
+    _VAE = torch.distributed.new_group(ranks=vae_ranks, backend=backend)
+
+
+def destroy_model_parallel() -> None:
+    """Set the groups to none and destroy them."""
+    global _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE
+
+    for group in (_TP, _SP, _DP, _CFG, _PP, _VAE_DECODE):
+        if group is not None:
+            group.destroy()
+
+    for group in (_DIT, _VAE):
+        if group is not None:
+            torch.distributed.destroy_process_group(group)
+
+    _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE = (None,) * 8

python/sglang/multimodal_gen/runtime/distributed/sp_shard_utils.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Unified sequence-parallel shard / pad / gather helpers.
+
+Layout invariant: padding always sits at the end of the LAST rank's local
+chunk, so the ulysses-gathered sequence carries one contiguous pad block at its
+global tail. `tail_attn_meta` then lets attention skip that block for free
+(the pad becomes its own varlen segment - no repacking, no mask compute).
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    sequence_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_ring_parallel_world_size,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+)
+
+# Text shorter than this stays replicated instead of SP-sharded (see
+# plan_text_strategy). 0 = always shard when legal; H100 bench showed sharding
+# wins from trivial lengths on, so the knob exists only as an escape hatch.
+_TEXT_SHARD_MIN = int(os.environ.get("SGLANG_SP_TEXT_SHARD_MIN", "0"))
+
+
+@dataclass(frozen=True)
+class SpShard:
+    """Facts of one tail-padded even shard, shared by tensors of that stream."""
+
+    orig_len: int  # real tokens (global)
+    local_len: int  # per-rank chunk length (equal on every rank)
+    num_pad: int  # pad tokens, all at the last rank's local tail
+    sp_size: int
+    sp_rank: int
+
+    @property
+    def local_pad(self) -> int:
+        """Pad rows inside THIS rank's chunk (tail rows of the last rank)."""
+        return self.num_pad if self.sp_rank == self.sp_size - 1 else 0
+
+    @property
+    def local_real_len(self) -> int:
+        return self.local_len - self.local_pad
+
+
+def build_shard_plan(seq_len: int) -> SpShard:
+    """Shard math only; tensors are sliced separately via `shard_like`."""
+    sp_size = get_sp_world_size()
+    if sp_size <= 1:
+        return SpShard(seq_len, seq_len, 0, 1, 0)
+    local_len = (seq_len + sp_size - 1) // sp_size
+    return SpShard(
+        orig_len=seq_len,
+        local_len=local_len,
+        num_pad=local_len * sp_size - seq_len,
+        sp_size=sp_size,
+        sp_rank=get_sp_parallel_rank(),
+    )
+
+
+def shard_like(
+    x: torch.Tensor, shard: SpShard, dim: int = 1, pad_mode: str = "zeros"
+) -> torch.Tensor:
+    """Apply a planned shard to one tensor (RoPE caches use the same plan as
+    hidden states so their chunks stay aligned)."""
+    if shard.sp_size <= 1:
+        return x
+    if shard.num_pad > 0:
+        if pad_mode == "repeat_last":
+            pad = x.narrow(dim, x.shape[dim] - 1, 1)
+            pad = pad.expand(
+                *[shard.num_pad if i == dim else -1 for i in range(x.dim())]
+            )
+            x = torch.cat([x, pad], dim=dim)
+        else:
+            # F.pad pads dims last-to-first: (left, right) pairs from dim -1.
+            pads = [0, 0] * (x.dim() - 1 - dim) + [0, shard.num_pad]
+            x = F.pad(x, pads)
+    return x.narrow(dim, shard.sp_rank * shard.local_len, shard.local_len)
+
+
+def shard_seq(
+    x: torch.Tensor, dim: int = 1, pad_mode: str = "zeros"
+) -> tuple[torch.Tensor, SpShard]:
+    """
+    mode:
+        zeroes: pad with zeroes at tail
+        repeat_last: repeat the last token, only for rotary embedding
+    """
+    shard = build_shard_plan(x.shape[dim])
+    return shard_like(x, shard, dim=dim, pad_mode=pad_mode), shard
+
+
+def gather_seq(local: torch.Tensor, orig_len: int, dim: int = 1) -> torch.Tensor:
+    """All-gather an SP-sharded sequence and trim the tail padding"""
+    if get_sp_world_size() <= 1:
+        return local
+    full = sequence_model_parallel_all_gather(local.contiguous(), dim=dim)
+    if full.shape[dim] > orig_len:
+        full = full.narrow(dim, 0, orig_len)
+    return full
+
+
+def shard_seq_prefix(
+    x: torch.Tensor, prefix_len: int, shard: SpShard, dim: int = 0
+) -> torch.Tensor:
+    """Shard only the leading ``prefix_len`` rows (e.g. the text segment of a
+    joint RoPE cache) with an existing plan; the remainder is kept as-is."""
+    rest = x.shape[dim] - prefix_len
+    return torch.cat(
+        [
+            shard_like(x.narrow(dim, 0, prefix_len), shard, dim=dim),
+            x.narrow(dim, prefix_len, rest),
+        ],
+        dim=dim,
+    )
+
+
+def join_seqs(
+    prefix: torch.Tensor, body: torch.Tensor, local_pad: int, dim: int = 1
+) -> torch.Tensor:
+    """Concatenate local sharded ``[prefix (txt tokens, padding tokens), body (img tokens)]`` for joint attention, while relocating the
+     prefix's ``local_pad`` tail rows behind the body.
+
+    Why leave the padding at tail: the shard pads the *text* chunk, but the local joint layout is
+    [text, image].
+
+    In naive implementation, after the ulysses all-to-all, that pad would sit mid-sequence (of last rank)
+    ([... txt_last, PAD, img_last]), which required further mem copy (for the padding tokens), inefficient in this case
+
+    With the pad relocated behind the image, the padding forms one global-tail block that the zero-copy varlen
+    path (tail_attn_meta, implemented in USPAttention.forward) skips for free
+    """
+    if local_pad > 0:
+        real = prefix.shape[dim] - local_pad
+        return torch.cat(
+            [
+                # txt tokens
+                prefix.narrow(dim, 0, real),
+                body,
+                # leave the padding at global-tail
+                prefix.narrow(dim, real, local_pad),
+            ],
+            dim=dim,
+        )
+    return torch.cat([prefix, body], dim=dim)
+
+
+def split_seqs(
+    joint: torch.Tensor, prefix_len: int, local_pad: int, dim: int = 1
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Inverse of ``join_seqs``: recover ``(prefix, body)`` from the joint output, with the pad rows rejoining the prefix tail so the residual text
+    stream keeps its per-rank shape.
+
+     ([... txt_last, PAD, img_last]) -> prefix (txt + pad), body (img)
+    """
+    total = joint.shape[dim]
+    if local_pad > 0:
+        real = prefix_len - local_pad
+        body_end = total - local_pad
+        prefix = torch.cat(
+            [joint.narrow(dim, 0, real), joint.narrow(dim, body_end, local_pad)],
+            dim=dim,
+        )
+        return prefix, joint.narrow(dim, real, body_end - real)
+    return (
+        joint.narrow(dim, 0, prefix_len),
+        joint.narrow(dim, prefix_len, total - prefix_len),
+    )
+
+
+def should_shard_text(txt_len: int) -> bool:
+    """True when the joint-attention text stream should be SP-sharded here
+    (see plan_text_strategy for the policy)."""
+    return get_sp_world_size() > 1 and plan_text_strategy(txt_len) == "shard"
+
+
+def tail_attn_meta(
+    shard: SpShard,
+    batch_size: int,
+    device: torch.device,
+    image_seq_len: int = 0,
+) -> dict | None:
+    """Per-request attention meta for a tail-padded shard: `cu_seqlens_tail`
+    splits each batch row into [valid | pad] varlen segments over the gathered
+    layout, so USPAttention runs varlen FA on the padded q/k/v with zero
+    repacking. Built once per request, reused by every block."""
+    if shard.sp_size <= 1 or shard.num_pad == 0:
+        return None
+    seq = shard.sp_size * (shard.local_len + image_seq_len)
+    valid = seq - shard.num_pad
+    row = torch.tensor([valid, shard.num_pad], dtype=torch.int32, device=device)
+    seglens = row.repeat(batch_size)
+    cu_seqlens = torch.zeros(2 * batch_size + 1, dtype=torch.int32, device=device)
+    cu_seqlens[1:] = torch.cumsum(seglens, dim=0)
+    return {
+        "pad_start": valid,
+        "pad_end": seq,
+        "local_pad": shard.local_pad,
+        "cu_seqlens_tail": cu_seqlens,
+        "max_seqlen_tail": max(valid, shard.num_pad),
+    }
+
+
+def plan_text_strategy(txt_len: int) -> str:
+    """Choose "shard" or "replicate" for the joint-attention text stream.
+
+    Prefer "shard" by default. for small sequence (shorter than SGLANG_SP_TEXT_SHARD_MIN), choose "replicate" for better performance
+
+    """
+    sp_size = get_sp_world_size()
+    if sp_size <= 1:
+        return "replicate"
+    if txt_len % sp_size != 0 and get_ring_parallel_world_size() > 1:
+        return "replicate"
+    if txt_len < _TEXT_SHARD_MIN:
+        return "replicate"
+    return "shard"

python/sglang/multimodal_gen/runtime/distributed/utils.py

@@ -0,0 +1,196 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/utils.py
+
+# Copyright 2023 The vLLM team.
+# Adapted from
+# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+import dataclasses
+import pickle
+import time
+from collections import deque
+from collections.abc import Sequence
+from typing import Any
+
+import torch
+from torch.distributed import TCPStore
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def ensure_divisibility(numerator, denominator) -> None:
+    """Ensure that numerator is divisible by the denominator."""
+    assert numerator % denominator == 0, "{} is not divisible by {}".format(
+        numerator, denominator
+    )
+
+
+def divide(numerator: int, denominator: int) -> int:
+    """Ensure that numerator is divisible by the denominator and return
+    the division value."""
+    ensure_divisibility(numerator, denominator)
+    return numerator // denominator
+
+
+def split_tensor_along_last_dim(
+    tensor: torch.Tensor,
+    num_partitions: int,
+    contiguous_split_chunks: bool = False,
+) -> Sequence[torch.Tensor]:
+    """Split a tensor along its last dimension.
+
+    Arguments:
+        tensor: input tensor.
+        num_partitions: number of partitions to split the tensor
+        contiguous_split_chunks: If True, make each chunk contiguous
+                                 in memory.
+
+    Returns:
+        A list of Tensors
+    """
+    # Get the size and dimension.
+    last_dim = tensor.dim() - 1
+    last_dim_size = divide(tensor.size()[last_dim], num_partitions)
+    # Split.
+    tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
+    # NOTE: torch.split does not create contiguous tensors by default.
+    if contiguous_split_chunks:
+        return tuple(chunk.contiguous() for chunk in tensor_list)
+
+    return tuple(tensor_list)
+
+
+@dataclasses.dataclass
+class StatelessProcessGroup:
+    """A dataclass to hold a metadata store, and the rank, world_size of the
+    group. Only use it to communicate metadata between processes.
+    For data-plane communication, create NCCL-related objects.
+    """
+
+    rank: int
+    world_size: int
+    store: torch._C._distributed_c10d.Store
+    data_expiration_seconds: int = 3600  # 1 hour
+
+    # dst rank -> counter
+    send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict)
+    # src rank -> counter
+    recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict)
+    broadcast_send_counter: int = 0
+    broadcast_recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict)
+
+    # A deque to store the data entries, with key and timestamp.
+    entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque)
+
+    def __post_init__(self):
+        assert self.rank < self.world_size
+        self.send_dst_counter = {i: 0 for i in range(self.world_size)}
+        self.recv_src_counter = {i: 0 for i in range(self.world_size)}
+        self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)}
+
+    def send_obj(self, obj: Any, dst: int):
+        """Send an object to a destination rank."""
+        self.expire_data()
+        key = f"send_to/{dst}/{self.send_dst_counter[dst]}"
+        self.store.set(key, pickle.dumps(obj))
+        self.send_dst_counter[dst] += 1
+        self.entries.append((key, time.perf_counter()))
+
+    def expire_data(self) -> None:
+        """Expire data that is older than `data_expiration_seconds` seconds."""
+        while self.entries:
+            # check the oldest entry
+            key, timestamp = self.entries[0]
+            if time.perf_counter() - timestamp > self.data_expiration_seconds:
+                self.store.delete_key(key)
+                self.entries.popleft()
+            else:
+                break
+
+    def recv_obj(self, src: int) -> Any:
+        """Receive an object from a source rank."""
+        obj = pickle.loads(
+            self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}")
+        )
+        self.recv_src_counter[src] += 1
+        return obj
+
+    def broadcast_obj(self, obj: Any | None, src: int) -> Any:
+        """Broadcast an object from a source rank to all other ranks.
+        It does not clean up after all ranks have received the object.
+        Use it for limited times, e.g., for initialization.
+        """
+        if self.rank == src:
+            self.expire_data()
+            key = f"broadcast_from/{src}/" f"{self.broadcast_send_counter}"
+            self.store.set(key, pickle.dumps(obj))
+            self.broadcast_send_counter += 1
+            self.entries.append((key, time.perf_counter()))
+            return obj
+        else:
+            key = f"broadcast_from/{src}/" f"{self.broadcast_recv_src_counter[src]}"
+            recv_obj = pickle.loads(self.store.get(key))
+            self.broadcast_recv_src_counter[src] += 1
+            return recv_obj
+
+    def all_gather_obj(self, obj: Any) -> list[Any]:
+        """All gather an object from all ranks."""
+        gathered_objs = []
+        for i in range(self.world_size):
+            if i == self.rank:
+                gathered_objs.append(obj)
+                self.broadcast_obj(obj, src=self.rank)
+            else:
+                recv_obj = self.broadcast_obj(None, src=i)
+                gathered_objs.append(recv_obj)
+        return gathered_objs
+
+    def barrier(self):
+        """A barrier to synchronize all ranks."""
+        for i in range(self.world_size):
+            if i == self.rank:
+                self.broadcast_obj(None, src=self.rank)
+            else:
+                self.broadcast_obj(None, src=i)
+
+    @staticmethod
+    def create(
+        host: str,
+        port: int,
+        rank: int,
+        world_size: int,
+        data_expiration_seconds: int = 3600,
+    ) -> "StatelessProcessGroup":
+        """A replacement for `torch.distributed.init_process_group` that does not
+        pollute the global state.
+
+        If we have process A and process B called `torch.distributed.init_process_group`
+        to form a group, and then we want to form another group with process A, B, C,
+        D, it is not possible in PyTorch, because process A and process B have already
+        formed a group, and process C and process D cannot join that group. This
+        function is a workaround for this issue.
+
+        `torch.distributed.init_process_group` is a global call, while this function
+        is a stateless call. It will return a `StatelessProcessGroup` object that can be
+        used for exchanging metadata. With this function, process A and process B
+        can call `StatelessProcessGroup.create` to form a group, and then process A, B,
+        C, and D can call `StatelessProcessGroup.create` to form another group.
+        """  # noqa
+        store = TCPStore(
+            host_name=host,
+            port=port,
+            world_size=world_size,
+            is_master=(rank == 0),
+        )
+
+        return StatelessProcessGroup(
+            rank=rank,
+            world_size=world_size,
+            store=store,
+            data_expiration_seconds=data_expiration_seconds,
+        )