feat(flash-attention): add varlen packed-sequence support

Browse Source

Implements variable-length (varlen) FlashAttention that processes
mixed-length batches without padding waste:

- New CUDA kernel flash_attention_varlen_forward with BLOCK_Q=64 /
  BLOCK_K=64 tiling; grid=(ceil(max_seqlen_q/64), num_heads, 1).
  Each block uses a linear scan over cu_seqlens_q to identify its
  owning sequence and exits early when past sequence end.
- New Rust module flash_varlen_forward: always-compiled CPU simulation
  (varlen_attention_cpu) for testing + #[cfg(cuda)] FlashVarlenKernel.
- SdpaBackend::VarLen variant added to backend_selector.
- 8 new CPU-only tests; total test count: 50.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

...

This commit is contained in:

Omar Sobh

2026-06-27 01:09:39 +00:00

co-authored by Claude Sonnet 4.6

parent 311eb23dbd

commit 80d7c7fb6c

5 changed files with 2121 additions and 0 deletions

Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL

Download Patch File Download Diff File Expand all files Collapse all files

docs/superpowers/plans/2026-06-27-varlen-flash-attention.md

+1220

File diff suppressed because it is too large Load Diff