Skip to content

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:

@george-mako
Copy link

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

@ai-nikolai
Copy link

@PaliC is this still up to date, or are there big changes in KernelBench since this PR was drafted?

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.

3 participants