Skip to content

[ONNX FE] optimize GQA decomposition path #36173

Open
Kotomi-Du wants to merge 13 commits into
openvinotoolkit:masterfrom
Kotomi-Du:gqa_mask
Open

[ONNX FE] optimize GQA decomposition path #36173
Kotomi-Du wants to merge 13 commits into
openvinotoolkit:masterfrom
Kotomi-Du:gqa_mask

Conversation

@Kotomi-Du

@Kotomi-Du Kotomi-Du commented Jun 2, 2026

Copy link
Copy Markdown
Contributor

Details:

SDPA in OV supports passing QKV information without explicit attention mask. Currently GQA decomposition manually generates causal mask which introduce unnecessary overhead (~7ms in Phi4 GPU). This PR is to replace the path with SDPA causal flag.

Besides, a conformance issue in GPU plugin is identified after enabling the path. Related fix is also applied in the PR.

Tickets:

AI Assistance:

  • AI assistance used: yes
  • AI helped review the code

@Kotomi-Du
Kotomi-Du requested review from Copilot, e-ddykim and sgbihu June 2, 2026 00:01
@Kotomi-Du
Kotomi-Du requested review from a team as code owners June 2, 2026 00:01
@github-actions github-actions Bot added category: GPU OpenVINO GPU plugin category: transformations OpenVINO Runtime library - Transformations labels Jun 2, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

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 optimizes the ONNX Frontend (GQA) decomposition by relying on ScaledDotProductAttention’s built-in causal flag instead of materializing a causal attention mask, and includes a GPU kernel fix needed for correctness after enabling that path.

Changes:

  • Replace manual causal mask construction in GroupQueryAttentionDecomposition with SDPA causal=true usage when no external attention mask is provided.
  • Adjust GPU SDPA micro-kernel causal offset calculation to fix a conformance issue in a non-paged-attention path.

Reviewed changes

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

File Description
src/common/transformations/src/transformations/op_conversions/group_query_attention_decomposition.cpp Switches no-mask GQA decomposition to SDPA causal flag path (removes explicit causal mask construction).
src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_micro.cl Fixes causal column offset computation to align masking behavior for non-paged attention.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

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 3 comments.

@Kotomi-Du
Kotomi-Du requested a review from a team as a code owner June 2, 2026 00:37
@github-actions github-actions Bot added the category: Core OpenVINO Core (aka ngraph) label Jun 2, 2026
@t-jankowski
t-jankowski requested review from mitruska and removed request for a team June 2, 2026 06:39

@mitruska mitruska left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The changes LGTM, but tests are missing.

col_offset += k - q;
#endif
#else
col_offset += k - q;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Could you explain when this change is necessary? Please add tests for the issue case.

@Kotomi-Du Kotomi-Du Jun 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

because SDPA non-PA path with causal as true will be failed without this change. Will update the PR by adding the tests.

@CuriousPanCake
CuriousPanCake requested a review from mvafin June 10, 2026 09:28

@CuriousPanCake CuriousPanCake left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Do we need tests for these changes or are they already tested?

@Kotomi-Du

Kotomi-Du commented Jun 11, 2026

Copy link
Copy Markdown
Contributor Author

Do we need tests for these changes or are they already tested?

I am preparing the tests, specifically for GPU. The onnx tests already include cases without mask. I also found even though SDPA as a general API to support is_casual as true, but it seems CPU inference failed in this case. I need to confirm if CPU plugin really supports.

Set the PR as Draft for now till I finish the tasks above.

@Kotomi-Du
Kotomi-Du marked this pull request as draft June 11, 2026 06:28
@github-actions github-actions Bot removed the category: Core OpenVINO Core (aka ngraph) label Jun 11, 2026
@Kotomi-Du

Copy link
Copy Markdown
Contributor Author

build_jenkins

@Kotomi-Du
Kotomi-Du marked this pull request as ready for review June 12, 2026 01:57
@Kotomi-Du

Copy link
Copy Markdown
Contributor Author

Open for CI-test, not ready for CPU plugin enabling yet.

@Kotomi-Du Kotomi-Du closed this Jun 15, 2026
@Kotomi-Du
Kotomi-Du deleted the gqa_mask branch June 15, 2026 20:54
@Kotomi-Du
Kotomi-Du restored the gqa_mask branch June 15, 2026 20:58
@Kotomi-Du Kotomi-Du reopened this Jun 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

category: GPU OpenVINO GPU plugin category: transformations OpenVINO Runtime library - Transformations

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants