Skip to content

Add CP KDA backward megakernel #18

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fdz-1999 wants to merge 13 commits into
base: kda-mega-backward-packed
Choose a base branch
from
Open

Add CP KDA backward megakernel #18

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
c4a7349
Fuse CP saved-state local KDA backward after communication
fdz-1999 Sep 7, 2026
88bb194
Enable CP local backward fusion with rematerialized states
fdz-1999 Sep 7, 2026
27b61e1
Port the complete context-parallel KDA backward megakernel
fdz-1999 Sep 7, 2026
e53b613
chore: keep CP migration focused on code and tests
fdz-1999 Sep 22, 2026
941c1b0
fix: keep CP megakernel on its BF16 contract
fdz-1999 Sep 23, 2026
bc1eb4d
test: validate CP megakernel within its BF16 contract
fdz-1999 Sep 23, 2026
6ae6a0a
fix: address review P1s in CP backward megakernel
fdz-1999 Sep 24, 2026
67f603f
fix: gate CP megakernel dh0 on the caller-reported state flag
fdz-1999 Sep 24, 2026
fb79285
Fix CP recurrence-seed regression and shard_map out_specs in tests
fdz-1999 Sep 24, 2026
bd34b7e
Return a uniform pytree from the initial-state grad test
fdz-1999 Sep 24, 2026
f99f47a
fix: preserve tiled beta load in CP backward stack
fdz-1999 Sep 24, 2026
51a4139
bench: parameterize KDA training shape comparisons
fdz-1999 Sep 29, 2026
00ebb27
Merge remote-tracking branch 'origin/kda-mega-backward-packed' into k…
fdz-1999 Sep 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .github/ci/shards.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -537,6 +537,8 @@ class Spec(_RequiredSpec, total=False):
'tokamax/_src/ops/experimental/kda/pallas_mosaic_tpu_remat_test.py',
'tokamax/_src/ops/experimental/kda/pallas_mosaic_tpu_packed_bwd_test.py',
'tokamax/_src/ops/experimental/kda/pallas_mosaic_tpu_packed_gradients_test.py',
'tokamax/_src/ops/experimental/kda/pallas_mosaic_tpu_cp_fused_test.py',
'tokamax/_src/ops/experimental/kda/pallas_mosaic_tpu_cp_megakernel_test.py',
'tokamax/_src/ops/experimental/kda/pallas_mosaic_tpu_bwd_fused_test.py',
),
minutes=16,
Expand Down
10 changes: 9 additions & 1 deletion tokamax/_src/ops/experimental/kda/pallas_mosaic_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,11 @@ class Config:
normalization backward; beta retains its existing gather.
`fuse_backward=True` also fuses the non-CP saved-state backward. Manual
state rematerialization is fused separately with `fuse_rematerialization=True`;
CP retains its existing backward path.
`fuse_cp_backward=True` fuses the saved-state local reverse pass after
CP communication. With `fuse_rematerialization=True`, this also applies
to CP rematerialization, preserving its staged state reconstruction.
`cp_megakernel=True` selects the full CP core (including remote DMA)
instead, for 128-aligned K/V. It is opt-in pending TPU validation.
"""

chunk_size: Annotated[int, pydantic.Field(gt=0)] = 64
Expand All @@ -79,6 +83,8 @@ class Config:
packed_backward: bool = False
packed_gradients: bool = False
fuse_backward: bool = True
fuse_cp_backward: bool = False
cp_megakernel: bool = False
# Opt-in until TPU compilation and device-time validation completes.
fuse_rematerialization: bool = False

Expand Down Expand Up @@ -578,6 +584,8 @@ def _fwd(
fuse_rematerialization=config.fuse_rematerialization,
packed_inputs=packed_inputs,
packed_gradients=config.packed_gradients,
fuse_cp_backward=config.fuse_cp_backward,
cp_megakernel=config.cp_megakernel,
)

grads = {
Expand Down
Loading