Presentation
UltraAttn: Efficiently Parallelizing Attention Through Hierarchical Context-Tiling
DescriptionLong-context comprehension is a crucial capability for LLM. Context parallelism and irregular block sparse attention are two key technologies to accelerate long contextual training and inference. Existing context parallelism techniques for attention suffer from poor scalability, owing to their common characteristics: the striped-like partition pattern. The striped-like partition pattern causes high communication traffic and inflexible kernel granularity, which in turn results in low single-kernel device utilization.
To address these problems, we propose UltraAttn, a novel context parallelism solution for irregular attention. UltraAttn hierarchically tiles the context to reduce communication cost. UltraAttn also performs context-tiling at the kernel level to adjust the granularity of kernels to trade off between kernel overlap and single-kernel device utilization. UltraAttn executes distributed attention with an ILP-based runtime to optimize latency. We evaluate UltraAttn on 64 GPUs. UltraAttn achieves 5.5× speedup on average in different types of irregular attention over the state-of-the-art context parallelism techniques.
To address these problems, we propose UltraAttn, a novel context parallelism solution for irregular attention. UltraAttn hierarchically tiles the context to reduce communication cost. UltraAttn also performs context-tiling at the kernel level to adjust the granularity of kernels to trade off between kernel overlap and single-kernel device utilization. UltraAttn executes distributed attention with an ILP-based runtime to optimize latency. We evaluate UltraAttn on 64 GPUs. UltraAttn achieves 5.5× speedup on average in different types of irregular attention over the state-of-the-art context parallelism techniques.
Event Type
Paper
TimeThursday, 20 November 202511:15am - 11:38am CST
Location275
HPC for Machine Learning