4min snip

Latent Space: The AI Engineer Podcast — Practitioners talking LLMs, CodeGen, Agents, Multimodality, AI UX, GPU Infra and all things Software 3.0 cover image

FlashAttention 2: making Transformers 800% faster w/o approximation - with Tri Dao of Together AI

Latent Space: The AI Engineer Podcast — Practitioners talking LLMs, CodeGen, Agents, Multimodality, AI UX, GPU Infra and all things Software 3.0

NOTE

Challenges in Kernel Fusion for Attention with Softmax Operation

The system and machine learning communities have been separately exploring ideas to optimize attention operations./nKernel fusion, which combines multiple operations into a single subroutine, has been a challenge due to dependencies and the softmax operation./nThe online softmax trick, introduced in papers by NVIDIA and Google, allows breaking up the softmax into smaller pieces and rescaling for optimization./nCombining ideas from both sides, such as fusion techniques and the online softmax trick, can lead to effective optimization strategies for attention./nKernel fusion can reduce memory usage but may limit flexibility for researchers who want to experiment with modifications to attention./nCompiler advancements, like embedding in PyTorch, are being explored to automatically fuse kernels and optimize code./nOptimizing attention operations with compilers is still ongoing, and attention-specific challenges like softmax rewriting add complexity./nIn the future, compilers may be able to perform these optimizations, saving time and effort in manual optimization.

00:00

Get the Snipd
podcast app

Unlock the knowledge in podcasts with the podcast player of the future.
App store bannerPlay store banner

AI-powered
podcast player

Listen to all your favourite podcasts with AI-powered features

Discover
highlights

Listen to the best highlights from the podcasts you love and dive into the full episode

Save any
moment

Hear something you like? Tap your headphones to save it with AI-generated key takeaways

Share
& Export

Send highlights to Twitter, WhatsApp or export them to Notion, Readwise & more

AI-powered
podcast player

Listen to all your favourite podcasts with AI-powered features

Discover
highlights

Listen to the best highlights from the podcasts you love and dive into the full episode