Skip to content

Repository files navigation

CoCoDiff — artifact for the communication library

Artifact for "CoCoDiff: Optimizing Collective Communications for Distributed Diffusion Transformer Inference Under Ulysses Sequence Parallelism" (Bin Ma, Xingjian Ding, Tekin Bicer, Pengfei Su, Dong Li).

The artifact is split across two repositories, and both are required:

Repository Contents
Library binkma-v/sp-aurora (this one) TAPA, V-First and V-Major themselves
Engine binkma-v/xDit-intel the xfuser inference engine, model drivers, jobs

This repository is sp_aurora, the sequence-parallel communication library for Intel GPUs. The engine imports it at run time; if the import fails, the engine silently falls back to a flat all-to-all and every configuration measures the same latency, so verify it before collecting data.

sp_aurora began as a fork of ring-flash-attention; its original library README is preserved as README_sp_aurora_library.md.

The three mechanisms

Mechanism Where
TAPA — splits each Ulysses all-to-all into an intra-GPU tile-to-tile phase (peer = rank XOR 1) and an inter-GPU phase run as two parallel 6-rank Xe Link rings, split by tile parity sp_aurora/comm/hierarchical_alltoall.py
V-First scheduling — launches V's phase-1 exchange on a background thread so it overlaps Q/K normalisation and RoPE sp_aurora/comm/hierarchical_alltoall.py (hierarchical_phase1_async)
V-Major selective communication — an L1 metric over cached v selects the active token projections, their indices are shared by all-gather, and only that subset crosses the inter-GPU links sp_aurora/comm/phase2_cache.py

Aurora's hierarchy is 185 GB/s between the two tiles of a GPU, 15 GB/s between GPUs over Xe Link, and 25 GB/s per Slingshot-11 NIC between nodes — a 12× spread that a flat all-to-all ignores.

Installation

Requires Aurora (or another Intel Data Center GPU Max system) with the oneAPI and frameworks modules.

git clone https://github.com/binkma-v/sp-aurora.git sp_aurora
cd sp_aurora

module use /soft/modulefiles/
module load frameworks/2025.3.1
module load oneapi/release/2025.3.1

# --system-site-packages matters: torch, IPEX and oneCCL come from the module
python -m venv --system-site-packages $HOME/venv/diff
source $HOME/venv/diff/bin/activate

bash sp_aurora/sycl/build.sh
BUILD_SYCL=1 pip install -e .

The SYCL kernels are optional. sp_aurora/intel_flash_attn_sycl.py wraps their import in try/except and falls back to the oneDNN and PyTorch paths, so the package works without building them — you will simply see a warning.

Configuration

Everything is selected through environment variables; see the engine repository's README for the full table. The essentials:

XDIT_USE_HIERARCHICAL_ALLTOALL=true    # TAPA
XDIT_ENABLE_VFIRST_OVERLAP=true        # V-First
XDIT_PHASE2_CACHE_ENABLED=true         # V-Major
XDIT_PHASE2_CACHE_MAJOR=V
XDIT_PHASE2_CACHE_RATIO=0.0-1.0        # or a fixed value such as 0.9
XDIT_NUM_INFERENCE_STEPS=50            # required for the range form

benchmark/ holds the standalone all-to-all comparisons.

License

Apache 2.0 — see LICENSE. This project derives from ring-flash-attention and bundles components originally under the MIT License; those notices are preserved in NOTICE.

About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages