#29 · Flash Attention 3

running

◷train_gpt.py▸flash_attn_varlen_func◎learner＄no GPU

Blocks

Quick summary

One import swaps the entire attention kernel. flash_attn_varlen_func (Flash Attention 3) handles variable-length…

Full explanation below the code →

fba-lab — train_gpt.py · flash_attn_varlen_funcexecuting

// block: flash_attn_varlen_func · lines 23–23$ study train_gpt.py --block fa3-importOne import swaps the entire attention kernel. flash_attn_varlen_func (Flash Attention 3) handles var… ✓

Explanation

One import swaps the entire attention kernel. flash_attn_varlen_func (Flash Attention 3) handles variable-length sequences packed together — no padding tokens, no wasted compute.

Think about

What does 'varlen' mean here and why does it matter for a speedrun?

// architecture

Live diagram

Record #29 · PR #118 · 163.84 s on 2× H100

Flash Attention 3

One import cut attention from 3 ms to 0.34 ms

Before this PR, the model called Flash Attention 2 on padded sequences. Padding tokens are zeros — fillers added so every sequence reaches the same fixed length. The GPU paid full compute cost to process them anyway. Every training step burned cycles on nothing.

Flash Attention 3's varlen API packs multiple real sequences end-to-end into one flat token buffer. No padding. Every GPU cycle touches a real token. The attention kernel dropped from 3 ms to 0.34 ms — barely visible in the profiler.

# line 23 — one import swaps the entire attention kernel
from flash_attn_interface import flash_attn_varlen_func

What is varlen packing?

Standard attention pads all sequences to the same fixed length. varlen packs them end-to-end into one flat buffer. cu_seqlens marks where each sequence starts. The kernel processes only real tokens — no zero-padding ever reaches the GPU.

Where does this appear in the profiler?

On the profiler trace beat, click flash_attention — it reads 0.34 ms total across all 11 transformer layers. The kernel is so fast it barely registers. Before FA3, attention was ~3 ms and clearly visible. Now it is not the bottleneck.

One import cut attention from 3 ms to 0.34 ms

# line 23 — one import swaps the entire attention kernel
from flash_attn_interface import flash_attn_varlen_func

What is varlen packing?

Where does this appear in the profiler?