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.

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

Add support for fused update_cache #15202

Open
wants to merge 9 commits into
base: main
Choose a base branch
from
Open

Add support for fused update_cache #15202

wants to merge 9 commits into from

Conversation

kpaigwar
Copy link
Contributor

@kpaigwar kpaigwar commented Nov 19, 2024 •
edited
Loading

Problem description

For perf optimization, we need a new fused Op to update cache in parallel for two input tensors sharded on different core_grid

What's changed

  1. Added program factory for new Op-type FUSED_UPDATE
  2. Refactored validation for paged_cache_operations
  3. Added pytest for new Op
  4. fixed df reconfiguration for update_cache compute_kernel

Notes

  1. The Op is currently only functional for BFP16 inputs anf BFP16 caches. Will add support for BFP8_b caches in next PR as we fix the bug.

Checklist

@kpaigwar kpaigwar changed the title Kpaigwar/fuse update cache Support fused update cache Nov 19, 2024
@kpaigwar kpaigwar self-assigned this Nov 19, 2024
@kpaigwar kpaigwar changed the title Support fused update cache Add support for fused update_cache Nov 19, 2024
@kpaigwar kpaigwar marked this pull request as ready for review November 19, 2024 16:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ayerofieiev-tt ayerofieiev-tt Awaiting requested review from ayerofieiev-tt ayerofieiev-tt is a code owner

@dmakoviichuk-tt dmakoviichuk-tt Awaiting requested review from dmakoviichuk-tt dmakoviichuk-tt is a code owner

@yan-zaretskiy yan-zaretskiy Awaiting requested review from yan-zaretskiy yan-zaretskiy is a code owner

@rfurko-tt rfurko-tt Awaiting requested review from rfurko-tt rfurko-tt is a code owner

@cfjchu cfjchu Awaiting requested review from cfjchu cfjchu is a code owner

@TT-BrianLiu TT-BrianLiu Awaiting requested review from TT-BrianLiu TT-BrianLiu is a code owner

@razorback3 razorback3 Awaiting requested review from razorback3 razorback3 is a code owner

@dongjin-na dongjin-na Awaiting requested review from dongjin-na dongjin-na is a code owner

@cglagovichTT cglagovichTT Awaiting requested review from cglagovichTT

@uaydonat uaydonat Awaiting requested review from uaydonat

@johanna-rock-tt johanna-rock-tt Awaiting requested review from johanna-rock-tt

At least 0 approving reviews are required to merge this pull request.

Assignees

@kpaigwar kpaigwar

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@kpaigwar