Code

Add support for small page sizes#824

Code

Open

skrider wants to merge 19 commits intoDao-AILab:mainDao-AILab/flash-attention:mainfrom skrider:small-pagesskrider/flash-attention:small-pagesCopy head branch name to clipboard

+195-68

Commits

Commits on Mar 26, 2024

• add print statements for debugging

  skrider

  committed
  8efeb7f

• add print statements for debugging

  skrider

  committed
  ac5e78a

• reshape gmem copy

  skrider

  committed
  14b190b

• only test trivial block size

  skrider

  committed
  409431b

• implement kv page iteration functions

  skrider

  committed
  70dd049

• rearrange initial offset computation

  skrider

  committed
  59e76be

• tests passing for single page k

  skrider

  committed
  175369f

• paged copy refactor working for page size 256

  skrider

  committed
  3691677

• allow small page sizes in flash api

  skrider

  committed
  c05b857

• remove print statements

  skrider

  committed
  347a625

• tidy flash_fwd_kernel

  skrider

  committed
  3bb71a9

• compiles for all h but 128

  skrider

  committed
  fa13c6b

• all working except rotary embedding

  skrider

  committed
  bde5aec

• add page size 16 to tests

  skrider

  committed
  bc66858

• reshape rotary sin/cos copy to align with paged KV copy

  skrider

  committed
  0f5a45e

• revert hardcoded rotcossin thread layout

  skrider

  committed
  802dd6a

• resolve page offsets absolutely not relatively

  skrider

  committed
  135a1da

• add test for page table overflow

  skrider

  committed
  a63157e

• allow smaller page sizes in varlen api

  skrider

  committed
  7968148