Skip to content

[WIP] [StreamK]#782

Draft
zhanglx13 wants to merge 2 commits into
main_perffrom
play_streamK
Draft

[WIP] [StreamK]#782
zhanglx13 wants to merge 2 commits into
main_perffrom
play_streamK

Conversation

@zhanglx13

Copy link
Copy Markdown

No description provided.


if STREAMK_TILES == 0:
return

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

in fact, I was trying to use STREAK_TILES to avoid this "if" with line 113 and 115, otherwise, we could use full tiles and total tiles for this if ?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

In theory, when STREAM_TILES==0, start_iter == last_iter and we should not need the while loop. However, they are not constant, so compiler cannot optimize them out. I found this if a very nice way to early return :)

else:
print("Unknown target")
exit(0)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

how do we deal with MI308 80/64 CUs ?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

idk. There must be a way (hopefully a torch API) to query the number of cus from the GPU. Let me find it.

kpack = 1
else:
print("Unknown target")
exit(0)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

same questions for MI308

loop_k -= 1

tl.assume(loop_k > 1)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

nice catch!

acc_dtype = tl.float32 if C.type.element_ty != tl.int8 else tl.int32

for tile_id in range(pid, total_full_tiles, NUM_SMS):
for tile_id in tl.range(pid, total_full_tiles, NUM_SMS, flatten=True):

@xiaohuguo2023 xiaohuguo2023 Apr 29, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

have you observed any perf improvement by enable this loop fusion ?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

no, this is just an experiment and it doesn't compile. I filed a ticket for it https://github.com/ROCm/triton-internal/issues/784

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