Context Parallelism

Sequence/Context Parallelism

cache-dit is compatible with context parallelism. Currently, we support the use of Hybrid Cache + Context Parallelism scheme (via CACHE_DIT parallelism backend) in cache-dit. Users can use Context Parallelism to further accelerate the speed of inference. Legacy configs that still use NATIVE_DIFFUSER are still accepted and will be automatically converted to CACHE_DIT. For more details, please refer to 📚examples. Currently, cache-dit supported context parallelism for FLUX.1, 🔥FLUX.2, Qwen-Image, Qwen-Image-Lightning, LTXVideo, Wan 2.1, Wan 2.2, HunyuanImage-2.1, HunyuanVideo, CogVideoX 1.0, CogVideoX 1.5, CogView 3/4 and VisualCloze, etc. cache-dit will support more models in the future.

from cache_dit import ParallelismConfig

cache_dit.enable_cache(
  pipe_or_adapter, 
  cache_config=DBCacheConfig(...),
  # Set ulysses_size > 1 to enable ulysses style context parallelism.
  parallelism_config=ParallelismConfig(ulysses_size=2),
)

L20x1	Ulysses-2	Ulysses-4	+ compile
FLUX, 23.56s	13.80s	8.28s	7.27s

UAA: Ulysses Anything Attention

Any Sequence Length: We have implemented the UAA: Ulysses Anything Attention: An Ulysses Attention that supports arbitrary sequence length with zero padding and nearly zero theoretical communication overhead. The default Ulysses Attention requires that the sequence len of hidden states must be divisible by the number of devices. This imposes significant limitations on the practical application of Ulysses.

from cache_dit import ParallelismConfig

cache_dit.enable_cache(
  pipe_or_adapter, 
  cache_config=DBCacheConfig(...),
  # Set `ulysses_anything` as True to enable UAA
  parallelism_config=ParallelismConfig(
    ulysses_size=2,
    ulysses_anything=True,
  ),
)

For example, in the T2I and I2V tasks, the length of prompts input by users is often variable, and it is difficult to ensure that this length is divisible by the number of devices. To address this issue, we have developed a padding-free Ulysses Attention (UAA) for arbitrary sequence length, which enhances the versatility of Ulysses.

dist.init_process_group(backend="cpu:gloo,cuda:nccl")

Compared to Ulysses Attention, in UAA, we have only added an extra all-gather op for scalar types to gather the seq_len value of each rank. To avoid multiple forced CUDA sync caused by H2D and D2H transfers, please add the gloo backend in init_process_group. This will significantly reduce communication latency.

Any Sequence Length
U*: Ulysses Attention, UAA: Ulysses Anything Attenton, UAA*: UAA + Gloo, Device: NVIDIA L20
FLUX.1-Dev w/o CPU Offload, 28 steps; Qwen-Image w/ CPU Offload, 50 steps; Gloo: Extra All Gather w/ Gloo

CP2 w/ U*	CP2 w/ UAA*	CP2 w/ UAA	L20x1	CP2 w/ UAA*	CP2 w/ U*	L20x1	CP2 w/ UAA*
FLUX, 13.87s	🎉13.88s	14.75s	23.25s	🎉13.75s	Qwen, 132s	181s	🎉133s
1024x1024	1024x1024	1024x1024	1008x1008	1008x1008	1312x1312	1328x1328	1328x1328
✔️U* ✔️UAA	✔️U* ✔️UAA	✔️U* ✔️UAA	NO CP	❌U* ✔️UAA	✔️U* ✔️UAA	NO CP	❌U* ✔️UAA

Any Head Num: By the way, Ulysses Attention and UAA in cache-dit now support arbitrary numbers of heads via additional padding and unpadding operations implemented before and after all-to-all. The overhead incurred by these extra padding and unpadding steps can be partially hidden through asynchronous communication. This support for arbitrary head counts is automatically activated whenever the number of heads is not divisible by the world size. For Example:

Any Head Num
Ulysses: Ulysses Attention, FP8 Ulysses: Ulysses w/ FP8 All2All, Device: NVIDIA L20
🔥Z-Image (Head=30, ❌CAN NOT divisible by 4), 1024x1024, 9 steps.

Ulysses 2, L20	Ulysses 4	FP8 Ulysses 4	+ Cache	+ FP8 DQ
1024x1024, 3.19s	1024x1024, 1.98s	1024x1024, 1.89s	1024x1024, 1.63s	1024x1024, 1.23s

We have also implemented a padding-free version that support any head num. Please be informed that this solution cannot be used when seq len is not divisible by world size. Users can enable this feature through environment variables:

export CACHE_DIT_UNEVEN_HEADS_COMM_NO_PAD=1 # NOT WORK if seq len is also not divisible by world size

Important: Please note that Ulysses Anything Attention (UAA) is currently an experimental feature. It has not undergone large-scale testing, and may introduce a slight performance degradation while the cpu:gloo commucation backend is not available.

Async Ulysses QKV Projection

Inspired by ByteDance-Seed/VeOmni: Async Ulysses CP, we have also added support for Async Ulysses QKV Projection for certain models in cache-dit. This enables partial overlap of communication and computation, which can further enhance the performance of Ulysses style Context Parallelism. Currently, only the 🔥FLUX.1, 🔥Qwen-Image, 🔥Z-Image, 🔥Ovis-Image, 🔥FLUX.2 and 🔥LongCat-Image models are supported, and more models will be added in the future—stay tuned!

from cache_dit import ParallelismConfig

cache_dit.enable_cache(
  pipe_or_adapter, 
  cache_config=DBCacheConfig(...),
  # Set `ulysses_async` as True to enable Async Ulysses QKV Projection.
  parallelism_config=ParallelismConfig(
    ulysses_size=2,
    ulysses_async=True,
  ),
)

Ulysses: Standard Ulysses Attention, Async Ulysses: Ulysses Attention with Async QKV Projection

L20x2 w/ Ulysses	w/ Async Ulysses	w/ Ulysses + compile	w/ Async Ulysses + compile
FLUX.1, 13.87s	🎉13.20s	12.21s	🎉11.97s

Async FP8 Ulysses Attention

cache-dit has implemented Async FP8 Ulysses Attention for 🔥all supported DiTs. This optimization reduces communication latency while preserving high precision. Users can enable this feature by setting ulysses_float8=True. To maintain higher precision during softmax computation—where Softmax(Q@K^T) is sensitive to numerical instability—we currently retain K in FP16/BF16 format. Float8-optimized all_to_all communication is therefore only applied to Q, V, and O.

from cache_dit import ParallelismConfig

cache_dit.enable_cache(
  pipe_or_adapter, 
  cache_config=DBCacheConfig(...),
  # Set `ulysses_float8` as True to enable Async FP8 Ulysses Attention
  parallelism_config=ParallelismConfig(
    ulysses_size=2,
    ulysses_float8=True,
  ),
)

L20x2 w/ Ulysses	w/ Ulysses FP8	w/ Ulysses + compile	w/ Ulysses FP8 + compile
FLUX.1, 13.87s	🎉13.36s	12.21s	🎉11.54s

Ring Attention with Batched P2P

Currently, cache-dit support 2 ring_rotate_method, namely, allgather and p2p. allgather: Use allgather to gather the key and value tensors. p2p: Use batch_isend_irecv ops to rotate the key and value tensors (default). This method is more efficient due to th better overlap of communication and computation.

from cache_dit import ParallelismConfig

cache_dit.enable_cache(
  pipe_or_adapter, 
  cache_config=DBCacheConfig(...),
  # Set `ring_rotate_method` as 'p2p' to enable the faster ring attention implementation
  parallelism_config=ParallelismConfig(
    ring_size=2,
    ring_rotate_method='p2p', # default is 'p2p'
  ),
)

The FLUX.1-dev benchmark for Ring w/ allgather (AG), Ring w/ batched p2p (p2p) and Ulysses listed as following:

Ring-2 w/ AG	Ring-2 w/ p2p	Ulysses-2	Ring-4 w/ AG	Ring-4 w/ p2p	Ulysses-4
14.01s	13.43s	13.87s	9.82s	9.16s	8.22s

USP: Unified Sequence Parallelism

Unified Sequence Parallelism combines Ring Attention and Ulysses Attention into a single approach for efficient long-sequence processing. It applies Ulysses’s all-to-all communication first to redistribute heads and sequence tokens, then uses Ring Attention to process the redistributed data, and finally reverses the all-to-all to restore the original layout.

Ulysses Attention efficiently parallelizes across attention heads. Ring Attention handles very long sequences with minimal memory overhead. Together, they enable 2D parallelization across both heads and sequence dimensions

from cache_dit import ParallelismConfig

cache_dit.enable_cache(
  pipe_or_adapter, 
  cache_config=DBCacheConfig(...),
  # Set both `ulysses_size` and `ring_size` greater than 1
  parallelism_config=ParallelismConfig(
    ulysses_size=2, ring_size=2,
  ),
)

From the below table (FLUX.1-dev), it’s clear that Ulysses provides better throughput, but the number of devices it can use remains limited to the number of attention heads, a limitation that is solved by unified attention.

Ring-4 w/ allgather	Ring-4 w/ batched p2p	Ulysses-4	USP-2-2
9.82s	9.16s	8.22s	8.55s