[ONNX FE] optimize GQA decomposition path #36173
Conversation
There was a problem hiding this comment.
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
GroupQueryAttentionDecompositionwith SDPAcausal=trueusage 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. |
mitruska
left a comment
There was a problem hiding this comment.
The changes LGTM, but tests are missing.
| col_offset += k - q; | ||
| #endif | ||
| #else | ||
| col_offset += k - q; |
There was a problem hiding this comment.
Could you explain when this change is necessary? Please add tests for the issue case.
There was a problem hiding this comment.
because SDPA non-PA path with causal as true will be failed without this change. Will update the PR by adding the tests.
CuriousPanCake
left a comment
There was a problem hiding this comment.
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 Set the PR as Draft for now till I finish the tasks above. |
|
build_jenkins |
|
Open for CI-test, not ready for CPU plugin enabling yet. |
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: