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

[PyTorch] Adjusted the logic of MHA and DPA to enable speculative decoding #668

Merged
merged 14 commits into from
Mar 6, 2024
Merged

[PyTorch] Adjusted the logic of MHA and DPA to enable speculative decoding #668

merged 14 commits into from
Mar 6, 2024

Conversation

Oleg-Goncharov
Copy link
Collaborator

@Oleg-Goncharov Oleg-Goncharov commented Feb 15, 2024 •
edited
Loading

This PR modifies the logic of MHA and DPA by using InferenceParams (KV-cache), and extends the supported shapes of matrices in unfused softmax to rectangular to enable speculative decoding.

  • Added test cases into test_numerics.py to check if the outputs of TransformerLayer or MultiheadAttention generated in a single full run match with those generated incrementally using KV-cache
  • The longer the input sequence, the larger the divergence of the output due to accumulation of numerical errors in a forward pass

@Oleg-Goncharov Oleg-Goncharov added the enhancement New feature or request label Feb 16, 2024
@Oleg-Goncharov Oleg-Goncharov changed the title [Pytorch] Adjusted the logic of MHA and DPA to enable speculative decoding [PyTorch] Adjusted the logic of MHA and DPA to enable speculative decoding Feb 16, 2024
@Oleg-Goncharov
Copy link
Collaborator Author

Oleg-Goncharov commented Feb 16, 2024 •
edited
Loading

/te-ci

@ptrendx
Copy link
Member

ptrendx commented Feb 16, 2024

/te-ci pytorch

timmoon10
timmoon10 reviewed Feb 17, 2024
View reviewed changes
transformer_engine/pytorch/softmax.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/softmax.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/attention.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/attention.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/attention.py Show resolved Hide resolved
transformer_engine/pytorch/attention.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/attention.py Outdated Show resolved Hide resolved
transformer_engine/pytorch/attention.py Outdated Show resolved Hide resolved
@timmoon10 timmoon10 self-requested a review February 17, 2024 01:08
Oleg-Goncharov and others added 3 commits February 19, 2024 17:47
@Oleg-Goncharov @timmoon10
Apply suggestions from code review
87bc3cd 
Co-authored-by: Tim Moon <[email protected]>
Signed-off-by: Oleg Goncharov <[email protected]>
@Oleg-Goncharov @timmoon10
Apply suggestions from code review [sequence start-end]
97f315f 
Co-authored-by: Tim Moon <[email protected]>
Signed-off-by: Oleg Goncharov <[email protected]>
@Oleg-Goncharov @timmoon10
Apply suggestions from code review [inference_params offset update]]
65070c1 
Co-authored-by: Tim Moon <[email protected]>
Signed-off-by: Oleg Goncharov <[email protected]>
timmoon10
timmoon10 approved these changes Feb 26, 2024
View reviewed changes
Copy link
Collaborator

@timmoon10 timmoon10 left a comment

Choose a reason for hiding this comment

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

LGTM, pending CI.

@timmoon10
Copy link
Collaborator

/te-ci pytorch

@Oleg-Goncharov
Cleaned up the files per pylint
4f8e7a9 
Signed-off-by: Oleg Goncharov <[email protected]>
@timmoon10
Copy link
Collaborator

/te-ci pytorch

@ksivaman
Copy link
Member

/te-ci pytorch

@ksivaman
Copy link
Member

ksivaman commented Mar 6, 2024

/te-ci pytorch

@ksivaman ksivaman merged commit b459ccc into NVIDIA:main Mar 6, 2024
20 checks passed
@Oleg-Goncharov Oleg-Goncharov deleted the pr_inference_params branch March 6, 2024 20:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ptrendx ptrendx ptrendx left review comments

@timmoon10 timmoon10 timmoon10 approved these changes

Assignees
No one assigned
Labels
enhancement New feature or request
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@Oleg-Goncharov @ptrendx @timmoon10 @ksivaman