Repository navigation

Add CP KDA backward megakernel - #18

Open

fdz-1999 wants to merge 13 commits into

kda-mega-backward-packedfrom

kda-mega-cp-backward-v2

Commits on Sep 23, 2026
- Fuse CP saved-state local KDA backward after communication
  
  fdz-1999
  and
  Fred33146
  committed
  c4a7349
- Enable CP local backward fusion with rematerialized states
  
  fdz-1999
  and
  Fred33146
  committed
  88bb194
- Port the complete context-parallel KDA backward megakernel
  
  fdz-1999
  and
  Fred33146
  committed
  27b61e1
- chore: keep CP migration focused on code and tests
  fdz-1999
  committed
  e53b613
- fix: keep CP megakernel on its BF16 contract
  fdz-1999
  committed
  941c1b0
- test: validate CP megakernel within its BF16 contract
  fdz-1999
  committed
  bc1eb4d
Commits on Sep 24, 2026
- fix: address review P1s in CP backward megakernel
  fdz-1999
  committed
  6ae6a0a
- fix: gate CP megakernel dh0 on the caller-reported state flag
  fdz-1999
  committed
  67f603f
- Fix CP recurrence-seed regression and shard_map out_specs in tests
  fdz-1999
  committed
  fb79285
- Return a uniform pytree from the initial-state grad test
  fdz-1999
  committed
  bd34b7e
- fix: preserve tiled beta load in CP backward stack
  fdz-1999
  committed
  f99f47a
Commits on Sep 29, 2026
- bench: parameterize KDA training shape comparisons
  fdz-1999
  committed
  51a4139
- Merge remote-tracking branch 'origin/kda-mega-backward-packed' into kda-mega-cp-backward-v2
  fdz-1999
  committed
  00ebb27