Add an int64 path for mlp kernels #3614
Draft
+92
−12
Conversation
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
unsloth/kernels/geglu.py
Outdated
| device = gate.device | ||
| out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = device) | ||
| grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) | ||
| if n_elements <= (2**31) - 1024: |
Contributor
There was a problem hiding this comment.
Why -1024? Is it maybe hd?
Collaborator
Author
There was a problem hiding this comment.
yes I forgot to account for hd. The idea is that I wanted to add a buffer just to be safe.
unsloth/kernels/geglu.py
Outdated
| batch_seq_len, hd = e.shape | ||
| n_elements = e.numel() | ||
| grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) | ||
| if n_elements <= (2**31) - 1024: |
Contributor
There was a problem hiding this comment.
Maybe move (2**31) to a global var
unsloth/kernels/swiglu.py
Outdated
| e, | ||
| g, | ||
| n_elements, | ||
| BLOCK_SIZE: tl.constexpr, |
Contributor
There was a problem hiding this comment.
there is actually a way to use 1 kernel only and dispatch, but for now this is fine - we can refactor later
mmathew23
force-pushed
the
tiled/contextlen
branch
2 times, most recently
from
c008eca to
262ada3
Compare
mmathew23
force-pushed
the
tiled/contextlen
branch
from
262ada3 to
833d91f
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
The llama mlp kernels produce nans with extremely long context length. This is happens when the num_elements is greater than 2**31. In these cases it's best to calculate offsets with tl.int64 instead of int32. This PR will route to int64 kernels if the num_elements is big enough.