CP method in ParaAttention and USP in xDiT

[TuanNT-ZenAI]

Today, I tested Flux_schnell with 8 steps in the CP method of ParaAttention and the USP of XDiT using 2x H100SXM GPUs. I found that USP is faster than the CP method, even though their repository states that USP Flux is implemented based on ParaAttention. Could you explain the differences between these two methods?

Thanks,
Chengzeyi

[chengzeyi]

Today, I tested Flux_schnell with 8 steps in the CP method of ParaAttention and the USP of XDiT using 2x H100SXM GPUs. I found that USP is faster than the CP method, even though their repository states that USP Flux is implemented based on ParaAttention. Could you explain the differences between these two methods?

Thanks, Chengzeyi

Can you try running ParaAttention with torch.compile? ParaAttention is designed to work best with that.

[chengzeyi]

And a big difference is that ParaAttention aims to keep simple and universal. So we utilize the native torch sdpa instead of flash_attn. Also we don’t change the original diffusers code to eliminate concatenation. Instead, we expect a compiler like torch.compile to do these graph level optimizations automatically.

[TuanNT-ZenAI]

thanks @chengzeyi