Skip to content

Add Triton Backend #35

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

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Add Triton Backend #35

wants to merge 7 commits into from

Conversation

PaliC
Copy link
Collaborator

@PaliC PaliC commented Mar 18, 2025

This PR adds a triton backend to kernel bench. To invoke it simply add backend="triton" to the following 4 scripts (use them as normal otherwise)

scripts/generate_and_eval_single_sample_modal.py
scripts/generate_and_eval_single_sample.py
scripts/generate_samples.py
scripts/eval_from_generations.py

This PR also adds a {error_type}_name into the eval json. The reason for this is that it makes classifying errors (especially for triton) much easier. For example, from the error log it isn't obvious what an error is (ie. you might get at 37:15:\n h_start = pooled_row * stride - padding\n w_start = pooled_col * stride - padding\n\n # Initialize the max value\n max_val = tl.full((1,), float('-inf'), tl.float32)\n\n # Itera...). But if the error name is triton.compiler.errors.UnsupportedLanguageConstruct it's a lot more obvious.

Testing: I've tested the 4 scripts in both the triton and cuda variants and they seem to work normally. (outside of scripts/generate_and_eval_single_sample_modal.py which should be equivalent to scripts/generate_and_eval_single_sample.py)

Todo:

  • This PR is a little bit crude as it just adds the logic (and linting (sorry meta's IDE does this automatically). There are other things like updating the readme, and a bit of logical refactoring (use kernel instead of cuda for example). But I will leave that to a followup PR as this one is already 1000 lines.
  • Add fewshot for triton
  • Add CoT for triton

Below is the github copilot generated summary which is honestly pretty useful for navigating large PRs.

==========================================================================================
This pull request includes several changes to improve code readability and add new functionality to the scripts/eval_from_generations.py and scripts/generate_and_eval_single_sample.py files. The most notable changes include reformatting code for better readability, adding a new backend configuration option, and enhancing error logging.

Code readability improvements:

New functionality:

Error logging enhancements:

Miscellaneous:

@OpusManuum
Copy link

@simonguozirui Any update on when this will be merged to main?

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