[feature] help us implement unified attention

[sayakpaul]

Currently, we support ring and ulysses attention backends for context parallelism. More info: https://huggingface.co/docs/diffusers/main/en/api/parallel

Unified CP is a generalization of performing Ulysses and Ring together. Both those methods become special subcases of Unified attention. Paper: https://arxiv.org/abs/2405.07719v3

Before proceeding with a PR, please use this thread to discuss your proposal and we can take it from there.

[tolgacangoz]

I could handle this after our MAGI-1 PR, if no one else will work on it before the PR's completion. MAGI-1 proposes to include another CP-specialized attention module: MagiAttention.

[yiyixuxu]

hey @tolgacangoz
Thanks for offering! However, we'd like to give other community members an opportunity to work on this feature as you already have many PR open.
Appreciate your understanding!

[punitvara]

Hi @sayakpaul I would like to read paper and propose plan for this. I am not sure how difficult or easy for first timer in diffuser repo. I will try getting enough context and post questions here.

[KarthikSundar2002]

I would like to work on this, I will get back with a proposal by tmrw :).

[KarthikSundar2002]

So I have read the paper and looked into the repo a bit.

I feel this could be a high level plan:

1. Implement AllToAll Operator
2. Implement Load Balance Partition

I could see an EquipartitionSharder present already in the hooks/context_parallel.py
So, we just need to implement the AllToAll Operator, and make sure to use it when the ring degree and ulysses degree both are not equal to 1.

Am I going in the right direction?

[punitvara]

I also read the paper and understood both ring and ulyssess attension. I have yet to look which parts are missing in repo. I can collaborate with Karthik on that if his plans looks good. Anyway I will also check missing pieces of puzzle. I assume multi GPU setup would run CI/CD so developer doesn't need it right ?

[sayakpaul]

We can leave the CI/CD for now and focus on the high-level API design.

I think a good entrypoint to implementing this would be to study:

diffusers/src/diffusers/models/attention_dispatch.py

Line 1218 in 0c35b58

def _templated_context_parallel_attention(

We have TemplatedRingAttention and TemplatedUlyssesAttention classes that are called based on the parallel config being set by the user.

So, we could think of implementing a similar TemplatedUnifiedAttention (with appropriate config validation) and invoking it from _templated_context_parallel_attention().

[punitvara]

ok makes sense. I will study today and come back by tomorrow.

[KarthikSundar2002]

@sayakpaul Thanks for pointing me to the existing implementation for Ring and Ulysses Attention. It has helped quite a lot!

This is the current Logic I have for the forward and backward pass of the Unified Attention:

Forward Pass Logic

For each device in the 2D mesh (ring_degree × ulysses_degree):
Device has local Q, K, V tensors. Sequence is split across ring dimension similar to Ring Attention. Heads are split across ulysses dimension.

Step 1: Ulysses-style head splitting (within each ring group)
Reshape QKV: [B, S_local, H, D] -> [B, S_local, ulysses_degree, H_local, D] . Share all the sequence chunks across all the devices.

Step 2: Ring-style sequence rotation (within each ulysses group)
For each ring iteration:
1. All-gather KV chunks along ring dimension
2. Compute attention: Q_local × K_rotated × V_rotated
3. Merge outputs using online softmax
4. Rotate to next ring position

Step 3: Ulysses-style head gathering
Gather all the heads and reshape output

Backward Pass Logic

1. Split the grad_out across the head dimension
2. In each of ring iterations: compute the grad for the specific chunk, then accumulate the gradient across ring iterations, then move on to the next ring position to do the same.
   3.Sum Gradients Across the Ulysses Dimension.

Is my logic right?? I will have a class implementation PR by tomorrow. After implementing the class, there are a few bookkeeping to do like updating the dispatcher code.

Just a doubt with compute requirements, I currently don't have access to a multi-gpu cluster to test the implementation in. Is it possible to get the compute from HF?

[punitvara]

Looks like @KarthikSundar2002 is more interested in this issue. Should I pick any other issue ? @sayakpaul

[KarthikSundar2002]

@sayakpaul I have implemented the class in my fork, however I couldn't find a base branch of that commit to raise a PR.

You could find my implementation here - KarthikSundar2002@d2464ef

If it seems fine, I will fix the dispatch code as well.

[Bissmella]

Hi @sayakpaul , Karthik,
I would also like to collaborate on this if possible. I've gone through the paper, its reference implementation, and the current TemplatedUlyssesAttention and TemplatedRingAttention code in diffusers. I believe we can implement this with minimal changes and in very simple way by leveraging the existing TemplatedRingAttention implementation, minimizing error and in under 100 lines of code.

A quick overview of unified SP attention:

USP combines Ring and Ulysses attention in a 2D mesh configuration. For example, with 8 GPUs arranged as a 2×4 mesh:

[GPU 0, GPU 1, GPU 2, GPU 3]  ← Ulysses group (row 0)
[GPU 4, GPU 5, GPU 6, GPU 7]  ← Ulysses group (row 1)
   ↓      ↓      ↓      ↓
ring0  ring1  ring2   ring3 
 Ring groups (columns)

2 rows -> 2 Ulysses groups
4 columns -> 4 ring groups

Implementation strategy:

The key insight is that Ring attention does the heavy lifting, we just need to add Ulysses AllToAll operations before and after. The implementation requires adding:

1. all_to_all_double which scatters in one dimension and gathers on another with simple forward and backward.

def all_to_all_double(tensor, scatter_idx, gather_idx, group):
    # group here is the Ulysses process group
    # (bs, L/N, heads, dim) → (bs, L, heads/N, dim) when scatter_idx=1, gather_idx=2
    # its backward is similar but with reversed scatter and gater indices
    ...

2. TemplatedUnifiedAttention which uses the all_to_all_double and TemplatedRingAttention in minimal few lines. I have verified the tensor shapes match between the all_to_all_double and TemplatedRingAttention and it should work seamlessly.

class TemplatedUnifiedAttention(...):
    # Ulysses: partition heads, gather sequence
    q, k, v = all_to_all_double(q, k, v, scatter_idx=1, gather_idx=2, group=ulysses_group)

    # Ring attention (existing implementation)
    output = TemplatedRingAttention.apply(q, k, v, ...)

    # Reverse Ulysses: gather heads, partition sequence  
    output = all_to_all_double(output, scatter_idx=2, gather_idx=1, group=ulysses_group)
    return output

[sayakpaul]

@Bissmella that sounds like a terrific idea! I like it a lot. @KarthikSundar2002 would you like to collaborate with @Bissmella (maybe help with testing and other things like validation? Apologies for the inconvenience.

[KarthikSundar2002]

@Bissmella that sounds really great! Love the idea! I would love to collaborate. I would like to give a go at implementing your idea today, and you could review it and add the changes to it as well. After that we could move on to adding all the validation and testing. Does this sound good?