feat(kernels): SM90 TMA + mma.sync FlashAttention, forward + backward (causal, varlen, LSE) - #237

Open

Billy1900 wants to merge 6 commits into

RL-Align:mainfrom

Billy1900:feat/sm90-attention-varlen-fwd

Commits on Jul 19, 2026

feat(kernels): Triton FlashAttention LSE export + varlen packing
Billy1900
committed
3d5364b
fix(kernels): use DO's own strides in varlen backward, not Q's/Out's
Billy1900
committed
8276596

Commits on Jul 20, 2026

feat(kernels): SM90 TMA + mma.sync FlashAttention forward (causal, varlen, LSE)
Billy1900
committed
80ffd14
feat(kernels): SM90 backward (dQ/dK/dV) for varlen FlashAttention, fully differentiable
Billy1900
committed
013f4b4

Commits on Jul 27, 2026

Merge branch 'main' into feat/sm90-attention-varlen-fwd
Flink-ddd
authored
1f26356

Commits on Jul 28, 2026

fix(lint): black-format SM90 varlen attention files, shorten setup.py comment

Billy1900
and
claude
committed
79d5c72