[Diffusion] Add native SymmMem Fast Ulysses transport (#6340)

Signed-off-by: roG0d <baonudesifeizhai@gmail.com>

tests/config/test_omni_config.py

@@ -653,6 +653,7 @@ def test_sub_config_fields_match_structured_scopes():
         "ring_degree",
         "allgather_degree",
         "ulysses_mode",
+        "ulysses_a2a_permute",
         "cfg_parallel_size",
         "vae_patch_parallel_size",
         "vae_parallel_mode",
@@ -858,6 +859,23 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr
     assert stage.parallel_config.world_size == 2
 
 
+def test_from_pipeline_config_preserves_deployed_ulysses_a2a_permute() -> None:
+    pipeline = _resolve_pipeline_or_skip("dreamzero")
+    deploy = DeployConfig(
+        stages=[
+            StageDeployConfig(
+                stage_id=0,
+                ulysses_a2a_permute=True,
+            )
+        ]
+    )
+
+    stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
+
+    assert isinstance(stage, VllmOmniDiffusionStageConfig)
+    assert stage.parallel_config.ulysses_a2a_permute is True
+
+
 def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism():
     cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
 

tests/diffusion/distributed/test_a2a_permute.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed import a2a_permute
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_jit_build_includes_cuda_headers_from_nvidia_wheels(tmp_path, monkeypatch) -> None:
+    cu13_include = tmp_path / "nvidia" / "cu13" / "include"
+    nccl_include = tmp_path / "nvidia" / "nccl" / "include"
+    nccl_lib = tmp_path / "nvidia" / "nccl" / "lib" / "libnccl.so.2"
+    cu13_include.mkdir(parents=True)
+    nccl_include.mkdir(parents=True)
+    nccl_lib.parent.mkdir(parents=True)
+    (cu13_include / "cusparse.h").touch()
+    (nccl_include / "nccl.h").touch()
+    nccl_lib.touch()
+
+    load_kwargs = {}
+    monkeypatch.setattr(a2a_permute.sysconfig, "get_paths", lambda: {"purelib": str(tmp_path)})
+    monkeypatch.setattr(
+        torch.utils.cpp_extension,
+        "load",
+        lambda **kwargs: load_kwargs.update(kwargs),
+    )
+    monkeypatch.setattr(a2a_permute.symm_mem, "set_backend", lambda _backend: None)
+    monkeypatch.setattr(a2a_permute, "_BUILT", False)
+
+    a2a_permute.ensure_a2a_permute_available()
+
+    assert set(load_kwargs["extra_include_paths"]) == {str(cu13_include), str(nccl_include)}
+    assert load_kwargs["extra_ldflags"] == [str(nccl_lib)]
+
+
+@dataclass
+class _FakeAllocation:
+    size: int
+    device: torch.device
+
+
+class _FakeHandle:
+    rank = 0
+
+    def __init__(self, allocation: _FakeAllocation) -> None:
+        self.allocation = allocation
+        self.views: list[tuple[tuple[int, ...], torch.dtype]] = []
+
+    def get_buffer(self, rank, shape, dtype):
+        assert rank == self.rank
+        self.views.append((shape, dtype))
+        return self.allocation, shape, dtype
+
+
+def test_workspace_reuses_peak_capacity_across_shapes(monkeypatch) -> None:
+    allocations: list[_FakeAllocation] = []
+    all_reduce_calls: list[object] = []
+    synchronize_calls: list[torch.device] = []
+    stream_id = [1]
+    a2a_permute._SYMM_WORKSPACES.clear()
+
+    monkeypatch.setattr(a2a_permute.torch.cuda, "is_current_stream_capturing", lambda: False)
+    monkeypatch.setattr(
+        a2a_permute.torch.cuda,
+        "current_stream",
+        lambda _device: SimpleNamespace(cuda_stream=stream_id[0]),
+    )
+    monkeypatch.setattr(a2a_permute, "_resolve_process_group", lambda _name: object())
+    monkeypatch.setattr(a2a_permute.torch, "ones", lambda *_args, **_kwargs: object())
+    monkeypatch.setattr(
+        a2a_permute.dist,
+        "all_reduce",
+        lambda _tensor, group: all_reduce_calls.append(group),
+    )
+    monkeypatch.setattr(
+        a2a_permute.torch.accelerator,
+        "synchronize",
+        lambda device: synchronize_calls.append(device),
+    )
+
+    def empty(size, *, dtype, device):
+        assert dtype == torch.uint8
+        allocation = _FakeAllocation(size=size, device=torch.device(device))
+        allocations.append(allocation)
+        return allocation
+
+    monkeypatch.setattr(a2a_permute.symm_mem, "empty", empty)
+    monkeypatch.setattr(a2a_permute.symm_mem, "rendezvous", lambda allocation, _group: _FakeHandle(allocation))
+
+    device = torch.device("cuda:0")
+    first = a2a_permute._get_symm_buffer((2, 3), torch.float16, device, "group")
+    smaller = a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
+    larger = a2a_permute._get_symm_buffer((4, 3), torch.float16, device, "group")
+
+    assert first[0].size == 12
+    assert smaller[0] is first[0]
+    assert larger[0].size == 24
+    assert [allocation.size for allocation in allocations] == [12, 24]
+    assert len(a2a_permute._SYMM_WORKSPACES) == 1
+    assert len(all_reduce_calls) == 1
+    assert synchronize_calls == [device, device]
+
+    stream_id[0] = 2
+    with pytest.raises(RuntimeError, match="single CUDA stream"):
+        a2a_permute._get_symm_buffer((1, 4), torch.float16, device, "group")
+
+    a2a_permute.clear_a2a_permute_workspaces()
+    assert not a2a_permute._SYMM_WORKSPACES
+    assert synchronize_calls == [device, device, device]

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -33,6 +33,19 @@ def test_default_stage_config_includes_cache_backend():
     assert engine_args["model_stage"] == "diffusion"
 
 
+def test_default_stage_config_preserves_ulysses_a2a_permute() -> None:
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "ulysses_degree": 4,
+            "ulysses_a2a_permute": True,
+        }
+    )[0]
+
+    parallel_config = stage_cfg["engine_args"]["parallel_config"]
+    assert parallel_config.ulysses_degree == 4
+    assert parallel_config.ulysses_a2a_permute is True
+
+
 def test_default_stage_config_preserves_model_extras():
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0]
 

tests/entrypoints/test_serve.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Unit tests for the Omni serve CLI helpers."""
 
 from __future__ import annotations
@@ -76,6 +79,17 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None:
     assert args.cfg_parallel_size == 4
 
 
+def test_serve_parser_accepts_ulysses_a2a_permute() -> None:
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    OmniServeCommand().subparser_init(subparsers)
+
+    args = parser.parse_args(["serve", "fake-model", "--omni", "--ulysses-a2a-permute"])
+
+    assert args.ulysses_a2a_permute is True
+    assert args.get_explicit_kwargs_dict()["ulysses_a2a_permute"] is True
+
+
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",

vllm_omni/config/omni_config.py

@@ -213,6 +213,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False):
     ring_degree: int
     allgather_degree: int
     ulysses_mode: str
+    ulysses_a2a_permute: bool
     cfg_parallel_size: int
     vae_patch_parallel_size: int
     vae_parallel_mode: str
@@ -573,6 +574,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig):
     ring_degree: int = Field(default=1, ge=1)
     allgather_degree: int = Field(default=1, ge=1)
     ulysses_mode: str = "strict"
+    ulysses_a2a_permute: bool = False
     cfg_parallel_size: int = Field(default=1, ge=1)
     vae_patch_parallel_size: int = Field(default=1, ge=1)
     text_encoder_tp_size: int = Field(default=1, ge=1)

vllm_omni/config/stage_config.py

@@ -389,6 +389,7 @@ class StageDeployConfig:
     # Diffusion parallel_config deploy/runtime override fields.
     ulysses_degree: int | None = None
     ulysses_mode: str | None = None
+    ulysses_a2a_permute: bool | None = None
     ring_degree: int | None = None
     allgather_degree: int | None = None
     sequence_parallel_size: int | None = None

vllm_omni/diffusion/attention/parallel/factory.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -47,6 +47,7 @@ def build_parallel_attention_strategy(
     ulysses_degree = getattr(p, "ulysses_degree", 1)
     ring_degree = getattr(p, "ring_degree", 1)
     allgather_degree = getattr(p, "allgather_degree", 1)
+    ulysses_a2a_permute = getattr(p, "ulysses_a2a_permute", False)
 
     try:
         sp_group = get_sp_group()
@@ -83,6 +84,7 @@ def build_parallel_attention_strategy(
             scatter_idx=scatter_idx,
             gather_idx=gather_idx,
             use_sync=use_sync,
+            ulysses_a2a_permute=ulysses_a2a_permute,
         )
 
     # Pure Ring Attention

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
@@ -16,6 +16,11 @@ from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGr
 from vllm_omni.diffusion.forward_context import get_ulysses_mode
 
 
+def _a2a_permute_enabled(enabled: bool, scatter_idx: int, gather_idx: int, world_size: int) -> bool:
+    """Whether the fused permute-free all-to-all applies (strict layout only)."""
+    return enabled and world_size > 1 and scatter_idx == 2 and gather_idx == 1
+
+
 def _ceil_div(n: int, d: int) -> int:
     return (n + d - 1) // d
 
@@ -181,12 +186,23 @@ class UlyssesParallelAttention:
         scatter_idx: int,
         gather_idx: int,
         use_sync: bool,
+        ulysses_a2a_permute: bool = False,
     ) -> None:
         self._sp_group = sp_group
         self._ulysses_pg = sp_group.ulysses_group
         self._scatter_idx = scatter_idx
         self._gather_idx = gather_idx
         self._use_sync = use_sync
+        self._ulysses_a2a_permute = ulysses_a2a_permute
+        if _a2a_permute_enabled(
+            ulysses_a2a_permute,
+            scatter_idx,
+            gather_idx,
+            sp_group.ulysses_world_size,
+        ):
+            from vllm_omni.diffusion.distributed.a2a_permute import ensure_a2a_permute_available
+
+            ensure_a2a_permute_available()
 
     @property
     def enabled(self) -> bool:
@@ -347,9 +363,26 @@ class UlyssesParallelAttention:
                     )
 
             # (bs, seq_len/P, head_cnt, head_size) -> (bs, seq_len, head_cnt/P, head_size)
-            query = SeqAllToAll4D.apply(self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync)
-            key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync)
-            value = SeqAllToAll4D.apply(self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync)
+            if _a2a_permute_enabled(
+                self._ulysses_a2a_permute,
+                self._scatter_idx,
+                self._gather_idx,
+                ulysses_world_size,
+            ):
+                from vllm_omni.diffusion.distributed.a2a_permute import ulysses_qkv_fwd
+
+                group_name = self._ulysses_pg.group_name
+                query = ulysses_qkv_fwd(query, group_name, ulysses_world_size)
+                key = ulysses_qkv_fwd(key, group_name, ulysses_world_size)
+                value = ulysses_qkv_fwd(value, group_name, ulysses_world_size)
+            else:
+                query = SeqAllToAll4D.apply(
+                    self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync
+                )
+                key = SeqAllToAll4D.apply(self._ulysses_pg, key, self._scatter_idx, self._gather_idx, self._use_sync)
+                value = SeqAllToAll4D.apply(
+                    self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync
+                )
             seq_lens = []
             local_seq_len = 0
             orig_head_cnt = 0
@@ -456,6 +489,19 @@ class UlyssesParallelAttention:
                     orig_head_cnt=ctx.orig_head_cnt,
                     use_sync=ctx.use_sync,
                 )
+            elif _a2a_permute_enabled(
+                self._ulysses_a2a_permute,
+                ctx.scatter_idx,
+                ctx.gather_idx,
+                dist.get_world_size(ctx.ulysses_pg),
+            ):
+                from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev
+
+                output_img = ulysses_o_rev(
+                    output_img,
+                    ctx.ulysses_pg.group_name,
+                    dist.get_world_size(ctx.ulysses_pg),
+                )
             else:
                 output_img = SeqAllToAll4D.apply(
                     ctx.ulysses_pg, output_img, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync
@@ -488,4 +534,17 @@ class UlyssesParallelAttention:
                 orig_head_cnt=ctx.orig_head_cnt,
                 use_sync=ctx.use_sync,
             )
+        if _a2a_permute_enabled(
+            self._ulysses_a2a_permute,
+            ctx.scatter_idx,
+            ctx.gather_idx,
+            dist.get_world_size(ctx.ulysses_pg),
+        ):
+            from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev
+
+            return ulysses_o_rev(
+                attn_output,
+                ctx.ulysses_pg.group_name,
+                dist.get_world_size(ctx.ulysses_pg),
+            )
         return SeqAllToAll4D.apply(ctx.ulysses_pg, attn_output, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync)

vllm_omni/diffusion/data.py

@@ -223,6 +223,9 @@ class DiffusionParallelConfig:
       sequence shapes across the ring group.
     """
 
+    ulysses_a2a_permute: bool = False
+    """Use fused permute-free all-to-all for eligible strict Ulysses exchanges."""
+
     cfg_parallel_size: int = 1
     """Number of ranks used to execute guidance passes in parallel."""
 

vllm_omni/diffusion/distributed/a2a_permute.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Fused permute-free Ulysses all-to-all over NCCL symmetric memory.
+
+JIT-compiles the CUDA kernel from pytorch/pytorch#178230 (``all_to_all_permute``)
+and exposes it as two *functional* custom ops:
+
+    ulysses_qkv_fwd(x, group_name, world_size)  # (B, S/p, H, D) -> (B, S, H/p, D)
+    ulysses_o_rev(y, group_name, world_size)     # (B, S, H/p, D) -> (B, S/p, H, D)
+
+These replace the synchronous ``all_to_all_4D`` (permute + NCCL all_to_all_single)
+used by Ulysses SP. All symmetric-memory bookkeeping (a shape-keyed, rendezvoused
+buffer cache + the copy-in) lives *inside* the ops, so from torch.compile's point
+of view each op is an opaque ``input -> output`` function: no graph break, no
+visible buffer mutation, no Python control flow in the traced graph. A
+``register_fake`` provides output metadata so Dynamo keeps the op in-graph.
+"""
+
+from __future__ import annotations
+
+import glob
+import math
+import os
+import sysconfig
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+from torch.distributed.distributed_c10d import _resolve_process_group
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_BUILD_LOCK = threading.Lock()
+_BUILT = False
+_WORKSPACE_LOCK = threading.Lock()
+
+
+@dataclass(slots=True)
+class _SymmWorkspace:
+    allocation: torch.Tensor
+    handle: Any
+    capacity_bytes: int
+    stream_id: int
+
+
+# A single grow-only byte workspace per device/process-group. Shape changes
+# reuse a typed view into the workspace instead of retaining one allocation per
+# historical request shape.
+_SYMM_WORKSPACES: dict[tuple[torch.device, str], _SymmWorkspace] = {}
+
+
+def _ensure_built() -> None:
+    """JIT-compile + load the kernel once (process-wide). Cached by torch."""
+    global _BUILT
+    if _BUILT:
+        return
+    with _BUILD_LOCK:
+        if _BUILT:
+            return
+        from torch.utils.cpp_extension import load
+
+        here = os.path.dirname(__file__)
+        src = os.path.join(here, "csrc", "a2a_permute.cu")
+        site_packages = sysconfig.get_paths()["purelib"]
+        nvidia_root = os.path.join(site_packages, "nvidia")
+        include_paths = glob.glob(os.path.join(nvidia_root, "*", "include"))
+        nccl_libs = glob.glob(os.path.join(nvidia_root, "nccl", "lib", "libnccl.so*"))
+        if not nccl_libs:
+            raise RuntimeError("a2a_permute: could not locate nvidia-nccl libnccl.so")
+        if not any(os.path.isfile(os.path.join(path, "nccl.h")) for path in include_paths):
+            raise RuntimeError("a2a_permute: could not locate nvidia-nccl nccl.h")
+        load(
+            name="vllm_omni_a2a_permute",
+            sources=[src],
+            # PyTorch wheel-only installations keep headers such as cusparse.h
+            # under nvidia/{cu13,cusparse}/include rather than CUDA_HOME.
+            extra_include_paths=list(include_paths),
+            extra_cflags=["-DUSE_NCCL", "-DUSE_C10D_NCCL", "-O3"],
+            extra_cuda_cflags=["-DUSE_NCCL", "-DUSE_C10D_NCCL", "-O3", "--expt-relaxed-constexpr"],
+            extra_ldflags=[nccl_libs[0]],
+            is_python_module=False,
+            verbose=False,
+        )
+        symm_mem.set_backend("NCCL")
+        logger.info("[a2a_permute] JIT kernel built and loaded; symm-mem backend=NCCL")
+        _BUILT = True
+
+
+def ensure_a2a_permute_available() -> None:
+    """Build the extension during worker/model initialization, not a request."""
+    _ensure_built()
+
+
+def _required_nbytes(shape: tuple[int, ...], dtype: torch.dtype) -> int:
+    return math.prod(shape) * torch.empty((), dtype=dtype).element_size()
+
+
+def _get_symm_buffer(
+    symm_shape: tuple[int, ...], dtype: torch.dtype, device: torch.device, group_name: str
+) -> torch.Tensor:
+    """Get a typed view from a grow-only symmetric-memory workspace.
+
+    Allocation/growth and rendezvous are collective. Every rank reaches them in
+    lockstep because sequence parallel execution is symmetric. Steady-state
+    requests at or below the peak capacity take no synchronization path.
+    """
+    device = torch.device(device)
+    key = (device, group_name)
+    required_bytes = _required_nbytes(symm_shape, dtype)
+    stream_id = torch.cuda.current_stream(device).cuda_stream
+    with _WORKSPACE_LOCK:
+        workspace = _SYMM_WORKSPACES.get(key)
+        if workspace is not None and workspace.stream_id != stream_id:
+            raise RuntimeError("a2a_permute workspace must be used from a single CUDA stream")
+        if workspace is None or workspace.capacity_bytes < required_bytes:
+            if torch.cuda.is_current_stream_capturing():
+                capacity = 0 if workspace is None else workspace.capacity_bytes
+                raise RuntimeError(
+                    "a2a_permute: cannot grow the symmetric-memory workspace "
+                    f"from {capacity} to {required_bytes} bytes during CUDA graph capture; "
+                    "warm up the maximum request shape before capture"
+                )
+            if workspace is None:
+                # NCCL symmetric-memory rendezvous requires the communicator to
+                # exist before the first allocation.
+                pg = _resolve_process_group(group_name)
+                warm = torch.ones(1, device=device)
+                dist.all_reduce(warm, group=pg)
+            # Growth is rare and must not release the previous allocation while
+            # an earlier kernel is still consuming it.
+            torch.accelerator.synchronize(device)
+            allocation = symm_mem.empty(required_bytes, dtype=torch.uint8, device=device)
+            handle = symm_mem.rendezvous(allocation, group_name)
+            workspace = _SymmWorkspace(
+                allocation=allocation,
+                handle=handle,
+                capacity_bytes=required_bytes,
+                stream_id=stream_id,
+            )
+            _SYMM_WORKSPACES[key] = workspace
+        return workspace.handle.get_buffer(workspace.handle.rank, symm_shape, dtype)
+
+
+def clear_a2a_permute_workspaces() -> None:
+    """Release cached symmetric-memory workspaces during worker shutdown."""
+    with _WORKSPACE_LOCK:
+        if not _SYMM_WORKSPACES:
+            return
+        devices = {workspace.allocation.device for workspace in _SYMM_WORKSPACES.values()}
+        for device in devices:
+            torch.accelerator.synchronize(device)
+        count = len(_SYMM_WORKSPACES)
+        _SYMM_WORKSPACES.clear()
+    logger.info("[a2a_permute] Released %d symmetric-memory workspace(s)", count)
+
+
+# ---------------------------------------------------------------------------
+# Forward: (B, S_local, H, D) -> (B, S_global, H_local, D)
+#   scatter heads (dim 2), gather sequence (dim 1). Matches all_to_all_4D(x,2,1).
+# ---------------------------------------------------------------------------
+@torch.library.custom_op("vllm_omni_a2a::ulysses_qkv_fwd", mutates_args=(), device_types="cuda")
+def ulysses_qkv_fwd(x: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor:
+    _ensure_built()
+    p = world_size
+    B, s_local, H, D = x.shape
+    Hl = H // p
+    lc = Hl * D
+    rows = B * s_local
+    symm_in = _get_symm_buffer((rows, p, lc), x.dtype, x.device, group_name)
+    # (B, S_local, H, D) -> (rows, p, lc); H is row-major so column block r = heads [r*Hl:(r+1)*Hl]
+    symm_in.copy_(x.reshape(rows, p, lc))
+    out = torch.empty(p, rows, lc, device=x.device, dtype=x.dtype)
+    torch.ops.a2ap.all_to_all_permute(symm_in, out, 1, 0, group_name)
+    # (p, rows, lc) -> (B, S_global, H_local, D), sequence ordered rank-major
+    return out.reshape(p, B, s_local, Hl, D).permute(1, 0, 2, 3, 4).reshape(B, p * s_local, Hl, D).contiguous()
+
+
+@ulysses_qkv_fwd.register_fake
+def _(x: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor:
+    B, s_local, H, D = x.shape
+    return x.new_empty(B, s_local * world_size, H // world_size, D)
+
+
+# ---------------------------------------------------------------------------
+# Reverse: (B, S_global, H_local, D) -> (B, S_local, H, D)
+#   scatter sequence (dim 1), gather heads (dim 2). Matches all_to_all_4D(y,1,2).
+# ---------------------------------------------------------------------------
+@torch.library.custom_op("vllm_omni_a2a::ulysses_o_rev", mutates_args=(), device_types="cuda")
+def ulysses_o_rev(y: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor:
+    _ensure_built()
+    p = world_size
+    B, s_global, Hl, D = y.shape
+    s_local = s_global // p
+    H = Hl * p
+    cols = Hl * D
+    rows = B * s_local
+    symm_in = _get_symm_buffer((p, rows, cols), y.dtype, y.device, group_name)
+    # (B, p*S_local, Hl, D) -> (p, B*S_local, cols): block r = sequence shard destined to rank r
+    symm_in.copy_(y.reshape(B, p, s_local, Hl, D).permute(1, 0, 2, 3, 4).reshape(p, rows, cols))
+    out = torch.empty(rows, p, cols, device=y.device, dtype=y.dtype)
+    torch.ops.a2ap.all_to_all_permute(symm_in, out, 0, 1, group_name)
+    # (rows, p, cols) -> (B, S_local, H, D), heads ordered rank-major
+    return out.reshape(B, s_local, p, Hl, D).reshape(B, s_local, H, D).contiguous()
+
+
+@ulysses_o_rev.register_fake
+def _(y: torch.Tensor, group_name: str, world_size: int) -> torch.Tensor:
+    B, s_global, Hl, D = y.shape
+    return y.new_empty(B, s_global // world_size, Hl * world_size, D)

vllm_omni/diffusion/distributed/csrc/a2a_permute.cu

@@ -0,0 +1,234 @@
+// JIT build of pytorch/pytorch#178230 all_to_all_permute (Ulysses-style fused
+// permute-free all-to-all over NCCL symmetric memory).
+//
+// Kernel is verbatim from the PR; the host entry point is adapted to the
+// NCCLDevCommManager API shipped in torch 2.11 (we obtain the host ncclComm_t
+// via the manager and create our own ncclDevComm with enough LSA barriers,
+// cached per group — the PR's keyed-devcomm manager API is not in 2.11).
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/macros/Macros.h>
+#include <ATen/native/cuda/MemoryAccess.cuh>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/csrc/distributed/c10d/NCCLUtils.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp>
+#include <torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp>
+#include <torch/library.h>
+
+#include <map>
+#include <mutex>
+#include <string>
+
+namespace {
+
+using namespace c10d::symmetric_memory;
+
+#ifndef NCCL_HAS_SYMMEM_DEVICE_SUPPORT
+#error "Build requires NCCL >= 2.28 with symmetric-memory device support"
+#endif
+
+// ---- kernel (verbatim from PR) ----------------------------------------------
+__device__ inline void copy_bytes_vec16_aligned(
+    const char* src, char* dst, size_t nbytes, size_t tid, size_t stride) {
+  const size_t n_vec = nbytes / 16;
+  constexpr int kUnroll = 4;
+  size_t vec_idx = tid;
+  for (; vec_idx + static_cast<size_t>(kUnroll - 1) * stride < n_vec;
+       vec_idx += static_cast<size_t>(kUnroll) * stride) {
+    at::native::memory::Vec<16> chunk[kUnroll];
+#pragma unroll 4
+    for (int u = 0; u < kUnroll; ++u) {
+      const size_t i = vec_idx + static_cast<size_t>(u) * stride;
+      chunk[u] = at::native::memory::ld_vec<16>(src + i * 16);
+    }
+#pragma unroll 4
+    for (int u = 0; u < kUnroll; ++u) {
+      const size_t i = vec_idx + static_cast<size_t>(u) * stride;
+      at::native::memory::st_vec<16>(dst + i * 16, chunk[u]);
+    }
+  }
+  for (; vec_idx < n_vec; vec_idx += stride) {
+    auto v = at::native::memory::ld_vec<16>(src + vec_idx * 16);
+    at::native::memory::st_vec<16>(dst + vec_idx * 16, v);
+  }
+}
+
+constexpr int A2A_MAX_SLOTS = 64;
+constexpr int A2A_MAX_CTAS_PER_SLOT = 128;
+constexpr int A2A_THREADS_PER_CTA = 128;
+constexpr int A2A_MAX_CTA_COUNT = A2A_MAX_SLOTS * A2A_MAX_CTAS_PER_SLOT;
+
+__global__ void all_to_all_lsa_kernel(
+    ncclWindow_t window,
+    size_t base_src_byte_offset,
+    unsigned char* out,
+    int num_rows,
+    size_t src_row_stride_bytes,
+    size_t copy_row_bytes,
+    size_t peer_stride_bytes,
+    size_t dst_row_stride_bytes,
+    ncclDevComm devComm) {
+  const int peer_idx = blockIdx.x;
+  const int local_block = blockIdx.y;
+  const ncclCoopCta coop{};
+
+  ncclLsaBarrierSession<ncclCoopCta> bar{
+      coop, devComm, ncclTeamLsa(devComm), devComm.lsaBarrier,
+      blockIdx.x * gridDim.y + blockIdx.y};
+  bar.sync(coop, cuda::memory_order_acquire);
+
+  unsigned char* dst_peer_base =
+      out + static_cast<size_t>(peer_idx) * peer_stride_bytes;
+  CUDA_KERNEL_ASSERT((base_src_byte_offset & 15) == 0);
+  CUDA_KERNEL_ASSERT((reinterpret_cast<uintptr_t>(dst_peer_base) & 15) == 0);
+
+  for (int i = local_block; i < num_rows; i += gridDim.y) {
+    const size_t row_byte_offset =
+        base_src_byte_offset + static_cast<size_t>(i) * src_row_stride_bytes;
+    const void* src_row = ncclGetLsaPointer(window, row_byte_offset, peer_idx);
+    unsigned char* dst_row =
+        dst_peer_base + static_cast<size_t>(i) * dst_row_stride_bytes;
+    copy_bytes_vec16_aligned(
+        reinterpret_cast<const char*>(src_row),
+        reinterpret_cast<char*>(dst_row),
+        copy_row_bytes, coop.thread_rank(), coop.size());
+  }
+  bar.sync(coop, cuda::memory_order_release);
+}
+
+// ---- own devcomm cache (replaces PR's keyed manager API) --------------------
+ncclDevComm get_or_create_devcomm(ncclComm_t comm, const std::string& group_name) {
+  static std::mutex mu;
+  static std::map<std::string, ncclDevComm> cache;
+  std::lock_guard<std::mutex> lk(mu);
+  auto it = cache.find(group_name);
+  if (it != cache.end()) return it->second;
+#ifdef NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER
+  ncclDevCommRequirements reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER;
+#else
+  ncclDevCommRequirements reqs;
+  memset(&reqs, 0, sizeof(ncclDevCommRequirements));
+#endif
+  reqs.lsaBarrierCount = A2A_MAX_CTA_COUNT;
+  ncclDevComm dc;
+  ncclResult_t res = ncclDevCommCreate(comm, &reqs, &dc);
+  TORCH_CHECK(res == ncclSuccess, "ncclDevCommCreate (a2a_permute) failed: ", ncclGetErrorString(res));
+  cache.emplace(group_name, dc);
+  return dc;
+}
+
+// ---- host entry -------------------------------------------------------------
+void all_to_all_permute(
+    const at::Tensor& input,
+    at::Tensor& out,
+    int64_t scatter_dim,
+    int64_t gather_dim,
+    std::string group_name) {
+  TORCH_CHECK(input.stride(-1) == 1, "a2a_permute: innermost dim must be contiguous");
+  const bool col_scatter = (scatter_dim == 1 && gather_dim == 0);
+  const bool row_scatter = (scatter_dim == 0 && gather_dim == 1);
+  TORCH_CHECK(col_scatter || row_scatter,
+      "a2a_permute: unsupported (scatter_dim, gather_dim)=(", scatter_dim, ",", gather_dim, ")");
+
+  auto symm_mem = c10d::symmetric_memory::rendezvous(input, group_name);
+  TORCH_CHECK(symm_mem != nullptr, "a2a_permute: input must be NCCL symmetric memory");
+  auto* nccl_hdl = dynamic_cast<NCCLSymmetricMemory*>(symm_mem.get());
+  TORCH_CHECK(nccl_hdl != nullptr, "a2a_permute: requires NCCL symmetric memory backend");
+
+  c10::cuda::CUDAGuard guard(input.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+  auto device = input.device();
+
+  auto& manager = NCCLDevCommManager::get(device);
+  ncclComm_t comm = manager.get_comm(group_name);
+  ncclDevComm devcomm = get_or_create_devcomm(comm, group_name);
+
+  const int my_rank = nccl_hdl->get_rank();
+  const int p = nccl_hdl->get_world_size();
+  TORCH_CHECK(p <= A2A_MAX_SLOTS, "a2a_permute: group size ", p, " exceeds ", A2A_MAX_SLOTS);
+  TORCH_CHECK(out.is_contiguous(), "a2a_permute: out must be contiguous");
+  TORCH_CHECK(out.scalar_type() == input.scalar_type(), "a2a_permute: dtype mismatch");
+
+  auto window = nccl_hdl->get_window();
+  TORCH_CHECK(window != nullptr, "a2a_permute: NCCL window is null");
+
+  const size_t window_base_offset = nccl_hdl->get_offset();
+  const int64_t esize = input.element_size();
+  const size_t tensor_leading_offset =
+      window_base_offset + static_cast<size_t>(input.storage_offset()) * static_cast<size_t>(esize);
+  TORCH_CHECK(tensor_leading_offset % 16 == 0, "a2a_permute: tensor byte offset must be 16B aligned");
+  TORCH_CHECK(reinterpret_cast<uintptr_t>(input.data_ptr()) % 16 == 0, "a2a_permute: input ptr must be 16B aligned");
+  TORCH_CHECK(reinterpret_cast<uintptr_t>(out.data_ptr()) % 16 == 0, "a2a_permute: out ptr must be 16B aligned");
+
+  const int unroll = 4 * 16 / static_cast<int>(esize);
+  const int elems_per_cta = A2A_THREADS_PER_CTA * unroll;
+  int ctas_per_slot = 1;
+
+  if (col_scatter) {
+    const int rows = static_cast<int>(input.size(0));
+    int64_t total_cols = 0; int local_cols = 0;
+    if (input.dim() == 2) {
+      total_cols = input.size(1);
+      TORCH_CHECK(total_cols % p == 0, "a2a_permute: cols must divide group size");
+      local_cols = static_cast<int>(total_cols / p);
+    } else {
+      TORCH_CHECK(input.dim() == 3, "a2a_permute: input must be 2-D or 3-D");
+      TORCH_CHECK(input.size(1) == p, "a2a_permute: 3-D input dim1 must equal group size");
+      const int64_t lc = input.size(2);
+      total_cols = static_cast<int64_t>(p) * lc;
+      TORCH_CHECK(input.stride(1) == lc && input.stride(0) == total_cols,
+          "a2a_permute: 3-D input must be row-major contiguous");
+      local_cols = static_cast<int>(lc);
+    }
+    const bool ok3 = out.dim() == 3 && out.size(0) == p && out.size(1) == rows && out.size(2) == local_cols;
+    const bool ok2 = out.dim() == 2 && out.size(0) == (int64_t)p * rows && out.size(1) == local_cols;
+    TORCH_CHECK(ok3 || ok2, "a2a_permute: bad out shape for (1,0)");
+    const size_t row_bytes = (size_t)local_cols * (size_t)esize;
+    TORCH_CHECK(row_bytes % 16 == 0, "a2a_permute: local_cols*esize must be 16B-divisible");
+    ctas_per_slot = std::max(1, std::min((rows * local_cols + elems_per_cta - 1) / elems_per_cta, A2A_MAX_CTAS_PER_SLOT));
+    const size_t esz = (size_t)esize;
+    const size_t base_src = tensor_leading_offset + (size_t)my_rank * (size_t)local_cols * esz;
+    all_to_all_lsa_kernel<<<dim3(p, ctas_per_slot), A2A_THREADS_PER_CTA, 0, stream>>>(
+        window, base_src, reinterpret_cast<unsigned char*>(out.data_ptr()), rows,
+        (size_t)total_cols * esz, row_bytes, (size_t)rows * local_cols * esz, (size_t)local_cols * esz, devcomm);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  } else {
+    int local_rows = 0; int cols = 0;
+    if (input.dim() == 2) {
+      const int64_t total_rows = input.size(0);
+      cols = static_cast<int>(input.size(1));
+      TORCH_CHECK(total_rows % p == 0, "a2a_permute: rows must divide group size");
+      local_rows = static_cast<int>(total_rows / p);
+    } else {
+      TORCH_CHECK(input.dim() == 3, "a2a_permute: input must be 2-D or 3-D");
+      TORCH_CHECK(input.size(0) == p, "a2a_permute: 3-D input dim0 must equal group size");
+      local_rows = static_cast<int>(input.size(1));
+      const int64_t c = input.size(2); cols = static_cast<int>(c);
+      const int64_t s01 = (int64_t)local_rows * c;
+      TORCH_CHECK(input.stride(1) == c && input.stride(0) == s01, "a2a_permute: 3-D input must be row-major contiguous");
+    }
+    const bool ok3 = out.dim() == 3 && out.size(0) == local_rows && out.size(1) == p && out.size(2) == cols;
+    const bool ok2 = out.dim() == 2 && out.size(0) == local_rows && out.size(1) == (int64_t)p * cols;
+    TORCH_CHECK(ok3 || ok2, "a2a_permute: bad out shape for (0,1)");
+    const size_t row_bytes = (size_t)cols * (size_t)esize;
+    TORCH_CHECK(row_bytes % 16 == 0, "a2a_permute: cols*esize must be 16B-divisible");
+    ctas_per_slot = std::max(1, std::min((local_rows * cols + elems_per_cta - 1) / elems_per_cta, A2A_MAX_CTAS_PER_SLOT));
+    const size_t esz = (size_t)esize; const size_t cu = (size_t)cols;
+    const size_t base_src = tensor_leading_offset + (size_t)(my_rank * local_rows) * cu * esz;
+    all_to_all_lsa_kernel<<<dim3(p, ctas_per_slot), A2A_THREADS_PER_CTA, 0, stream>>>(
+        window, base_src, reinterpret_cast<unsigned char*>(out.data_ptr()), local_rows,
+        cu * esz, row_bytes, cu * esz, (size_t)p * cu * esz, devcomm);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  }
+}
+
+} // namespace
+
+TORCH_LIBRARY_FRAGMENT(a2ap, m) {
+  m.def("all_to_all_permute(Tensor input, Tensor(a!) out, int scatter_dim, int gather_dim, str group_name) -> ()");
+}
+TORCH_LIBRARY_IMPL(a2ap, CUDA, m) {
+  m.impl("all_to_all_permute", TORCH_FN(all_to_all_permute));
+}

vllm_omni/diffusion/distributed/csrc/nccl_extension.hpp

@@ -0,0 +1,44 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include <c10/macros/Macros.h>
+
+namespace c10d::nccl_extension {
+
+TORCH_API bool is_nccl_symmem_available();
+
+TORCH_API void nccl_put(at::Tensor& tensor, const int64_t peer);
+
+TORCH_API void nccl_get(at::Tensor& tensor, const int64_t peer);
+
+TORCH_API void nccl_wait_for_signal(at::Tensor& sigpad, int64_t signal);
+
+TORCH_API void nccl_put_with_signal(
+    at::Tensor& tensor,
+    int64_t signal,
+    int64_t peer);
+
+// Simultaneously reduce N blocks of a 2-D input tensor from a shared symmetric
+// memory buffer, routing each to a specific destination rank. Blocks are
+// described by inclusive-prefix-sum offsets along `dim` (0 or 1); all blocks
+// must have equal size.
+TORCH_API void nccl_reduce_scatter_offset(
+    const at::Tensor& input,
+    at::TensorList out,
+    const std::string& group_name,
+    int64_t dim,
+    std::optional<at::IntArrayRef> offsets,
+    std::optional<at::IntArrayRef> dst_ranks,
+    const std::string& red_op);
+
+// Permute-free all-to-all: scatter shards along `scatter_dim`, gather peer
+// chunks along `gather_dim`. Supported: (scatter_dim=1, gather_dim=0) -> [rows,
+// p*lc] to [p, rows, lc]; (scatter_dim=0, gather_dim=1) -> [p*lr, cols] to [lr,
+// p, cols].
+TORCH_API void nccl_all_to_all_permute(
+    const at::Tensor& input,
+    at::Tensor& out,
+    int64_t scatter_dim,
+    int64_t gather_dim,
+    const std::string& group_name);
+} // namespace c10d::nccl_extension

vllm_omni/diffusion/envs.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2024 xDiT team.
 # Adapted from
 # https://github.com/xdit-project/xDiT/blob/main/xfuser/envs.py

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -13,6 +13,7 @@ import multiprocessing as mp
 import os
 import queue
 import signal
+import sys
 import threading
 import traceback
 import uuid
@@ -964,7 +965,14 @@ class DiffusionWorker:
                     if mgr is not None:
                         mgr.shutdown_prefetch()
         finally:
-            destroy_distributed_env()
+            try:
+                a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
+                if a2a_permute is not None:
+                    a2a_permute.clear_a2a_permute_workspaces()
+            except Exception:
+                logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
+            finally:
+                destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/arg_utils.py

@@ -548,6 +548,7 @@ class OrchestratorArgs:
     diffusers_call_kwargs: str = "{}"
     ulysses_degree: int | None = None
     ulysses_mode: str = "strict"
+    ulysses_a2a_permute: bool | None = None
     ring_degree: int | None = None
     allgather_degree: int | None = None
     diffusion_quantization_config: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1051,6 +1051,7 @@ class AsyncOmniEngine:
             ring_degree = normalized_kwargs.get("ring_degree") or 1
             allgather_degree = normalized_kwargs.get("allgather_degree") or 1
             ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict"
+            ulysses_a2a_permute = bool(normalized_kwargs.get("ulysses_a2a_permute", False))
             sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size")
             pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1
             data_parallel_size = normalized_kwargs.get("data_parallel_size")
@@ -1077,6 +1078,7 @@ class AsyncOmniEngine:
                 ring_degree=ring_degree,
                 allgather_degree=allgather_degree,
                 ulysses_mode=ulysses_mode,
+                ulysses_a2a_permute=ulysses_a2a_permute,
                 cfg_parallel_size=cfg_parallel_size,
                 vae_patch_parallel_size=vae_patch_parallel_size,
                 vae_parallel_mode=vae_parallel_mode,

vllm_omni/entrypoints/cli/serve.py

@@ -540,6 +540,15 @@ class OmniServeCommand(CLISubcommand):
             "'strict' keeps the original divisibility requirements; "
             "'advanced_uaa' enables the experimental UAA path for uneven sequence/head shapes.",
         )
+        omni_config_group.add_argument(
+            "--ulysses-a2a-permute",
+            action=argparse.BooleanOptionalAction,
+            default=None,
+            help=(
+                "Enable fused permute-free Ulysses all-to-all over NCCL symmetric memory. "
+                "Only strict Ulysses layouts are eligible. Defaults to disabled."
+            ),
+        )
         omni_config_group.add_argument(
             "--ring",
             "--ring-degree",