diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 1a9efcbed..7ed879d7b 100644 --- a/tests/config/test_omni_config.py +++ b/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) diff --git a/tests/diffusion/distributed/test_a2a_permute.py b/tests/diffusion/distributed/test_a2a_permute.py new file mode 100644 index 000000000..a21e29217 --- /dev/null +++ b/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] diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 3864d1d67..71a6e00a5 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/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] diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 6aa5d6855..e759e0f15 100644 --- a/tests/entrypoints/test_serve.py +++ b/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", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 3043a1309..58dc55fb6 100644 --- a/vllm_omni/config/omni_config.py +++ b/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) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 661c0bf6a..72641a1bc 100644 --- a/vllm_omni/config/stage_config.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py index 76f7a78da..5b644ca45 100644 --- a/vllm_omni/diffusion/attention/parallel/factory.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index fc903d050..69f6a810e 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/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) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index f29c59c34..582346828 100644 --- a/vllm_omni/diffusion/data.py +++ b/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.""" diff --git a/vllm_omni/diffusion/distributed/a2a_permute.py b/vllm_omni/diffusion/distributed/a2a_permute.py new file mode 100644 index 000000000..b4e951d34 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu b/vllm_omni/diffusion/distributed/csrc/a2a_permute.cu new file mode 100644 index 000000000..ce5dba770 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +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(kUnroll - 1) * stride < n_vec; + vec_idx += static_cast(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(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(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 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(peer_idx) * peer_stride_bytes; + CUDA_KERNEL_ASSERT((base_src_byte_offset & 15) == 0); + CUDA_KERNEL_ASSERT((reinterpret_cast(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(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(i) * dst_row_stride_bytes; + copy_bytes_vec16_aligned( + reinterpret_cast(src_row), + reinterpret_cast(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 cache; + std::lock_guard 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(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(input.storage_offset()) * static_cast(esize); + TORCH_CHECK(tensor_leading_offset % 16 == 0, "a2a_permute: tensor byte offset must be 16B aligned"); + TORCH_CHECK(reinterpret_cast(input.data_ptr()) % 16 == 0, "a2a_permute: input ptr must be 16B aligned"); + TORCH_CHECK(reinterpret_cast(out.data_ptr()) % 16 == 0, "a2a_permute: out ptr must be 16B aligned"); + + const int unroll = 4 * 16 / static_cast(esize); + const int elems_per_cta = A2A_THREADS_PER_CTA * unroll; + int ctas_per_slot = 1; + + if (col_scatter) { + const int rows = static_cast(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(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(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(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<<>>( + window, base_src, reinterpret_cast(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(input.size(1)); + TORCH_CHECK(total_rows % p == 0, "a2a_permute: rows must divide group size"); + local_rows = static_cast(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(input.size(1)); + const int64_t c = input.size(2); cols = static_cast(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<<>>( + window, base_src, reinterpret_cast(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)); +} diff --git a/vllm_omni/diffusion/distributed/csrc/nccl_extension.hpp b/vllm_omni/diffusion/distributed/csrc/nccl_extension.hpp new file mode 100644 index 000000000..e9fe0e19c --- /dev/null +++ b/vllm_omni/diffusion/distributed/csrc/nccl_extension.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include +#include + +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 offsets, + std::optional 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 diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index 8ed8fbefd..32e006f7e 100644 --- a/vllm_omni/diffusion/envs.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index de8111c40..6499beee5 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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: diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 12dbacefe..5fe5e9269 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 4abccb6fe..7e0d9f629 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index b3a1e91af..f4c890816 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/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",