CoCoDiff — artifact for the inference engine

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).

CoCoDiff restructures the all-to-all collectives that dominate distributed DiT inference on the Aurora supercomputer. It has three mechanisms:

• TAPA (tile-aware parallel all-to-all) splits each Ulysses all-to-all into a high-bandwidth intra-GPU phase and an inter-GPU phase organised as two parallel Xe Link rings, matching Aurora's 12× bandwidth hierarchy.
• V-First scheduling projects V first and runs its intra-GPU exchange on a background thread, hiding it behind Q/K normalisation and RoPE.
• V-Major selective communication sends only the token projections that changed appreciably since the previous denoising step.

Across four DiT models and 1–8 Aurora nodes (up to 96 Intel GPU tiles), CoCoDiff averages a 3.6× end-to-end speedup over the flat all-to-all baseline, peaking at 8.4×.

Repository layout

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

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

This repository is a fork of xDiT whose xfuser package is extended to drive the four evaluated models through CoCoDiff. The upstream project's README is preserved as README_xDiT_upstream.md. sp_aurora is imported at run time, not vendored.

xfuser/            the engine; CoCoDiff hooks live in
                     model_executor/layers/attention_processor.py
                     core/long_ctx_attention/hybrid/attn_layer.py
tests/             per-model latency drivers, inpainting evaluators, analysers
jobs/              PBS scripts for 1-8 nodes; see jobs/README.md
benchmarks/        bench_pure_alltoall.py, the flat-vs-TAPA micro-benchmark
tools/plots/       figure-generating scripts

Setup

Requires an Aurora allocation. Build and install sp-aurora first, then:

git clone https://github.com/binkma-v/xDit-intel.git xDiT
cd xDiT

module use /soft/modulefiles/
module load frameworks/2025.3.1
module load oneapi/release/2025.3.1
source $VENV/bin/activate

pip install -e ".[diffusers,test]"
export PYTHONPATH=$PWD:/path/to/sp_aurora:$PYTHONPATH

# Verify: without this import the engine silently falls back to flat all-to-all
python -c "import sp_aurora; print(sp_aurora.__file__)"

Running

The three configurations are selected by environment variables, not by a command-line flag:

Variable	Flat	TAPA	CoCoDiff
XDIT_USE_HIERARCHICAL_ALLTOALL	false	true	true
XDIT_ENABLE_VFIRST_OVERLAP	false	false	true
XDIT_PHASE2_CACHE_ENABLED	false	false	true
XDIT_PHASE2_CACHE_MAJOR	—	—	V
XDIT_PHASE2_CACHE_RATIO	—	—	0.0-1.0
XDIT_PHASE2_CACHE_WARMUP	—	—	5
XDIT_PHASE2_CACHE_FORCE_FULL_SYNC	—	—	10
XDIT_NUM_INFERENCE_STEPS	50	50	50

The job scripts call these Flat, Hier and HFV. XDIT_PHASE2_CACHE_RATIO takes either a fixed value (0.9) or a range (0.0-1.0) for the time-varying schedule; the range form needs XDIT_NUM_INFERENCE_STEPS set, or the ratio stays at the lower bound.

jobs/README.md maps each figure and table in the paper to the script that produces it, and lists the variables to export before qsub.

License

Apache 2.0, inherited from upstream xDiT — see LICENSE.txt.