Skip to content

[triton] Remove the zero initialization of qk_acc by directly writing the result#1288

Merged
merrymercy merged 3 commits intosgl-project:mainfrom
ByronHsu:byhsu/simplify
Sep 1, 2024
Merged

[triton] Remove the zero initialization of qk_acc by directly writing the result#1288
merrymercy merged 3 commits intosgl-project:mainfrom
ByronHsu:byhsu/simplify

Conversation

@ByronHsu
Copy link
Copy Markdown
Collaborator

@ByronHsu ByronHsu commented Sep 1, 2024

Motivation

Currently, we do tl.zeros + tl.dot just to write qk to acc, but this can be simplified thanks to https://triton-lang.org/main/python-api/generated/triton.language.dot.html#triton.language.dot. Maybe triton compiler already does the optimization, but it is good to keep the code simple.

Modifications

Original:

qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
qk += tl.dot(q, k)

New

qk = tl.dot(q, k, out_dtype=tl.float32)

Checklist

  • Format your code according to the Contributor Guide.
  • Add unit tests as outlined in the Contributor Guide.
  • Update documentation as needed, including docstrings or example tutorials.

@ByronHsu ByronHsu changed the title Remove the zero initialization of qk_acc by directly writing the result to qk_acc. Remove the zero initialization of qk_acc by directly writing the result Sep 1, 2024
@ByronHsu ByronHsu changed the title Remove the zero initialization of qk_acc by directly writing the result [triton] Remove the zero initialization of qk_acc by directly writing the result Sep 1, 2024
@zhyncs
Copy link
Copy Markdown
Collaborator

zhyncs commented Sep 1, 2024

A series of nice work! Would you like to join our Slack channel https://sgl-fru7574.slack.com/ssb/redirect? Thanks!

@ByronHsu
Copy link
Copy Markdown
Collaborator Author

ByronHsu commented Sep 1, 2024

hmm the link doesn't work. can you invite me at byronhsu1230@gmail.com?

@merrymercy merrymercy merged commit 00b19f1 into sgl-project:main Sep 1, 2024
timethink pushed a commit to timethink/sglang that referenced this pull request Mar 9, 2025
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