Skip to content

fix: Scale QK in vision/attention #1154

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 1 commit into from
May 9, 2025
Merged

fix: Scale QK in vision/attention #1154

merged 1 commit into from
May 9, 2025

Conversation

changlan
Copy link
Contributor

@changlan changlan commented May 4, 2025

WindowAttention layer in vision/attention.py overrides the forward calls directly so we need to scale the QK before computing logits.

@changlan changlan requested review from ruomingp, markblee and a team as code owners May 4, 2025 04:22
@changlan changlan enabled auto-merge May 4, 2025 05:03
ruomingp
ruomingp reviewed May 4, 2025
View reviewed changes
Copy link
Contributor

@ruomingp ruomingp left a comment

Choose a reason for hiding this comment

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

Thanks! Please also see my comments on the internal PR.

@changlan changlan requested a review from ruomingp May 5, 2025 00:14
@changlan
Copy link
Contributor Author

changlan commented May 8, 2025

Gentle ping @ruomingp

@changlan changlan added this pull request to the merge queue May 9, 2025
Merged via the queue into apple:main with commit f72d420 May 9, 2025
6 checks passed
@changlan changlan deleted the vision-scaleqk branch May 9, 2025 21:25
Steboss pushed a commit to Steboss/axlearn that referenced this pull request May 15, 2025
@changlan @Steboss
Scale QK in vision/attention (apple#1154)
d940725
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ruomingp ruomingp ruomingp approved these changes

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

Assignees

@ruomingp ruomingp

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@changlan @ruomingp