FlexCP: FlexAttention for Long Sequence Training...
summary
FlexAttention integrates with context parallelism to shard QKV plus the block mask and rewrite mask functions so local QKV indexes map to global ones; an automatic load balancer called process time run robin (PTR) reorders Q indices by row density to avoid imbalance. Benchmarks on 128k sequences show PTR gives >20% speedup on medium/large docs (density ~15%) and >50% on full causal documents, while small-docs (2% density) see minor overhead; ring-attention support is still pending and APIs land in PyTorch nightly.
AI-written summary · about
chapters (6)
- 0:00FlexAttention & Context Parallel intro
- 1:05FlexAttention recap & masking
- 2:53Context Parallel Shot for block masks
- 3:55Sparsity for load balancing
- 6:31Round Robin load balancing
- 9:45Benchmarks & future plans
related talks
09mFlexAttention + FlashAttention-4: Fast and Flexible
36mML Efficiency for Large Models: From Data Efficiency to Faster Transformers
26mEfficiently serving LLMs at scale.
17mCacheSlide: Unlocking Cross Position-Aware KV Cache Reuse for Accelerating LLM Serving
28mOptimizing Context Tuning for RAG Systems
20mFlexLLM: Token-Level Co-Serving of LLM Inference and Finetuning with SLO Guarantees
15mPUZZLE: Efficiently Aligning Large Language Models through Light-Weight Context...
42mAttention and Memory in Deep Learning Networks