Skip to content

Use shared compile cache for KDA backward - #247

Merged
drisspg merged 1 commit into
mainfrom
drisspg/stack/40
Aug 8, 2026
Merged

Use shared compile cache for KDA backward#247
drisspg merged 1 commit into
mainfrom
drisspg/stack/40

Conversation

@drisspg

@drisspg drisspg commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Stacked PRs:


Use shared compile cache for KDA backward

Human Note

Agent note
Migrate the existing SM100 KDA backward kernel from its local dictionary and direct cute.compile
call to the shared persistent cache and typed TVM-FFI compile boundary. Give each specialization a
stable profiler/artifact name and derive the persistent launch grid from parent-discovered target
metadata so future forked compiler workers never query an inherited CUDA context.

This commit intentionally migrates only the existing chunk_delta_h_bwd_v1 call site. Additional
KDA kernels can move to the same flow independently after the cache and tuning layers land.

Test Plan:

~/.venvs/dev/bin/ruff check attn_gym/linear/kda/bwd/cute/chunk_delta_h_bwd_v1.py
ATTN_GYM_CUTE_CACHE_DIR=$(mktemp -d) gpu-run 0 -- ~/.venvs/nightly/bin/python -c \
  'from attn_gym.linear.kda.bwd.cute.chunk_delta_h_bwd_v1 import _compile_bwd_dhu; assert callable(_compile_bwd_dhu(False, 4, 128, 128, 64, 16))'

Human Note

Agent note
Migrate the existing SM100 KDA backward kernel from its local dictionary and direct `cute.compile`
call to the shared persistent cache and typed TVM-FFI compile boundary. Give each specialization a
stable profiler/artifact name and derive the persistent launch grid from parent-discovered target
metadata so future forked compiler workers never query an inherited CUDA context.

This commit intentionally migrates only the existing `chunk_delta_h_bwd_v1` call site. Additional
KDA kernels can move to the same flow independently after the cache and tuning layers land.

Test Plan:

```bash
~/.venvs/dev/bin/ruff check attn_gym/linear/kda/bwd/cute/chunk_delta_h_bwd_v1.py
ATTN_GYM_CUTE_CACHE_DIR=$(mktemp -d) gpu-run 0 -- ~/.venvs/nightly/bin/python -c \
  'from attn_gym.linear.kda.bwd.cute.chunk_delta_h_bwd_v1 import _compile_bwd_dhu; assert callable(_compile_bwd_dhu(False, 4, 128, 128, 64, 16))'
```

stack-info: PR: #247, branch: drisspg/stack/40
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 8, 2026
@drisspg
drisspg marked this pull request as draft August 8, 2026 01:17
@drisspg
drisspg changed the base branch from drisspg/stack/39 to main August 8, 2026 01:17
@drisspg
drisspg changed the base branch from main to drisspg/stack/39 August 8, 2026 01:18
@drisspg
drisspg marked this pull request as ready for review August 8, 2026 01:18
drisspg added a commit that referenced this pull request Aug 8, 2026
Human Note

Agent note
Migrate the existing SM100 KDA backward kernel from its local dictionary and direct `cute.compile`
call to the shared persistent cache and typed TVM-FFI compile boundary. Give each specialization a
stable profiler/artifact name and derive the persistent launch grid from parent-discovered target
metadata so future forked compiler workers never query an inherited CUDA context.

This commit intentionally migrates only the existing `chunk_delta_h_bwd_v1` call site. Additional
KDA kernels can move to the same flow independently after the cache and tuning layers land.

Test Plan:

```bash
~/.venvs/dev/bin/ruff check attn_gym/linear/kda/bwd/cute/chunk_delta_h_bwd_v1.py
ATTN_GYM_CUTE_CACHE_DIR=$(mktemp -d) gpu-run 0 -- ~/.venvs/nightly/bin/python -c \
  'from attn_gym.linear.kda.bwd.cute.chunk_delta_h_bwd_v1 import _compile_bwd_dhu; assert callable(_compile_bwd_dhu(False, 4, 128, 128, 64, 16))'
```

stack-info: PR: #247, branch: drisspg/stack/40
@drisspg
drisspg force-pushed the drisspg/stack/40 branch 2 times, most recently from 7660d04 to 7a10b6a Compare August 8, 2026 03:20
@drisspg
drisspg changed the base branch from drisspg/stack/39 to main August 8, 2026 03:22
@drisspg
drisspg merged commit 5101ca5 into main Aug 8, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant