Skip to content

Adding JIT cache #5

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.

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

Merged
merged 18 commits into from
Apr 10, 2025
Merged

Adding JIT cache #5

merged 18 commits into from
Apr 10, 2025

Conversation

bringlein
Copy link
Collaborator

@bringlein bringlein commented Apr 10, 2025

The launch overhead of triton kernels is a well known problem (see e.g. 1, 2, 3). Parts of the launch overhead comes from the fact that the triton JIT checks very carefully if an existing binary is safe to use.

In many scenarios, these checks can be relaxed.
This PR adds such a cache with relaxed checks is implemented by triton_dejavu.jitcache. It is implemented as a decorator that could be used in front of the triton.jit decorator:

@triton_dejavu.jitcache(
    check_keys=["x", "BLOCK_SIZE", "USE_ALIBI_SLOPES", "SLIDING_WINDOW", "filter_by_query_len"],
)
@triton.jit
def kernel_paged_attention_.... 

Details see Readme.

Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
Signed-off-by: Burkhard Ringlein <[email protected]>
@bringlein bringlein merged commit ae5de6a into main Apr 10, 2025
1 check passed
@bringlein bringlein deleted the ngl_cache_exp branch April 10, 2025 13:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants