-
Notifications
You must be signed in to change notification settings - Fork 15
Add cross_entropy example and unit test #320
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
Conversation
stack-info: PR: #320, branch: yf225/stack/28
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.
Could we call the pytorch function for cross entropy (that will get decomposed) here to use less code?
# TritonBench configuration - adjust based on available GPU memory | ||
if get_gpu_memory_info()[0] < 16.0: | ||
# Low memory configuration for GPUs with less than 16GB | ||
TRITONBENCH_ARGS = {"B": 4, "T": 512, "v_range": "10,15"} |
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.
Where is this used?
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.
It’s for development on low-VRAM machines - switched to use HELION_DEV_LOW_VRAM to gate this
28a5978
to
0bbe404
Compare
@jansel I tried to just use the pytorch F.cross_entropy function within Helion tile loop but ran into some lowering issues. I will do another sweep later to test lowering for all common F functions and fix the issues. |
Stacked PRs:
Add cross_entropy example and unit test