Skip to content
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

Add command line argument parsing for reduction dimensions in Triton sum kernel #2284

Closed
wants to merge 1 commit into from

Conversation

jananisriram
Copy link
Contributor

Summary:
Add argument parsing for the command line in order to pass in dimension(s) across which the kernel reduces and enable more rigorous testing of different versions of the sum kernel, referencing torchbenchmark/operators/fb/flash_attention/operator.py.

Inherit the __init__ function from the parent class BenchmarkOperator in order to facilitate command line argument parsing.

Change dim type to list to avoid type issues resulting from tl.constexpr.

Modify equality checks in kernel and operator to satisfy type requirements for dim.

Reviewed By: xuzhao9

Differential Revision: D58212366

…sum kernel

Summary:
Add argument parsing for the command line in order to pass in dimension(s) across which the kernel reduces and enable more rigorous testing of different versions of the sum kernel, referencing [torchbenchmark/operators/fb/flash_attention/operator.py](https://www.internalfb.com/code/fbsource/[864a578ce44afdba619d50a352c8ca3b783e05ef]/fbcode/pytorch/benchmark/torchbenchmark/operators/fb/flash_attention/operator.py?lines=84).

Inherit the `__init__` function from the parent class `BenchmarkOperator` in order to facilitate command line argument parsing.

Change `dim` type to `list` to avoid type issues resulting from `tl.constexpr`.

Modify equality checks in kernel and operator to satisfy type requirements for `dim`.

Reviewed By: xuzhao9

Differential Revision: D58212366
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D58212366

@facebook-github-bot
Copy link
Contributor

This pull request has been merged in 9f2ab74.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants