Repository navigation

Add packed KDA forward megakernel - #19

Open

fdz-1999 wants to merge 13 commits into

mainfrom

kda-mega-forward-packed

Commits on Sep 22, 2026
- Fuse non-CP KDA forward stages in Mosaic TPU
  
  fdz-1999
  and
  Fred33146
  committed
  c0d3a76
- Read packed input windows in fused non-CP KDA forward
  
  fdz-1999
  and
  Fred33146
  committed
  3070e24
- Compact packed KDA outputs with an opt-in Pallas kernel
  
  fdz-1999
  and
  Fred33146
  committed
  a148c99
- fix: support singleton KDA forward tiles
  fdz-1999
  committed
  1e0b931
- fix: validate packed forward independently on TPU
  fdz-1999
  committed
  7b79701
- fix: zero compacted output tails before VMEM stores
  fdz-1999
  committed
  788d64d
Commits on Sep 23, 2026
- bench: compare KDA megakernel stages on TPU
  fdz-1999
  committed
  0b4bbc0
- bench: label inference baseline as existing Mosaic forward
  fdz-1999
  committed
  ec555a2
- bench: support pinned upstream KDA baseline
  fdz-1999
  committed
  32ce8f2
- bench: sweep CP megakernel sequence lengths
  fdz-1999
  committed
  8ee058c
Commits on Sep 24, 2026
- test: cover multi-head forward and VJP correctness
  fdz-1999
  committed
  64622d9
- test: broadcast the packed segment map across the mini batch
  fdz-1999
  committed
  4085015
Commits on Sep 29, 2026
- bench: parameterize KDA training shape comparisons
  fdz-1999
  committed
  7ad795a