FlexAttention + FlashAttention-4: Fast and Flexible
summary
FlexAttention integrates with FlashAttention-4 to run Flex-style score and mask hooks as JIT-instantiated CuTeDSL kernels via PyTorch Inductor, converting pointwise score mods into tensor-SSA and wiring them into FA4’s warp-specialized, block-sparse pipeline. Benchmarks show 1.2–3.2× speedups over Triton on compute-bound workloads, up to 2.7× forward / 3× backward on B200, and 1.3–2× on Hopper; practical notes: set Flash backend and adjust block size to 256×128 for Blackwell.
AI-written summary · about
related talks
11mFlex Attention for Inference
11mFlexCP: FlexAttention for Long Sequence Training...
08mHacks to Make LLM Training Faster
14mFlexPipe: Maximizing Training Efficiency for Transformer-based Models with...
12mLightning Talk: Accelerating Inference on CPU with Torch.Compile
20mFlexLLM: Token-Level Co-Serving of LLM Inference and Finetuning with SLO Guarantees
43mSpeeding up Whisper for speech recognition by 10x - Optimisations and walkthrough
29mEnabling Lightweight, High-Performance FSDP With NVIDIA GPU