Use shared compile cache for KDA backward - #247
Merged
Merged
Conversation
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
force-pushed
the
drisspg/stack/39
branch
from
August 8, 2026 01:16
24fd863 to
0dd6382
Compare
drisspg
force-pushed
the
drisspg/stack/40
branch
from
August 8, 2026 01:16
78e9670 to
221fce4
Compare
This was referenced Aug 8, 2026
drisspg
marked this pull request as draft
August 8, 2026 01:17
drisspg
force-pushed
the
drisspg/stack/40
branch
from
August 8, 2026 01:18
221fce4 to
d8bfec4
Compare
drisspg
marked this pull request as ready for review
August 8, 2026 01:18
drisspg
force-pushed
the
drisspg/stack/39
branch
from
August 8, 2026 03:19
0dd6382 to
9a34d18
Compare
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
force-pushed
the
drisspg/stack/40
branch
2 times, most recently
from
August 8, 2026 03:20
7660d04 to
7a10b6a
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.compilecall 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_v1call site. AdditionalKDA kernels can move to the same flow independently after the cache and tuning layers land.
Test Plan: