Skip to content

Implementing softmax_fa2 in partial_attn shared and causal #566

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
ksmusz merged 17 commits into from
Dec 4, 2025
Merged

Implementing softmax_fa2 in partial_attn shared and causal #566

ksmusz merged 17 commits into from
Dec 4, 2025

Conversation

@ksmusz
Copy link
Collaborator

@ksmusz ksmusz commented Nov 13, 2025

No description provided.

@ksmusz
Implementing softmax_fa2 in partial_attn shared and causal
ce31996 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
Copilot AI review requested due to automatic review settings November 13, 2025 13:56
@github-actions
Copy link

github-actions bot commented Nov 13, 2025

🚧 CI Blocked

The main CI workflow was not started for the following reason:

This is a Draft PR. Please mark it as 'Ready for Review' to trigger the CI.

Copilot AI reviewed Nov 13, 2025
View reviewed changes
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull Request Overview

This PR replaces manual softmax computation with the optimized torch.ops.hpu.softmax_fa2 operation in two attention functions (partial_attn_causal and partial_attn_shared) for improved performance on HPU devices.

Key Changes:

  • Replaced manual softmax implementation (max, exp, sum operations) with torch.ops.hpu.softmax_fa2 in both functions
  • Added initialization of inputM_hpu and inputL_hpu tensors required by the new operation
  • Consolidated bias addition into the initial matmul operation in partial_attn_shared

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@ksmusz
Attn shared debug code and tensor alignment
93eb77c 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
@github-actions
Copy link

github-actions bot commented Nov 17, 2025

🚧 CI Blocked

The main CI workflow was not started for the following reason:

This is a Draft PR. Please mark it as 'Ready for Review' to trigger the CI.

@ksmusz
Initial working solution for shared and causal attn
6ab6952 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
@github-actions
Copy link

github-actions bot commented Nov 19, 2025

🚧 CI Blocked

The main CI workflow was not started for the following reason:

This is a Draft PR. Please mark it as 'Ready for Review' to trigger the CI.

@ksmusz
Code optimizations
fef4fd2 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
@github-actions
Copy link

github-actions bot commented Nov 20, 2025

🚧 CI Blocked

The main CI workflow was not started for the following reason:

This is a Draft PR. Please mark it as 'Ready for Review' to trigger the CI.

@github-actions
Copy link

github-actions bot commented Nov 20, 2025

🚧 CI Blocked

The main CI workflow was not started for the following reason:

This is a Draft PR. Please mark it as 'Ready for Review' to trigger the CI.

@ksmusz ksmusz requested a review from Copilot November 20, 2025 15:36
@ksmusz ksmusz marked this pull request as ready for review November 20, 2025 15:36
Copilot AI reviewed Nov 20, 2025
View reviewed changes
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull Request Overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@github-actions
Copy link

github-actions bot commented Nov 24, 2025

🚧 CI Blocked

The main CI workflow was not started for the following reason:

Your branch is behind the base branch. Please merge or rebase to get the latest changes.

@ksmusz
Added checker for bf16
e47018e 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
@ksmusz ksmusz requested a review from kamil-kaczor as a code owner November 26, 2025 16:52
@github-actions
Copy link

github-actions bot commented Nov 26, 2025

🚧 CI Blocked

The main CI workflow was not started for the following reason:

Your branch is behind the base branch. Please merge or rebase to get the latest changes.

@ksmusz
Copy link
Collaborator Author

ksmusz commented Nov 26, 2025

/run-gaudi-tests

@github-actions
Copy link

github-actions bot commented Nov 26, 2025

✅ CI Passed

All checks passed successfully against the following vllm commit:
0353d2e162cbda776d9dbfe026e65303204a7f1f

@ksmusz
Copy link
Collaborator Author

ksmusz commented Nov 27, 2025

/run-gaudi-tests

@github-actions
Copy link

github-actions bot commented Nov 27, 2025

✅ CI Passed

All checks passed successfully against the following vllm commit:
0353d2e162cbda776d9dbfe026e65303204a7f1f

@github-actions
Copy link

github-actions bot commented Dec 1, 2025

✅ CI Passed

All checks passed successfully against the following vllm commit:
0353d2e162cbda776d9dbfe026e65303204a7f1f

ksmusz and others added 3 commits December 3, 2025 12:44
@ksmusz
Optimizing dummy tensor creation
be3b277 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
@kzawora-intel
Merge branch 'main' into dev/ksmusz/unified_attention_softmax_causal_…
ce0d788 
…shared
@ksmusz
Refactor last dim calc method call
1dd509e 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
@github-actions
Copy link

github-actions bot commented Dec 3, 2025

✅ CI Passed

All checks passed successfully against the following vllm commit:
3a7751485b71ce5ef927e4aa03b28602cb90811c

@ksmusz
Adding case for tensors prepare only when softmax_fa2 is enabled
aab5c21 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
@github-actions
Copy link

github-actions bot commented Dec 4, 2025

🚧 CI Blocked

The main CI workflow was not started for the following reason:

Your branch is behind the base branch. Please merge or rebase to get the latest changes.

michalkuligowski
michalkuligowski reviewed Dec 4, 2025
View reviewed changes
ksmusz added 2 commits December 4, 2025 13:00
@ksmusz
Cleaning files
3b1a4df 
Signed-off-by: Krzysztof Smusz <ksmusz@habana.ai>
@ksmusz
Merge branch 'main' into dev/ksmusz/unified_attention_softmax_causal_…
0758568 
…shared
@github-actions
Copy link

github-actions bot commented Dec 4, 2025

✅ CI Passed

All checks passed successfully against the following vllm commit:
899e2ef558e7345b99bc0d53c2e1c60ffdca7470

@ksmusz ksmusz merged commit 3fc22d0 into vllm-project:main Dec 4, 2025
44 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@michalkuligowski michalkuligowski michalkuligowski left review comments

Copilot code review Copilot Copilot left review comments

@adobrzyn adobrzyn adobrzyn approved these changes

@kzawora-intel kzawora-intel Awaiting requested review from kzawora-intel kzawora-intel is a code owner

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

@mswiniarsk mswiniarsk Awaiting requested review from mswiniarsk

@mgawarkiewicz-intel mgawarkiewicz-intel Awaiting requested review from mgawarkiewicz-intel mgawarkiewicz-intel is a code owner

@vivekgoe vivekgoe Awaiting requested review from vivekgoe

@afierka-intel afierka-intel Awaiting requested review from afierka-intel afierka-intel is a code owner

@iboiko-habana iboiko-habana Awaiting requested review from iboiko-habana iboiko-habana is a code owner

@kamil-kaczor kamil-kaczor Awaiting requested review from kamil-kaczor kamil-kaczor is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@ksmusz @michalkuligowski @adobrzyn @kzawora-intel