Skip to content

[example] add jagged_softmax example #480

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.

Sign up for GitHub

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

Jump to bottom
Merged
merged 11 commits into from
Aug 14, 2025
Merged

[example] add jagged_softmax example #480

merged 11 commits into from
Aug 14, 2025

Conversation

pianpwk
Copy link
Contributor

@pianpwk pianpwk commented Aug 8, 2025 •
edited
Loading

very new to this, but thought it was worth putting up, at least for feedback :)

results from running HELION_USE_DEFAULT_CONFIG=0 python benchmarks/run.py --kernel jagged_softmax --num-inputs 10 --metrics speedup:

  x_val    torch_jagged_softmax_torch_sum-speedup    triton_jagged_softmax_simple_fused-speedup    triton_jagged_softmax_variable_length_loop-speedup    torch_compile_nested_tensor_integration-speedup    helion_jagged_softmax_tritonbench-speedup
-------  ----------------------------------------  --------------------------------------------  ----------------------------------------------------  -------------------------------------------------  -------------------------------------------
      4                                   1.84622                                      0.974523                                              0.954934                                           0.322953                                      1.96138
      4                                   2.95117                                      1.45007                                               1.43405                                            0.495255                                      3.30609
      4                                   1.6221                                       0.813262                                              0.871952                                           0.307515                                      2.00381
      4                                   1.66235                                      0.848234                                              0.872248                                           0.314032                                      1.97156
      4                                   1.59061                                      0.856801                                              0.914473                                           0.293083                                      1.90051
      4                                   1.68989                                      0.848728                                              0.877103                                           0.302414                                      1.35153
      4                                   1.69289                                      0.831421                                              0.866428                                           0.308762                                      1.38123
      4                                   1.59673                                      0.861041                                              0.862671                                           0.301662                                      1.40576
      4                                   1.66998                                      0.872588                                              0.899465                                           0.293748                                      1.57284
      4                                   1.6537                                       0.866808                                              0.877925                                           0.306829                                      1.88681
average                                   1.79756                                      0.922347                                              0.943125                                           0.324625                                      1.87415

from running HELION_USE_DEFAULT_CONFIG=0 python benchmarks/run.py --kernel jagged_softmax --num-inputs 3 --metrics accuracy:

  x_val    torch_jagged_softmax_torch_sum-accuracy    triton_jagged_softmax_simple_fused-accuracy    triton_jagged_softmax_variable_length_loop-accuracy    torch_compile_nested_tensor_integration-accuracy    helion_jagged_softmax_tritonbench-accuracy
-------  -----------------------------------------  ---------------------------------------------  -----------------------------------------------------  --------------------------------------------------  --------------------------------------------
      4                                          1                                              1                                                      1                                                   1                                             1
      4                                          1                                              1                                                      1                                                   1                                             1
      4                                          1                                              1                                                      1                                                   1                                             1
average                                          1                                              1                                                      1                                                   1                                             1

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 8, 2025
@pianpwk pianpwk changed the title [WIP] add jagged_softmax example [WIP][example] add jagged_softmax example Aug 9, 2025
@pianpwk pianpwk changed the title [WIP][example] add jagged_softmax example [example] add jagged_softmax example Aug 9, 2025
@pianpwk pianpwk marked this pull request as ready for review August 9, 2025 01:19
@jansel jansel requested a review from yf225 August 9, 2025 15:53
@jansel
Copy link
Contributor

jansel commented Aug 9, 2025

lint/test failures?

pianpwk
pianpwk commented Aug 12, 2025
View reviewed changes
examples/jagged_softmax.py
[flat_indices],
block_out,
extra_mask=combined_mask,
)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wasn't sure how to write this without the store, and it seems this store breaks ref eager mode testing

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm that's weird - do you mind filing an issue for the ref eager mode problem? thanks!

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yep, filed one at #496

@yf225
Copy link
Contributor

yf225 commented Aug 13, 2025

@pianpwk would you like to check that --metrics accuracy shows 1 for the Helion kernel too? thanks!

@pianpwk
Copy link
Contributor Author

pianpwk commented Aug 13, 2025

@pianpwk would you like to check that --metrics accuracy shows 1 for the Helion kernel too? thanks!

ah yep, just updated the description :)

@pianpwk pianpwk mentioned this pull request Aug 13, 2025
Open
yf225
yf225 approved these changes Aug 13, 2025
View reviewed changes
Copy link
Contributor

@yf225 yf225 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks @pianpwk !

@yf225 yf225 merged commit 043d0f4 into pytorch:main Aug 14, 2025
13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@yf225 yf225 yf225 approved these changes

Assignees
No one assigned
Labels
CLA Signed This label is managed by the Meta Open Source bot.
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@pianpwk @jansel @yf225