[Feature]: Support Ulysses Sequence Parallelism for Diffusion Models

[wtomin]

🚀 The feature, motivation and pitch

For long sequence (image & video) generation, SP (Sequence Parallelism) is widely adopted to accelerate the inference speed. Especially, Ulysses-SP and Ring-Attention are two popular SP algorithms.

First, I would like to start implementing Ulysses-SP on Qwen-Image. Currently working on

• [Diffusion]: Diffusion Ulysses-Sequence-Parallelism support #189

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[SamitHuang]

I think we can support TP and CP for diffusion models at first by re-using the parallelism interfaces in vllm. Then we can verify whether the CP interfaces like sequence_parallel_chunk and tensor_model_parallel_all_gather are suitable for implementing Ulysses SP (sequence split + all-to-all)

[wtomin]

I think we can support TP and CP for diffusion models at first by re-using the parallelism interfaces in vllm. Then we can verify whether the CP interfaces like sequence_parallel_chunk and tensor_model_parallel_all_gather are suitable for implementing Ulysses SP (sequence split + all-to-all)

I will take some time to analyse it. Thanks for your suggestion!

[wtomin]

I agree that we can consider implementing TP+CP first, using the available interface in vllm.

As for Ulysses-SP, I am worried that the existing CP interfaces are entangled with TP. For example:

def sequence_parallel_chunk(x: torch.Tensor) -> torch.Tensor:
    return torch.ops.vllm.sequence_parallel_chunk_impl(x)


def sequence_parallel_chunk_impl(x: torch.Tensor) -> torch.Tensor:
    tp_size = get_tensor_model_parallel_world_size()
    tp_rank = get_tensor_model_parallel_rank()

    # all_gather needs the sequence length to be divisible by tp_size
    seq_len = x.size(0)
    remainder = seq_len % tp_size
    if remainder != 0:
        pad_len = tp_size - remainder
        y = nn.functional.pad(x, (0, 0, 0, pad_len))
    else:
        y = x

    chunk = y.shape[0] // tp_size
    start = tp_rank * chunk
    return torch.narrow(y, 0, start, chunk)

TP and CP share the same communication group. While for Ulysses-SP and Ring-Attention, TP is not a must.

[mxuax]

This existing SP logic is designed to reduce activation memory in MLP or MoE layers by splitting the sequence dimension across TP ranks.
e.g., llama4 moe block, the following forward function called sequence_parallel_chunk and tensor_model_parallel_all_gather.

    def forward(self, hidden_states):
        num_tokens = hidden_states.shape[0]
        if self.is_sequence_parallel:
            hidden_states = sequence_parallel_chunk(hidden_states)

        router_logits, _ = self.router(hidden_states)

        shared_out, routed_out = self.experts(
            hidden_states=hidden_states,
            router_logits=router_logits,
        )
        experts_out = routed_out + shared_out

        if self.is_sequence_parallel:
            experts_out = tensor_model_parallel_all_gather(experts_out, 0)
            experts_out = experts_out[:num_tokens]
        elif self.tp_size > 1:
            experts_out = self.experts.maybe_all_reduce_tensor_model_parallel(
                experts_out
            )

        return experts_out

In vLLM, this SP implementation is not used in the forward function of any attention block, which further proves it is an implementation/variation of Megatron-SP.
In contrast, Ulysses/Ring Attention (Context Parallelism) operates primarily on the Attention layer and requires a distinct CP process group that should ideally be orthogonal to TP.

To facilitate a deeper community discussion, we should establish a more detailed distinction between Sequence Parallelism (Megatron-SP) and Context Parallelism (CP), particularly regarding how they map to vLLM's architecture.

1. vLLM Terminology (Proposed)

In vLLM, it is more accurate to distinguish between:
-PCP (Prefill-Context-Parallel): This is conceptually similar to the "Sequence Parallelism" (SP) found in DiT (Diffusion Transformers).
-DCP (Decode-Context-Parallel): Parallelism applied specifically during the decoding phase.

2. External Framework Comparisons

To avoid confusion with existing implementations, here is how the industry generally defines these terms:

• Megatron Sequence Parallelism (SP):
  Essence: Reduces activation memory per GPU to minimize re-computation (activation checkpointing) and speed up training.
  Usage: Typically coupled tightly with Tensor Parallelism (TP).
• Megatron Context Parallelism (CP):
  Essence: An enhanced version of SP.
  Mechanism: It incorporates Ring-Attention-like techniques (performing Ring-Attention across ranks with identical TP-PP-DP positions) to support training in conjunction with Megatron's various hybrid parallelism strategies
• DeepSpeed Ulysses:
  Context: DeepSpeed ZeRO is essentially data parallelism with model parallel characteristics. Long sequences pressure single-GPU memory during the full MHA calculation.
  Mechanism: Ulysses splits the input along the sequence dimension and uses all-to-all communication to allow each GPU to compute only a subset of attention heads.
• Ring-Attention:
  Essence: Can be viewed as a distributed implementation of Flash Attention V2.
  Mechanism: Each GPU computes the Multi-Head Attention (MHA) only for the specific seq_chunk it maintains, passing KV blocks in a ring.