-
Notifications
You must be signed in to change notification settings - Fork 104
Simplified Thunderkittens Port #107
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
base: main
Are you sure you want to change the base?
Simplified Thunderkittens Port #107
Conversation
1ce1f1b to
faf0935
Compare
|
We should add instructions for installing the TK dir in local environments to the README. Think there's also a mismatch with the CUDA version in the current TK repo too. It's hard coded to 12.6 if we rely on TK instructions to call 'source env.src' |
@willhu-jpg Added instructions to the README, lmk what you think. Also, could you clarify the CUDA version point? Like the tk-v2 branch uses a different version than kernelbench right now? |
|
README changes look good. For the CUDA version, checkout the env.src in the TK repo. Typical setup is to call "source env.src" to set all the environment variables, but they have it hardcoded to 12.6. |
willhu-jpg
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
changes look good otherwise! let's merge it :)
README.md
Outdated
| tk_root = os.environ.get("THUNDERKITTENS_ROOT", "/root/ThunderKittens") | ||
| ``` | ||
|
|
||
| This allows the kernel to include the right TK primitives.s |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit typo
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sounds good, fixed the typo. @simonguozirui TK works on CUDA versions 12.6 - 12.9, rn in modal scripts the system uses CUDA 12.8 by default which should be fine. Locally users will set CUDA_HOME, default is 12.6 for TK's env.src. Lmk if this arrangement sounds fine.
Thunderkittens backend support, but with a simplified load_inline() implementation.
This turns out to be much simpler than using separate files, but has its own tradeoffs and implications.
Also, users don't have to put the TK repo in the root directory, it's automatically cloned to the Modal.