Skip to content

feat(0831): Localize constant ReduceSum axes inside ONNX subfunctions - #1294

Merged
vbaddi merged 10 commits into
quic:mainfrom
vbaddi:fix/einsum
Sep 11, 2026
Merged

feat(0831): Localize constant ReduceSum axes inside ONNX subfunctions#1294
vbaddi merged 10 commits into
quic:mainfrom
vbaddi:fix/einsum

Conversation

@vbaddi

@vbaddi vbaddi commented Aug 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

  • Added LocalizeFunctionReduceSumAxesTransform for ONNX subfunction exports.
  • Replaced the previous torch.einsum(... reduction ...) workarounds back to equivalent .sum(...) forms across QEff modeling/MoE/blocking code.
  • Enabled the ReduceSum axes localization only when use_onnx_subfunctions=True.

We previously used einsum in several reduction patterns to avoid ONNX subfunction export promoting constant ReduceSum axes into FunctionProto inputs. That workaround avoided compiler failures, but it could lower into heavier performance dip's due to lowering of BMM instead of the intended elementwise multiply plus ReduceSum path.

The desired PyTorch source should be able to stay as:

(query * query).sum(dim=-1, keepdim=True)
and export as:
Mul -> ReduceSum <- Constant([-1]) -- inside the ONNX function body.

ONNX Transform Details

LocalizeFunctionReduceSumAxesTransform it rewrites a function input when:

  • a node inside an ONNX FunctionProto is ReduceSum
  • the ReduceSum axes input is one of the function formal inputs
  • every top-level call site passes a compile-time constant for that formal input
  • the value is a valid integer scalar or 1-D axes tensor

When eligible, the transform:

  • inserts a local Constant node inside the function body
  • rewires each matching ReduceSum to use that local constant
  • removes the axes formal input from the function signature

@vbaddi vbaddi self-assigned this Aug 31, 2026
@vbaddi vbaddi added enhancement New feature or request 1.23 Release 1.23 Features labels Aug 31, 2026
@vbaddi

vbaddi commented Aug 31, 2026

Copy link
Copy Markdown
Contributor Author

CI-Ready

@vbaddi
vbaddi requested a review from quic-rishinr August 31, 2026 08:48
quic-rishinr
quic-rishinr approved these changes Sep 8, 2026
View reviewed changes

@quic-rishinr quic-rishinr 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.

LGTM Thanks!!

@vbaddi vbaddi mentioned this pull request Sep 8, 2026
Merged
@vbaddi

vbaddi commented Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

CI-Ready

vbaddi and others added 6 commits September 10, 2026 15:32
@vbaddi
feat(0831): Localize constant ReduceSum axes inside ONNX subfunctions
6019f2a 
  Replace QEff einsum-based reduction workarounds with equivalent torch.sum()
  forms now that ONNX subfunction ReduceSum axes are localized as constants.

  Add a narrow ONNX transform that detects ReduceSum nodes inside FunctionProto
  bodies where the axes input was promoted to a function formal input. When every
  top-level call site passes the same compile-time integer constant for that
  formal input, the transform inserts a local Constant node inside the function,
  rewires ReduceSum to use it, removes the formal function input, and removes the
  matching actual argument from each call site.

  This keeps PyTorch code such as:

      (query * query).sum(dim=-1, keepdim=True)

  lowering to Mul + ReduceSum instead of requiring an einsum workaround, while
  still presenting ReduceSum axes to the compiler as a compile-time constant
  inside the ONNX FunctionProto.

Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
@vbaddi
nit*0831): fix config for prefill in qwen3_5 moe
0c2b6f2 
Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
@kdulla @vbaddi
formatting fix
0dffe6c 
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
(cherry picked from commit 66b459b)
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
@ochougul @vbaddi
fix: return decode QPC key for CB decode-only compile
f5d5d47 
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
(cherry picked from commit c827393)
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
@ochougul @vbaddi
test: cover inferred prefill attention blocking
9fd2172 
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
@vbaddi
nit(test): disable ort tests for gemma (HQKV blocking), edge case.
0f4e5e1 
Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
Mohit Soni
fix: restore VLM CB prefill specialization selection
95834e7 
Signed-off-by: Mohit Soni <mohisoni@qti.qualcomm.com>
@mohiso22

mohiso22 commented Sep 10, 2026

Copy link
Copy Markdown
Contributor

CI-Ready

@quic-rishinr

quic-rishinr commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

CI-Ready

ochougul and others added 2 commits September 11, 2026 12:13
@ochougul
Merge branch 'main' into fix/einsum
72d9b83
@ochougul
fixing CI
59b0b41 
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
@ochougul

ochougul commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

CI-Ready

1 similar comment
@rmakar-qaic

rmakar-qaic commented Sep 11, 2026

Copy link
Copy Markdown

CI-Ready

@vbaddi
vbaddi merged commit 2dae7e3 into quic:main Sep 11, 2026
12 of 13 checks passed
@vbaddi vbaddi mentioned this pull request Sep 11, 2026
Closed
athavale-shivani added a commit to vaibverm/efficient-transformers-blocking-techniques that referenced this pull request Sep 15, 2026
@athavale-shivani
Skip ORT compare for Gemma qkv/hqkv/qkv_paged/hqkv_paged blocking modes
9307ccb 
      Extends the existing PR quic#1294 guard to the remaining
      call sites that hit the same issue: test_full_*, test_few_* (main
      loop + qkv_paged block), and test_dummy_*'s hqkv_paged block.

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>
athavale-shivani pushed a commit to vaibverm/efficient-transformers-blocking-techniques that referenced this pull request Sep 16, 2026
@vbaddi @kdulla
@ochougul Mohit Soni @athavale-shivani
feat(0831): Localize constant ReduceSum axes inside ONNX subfunctions (…
b859a0d 
…quic#1294)

## Summary

- Added `LocalizeFunctionReduceSumAxesTransform` for ONNX subfunction
exports.
- Replaced the previous `torch.einsum(... reduction ...)` workarounds
back to equivalent `.sum(...)` forms across QEff modeling/MoE/blocking
code.
- Enabled the ReduceSum axes localization only when
`use_onnx_subfunctions=True`.

We previously used `einsum` in several reduction patterns to avoid ONNX
subfunction export promoting constant `ReduceSum` axes into
`FunctionProto` inputs. That workaround avoided compiler failures, but
it could lower into heavier performance dip's due to lowering of BMM
instead of the intended elementwise multiply plus ReduceSum path.

  The desired PyTorch source should be able to stay as:

  ```python
  (query * query).sum(dim=-1, keepdim=True)
  and export as:
  Mul -> ReduceSum <- Constant([-1]) -- inside the ONNX function body.
  ```

  ## ONNX Transform Details

LocalizeFunctionReduceSumAxesTransform it rewrites a function input
when:

  - a node inside an ONNX FunctionProto is ReduceSum
  - the ReduceSum axes input is one of the function formal inputs
- every top-level call site passes a compile-time constant for that
formal input
  - the value is a valid integer scalar or 1-D axes tensor

  When eligible, the transform:

  - inserts a local Constant node inside the function body
  - rewires each matching ReduceSum to use that local constant
  - removes the axes formal input from the function signature

---------

Signed-off-by: vbaddi <vbaddi@qti.qualcomm.com>
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
Signed-off-by: ochougul <ochougul@qti.qualcomm.com>
Signed-off-by: Mohit Soni <mohisoni@qti.qualcomm.com>
Co-authored-by: Kushal Dulla <kdulla@qti.qualcomm.com>
Co-authored-by: ochougul <ochougul@qti.qualcomm.com>
Co-authored-by: Mohit Soni <mohisoni@qti.qualcomm.com>
Co-authored-by: Onkar Chougule <168134249+ochougul@users.noreply.github.com>
athavale-shivani added a commit to vaibverm/efficient-transformers-blocking-techniques that referenced this pull request Sep 16, 2026
@athavale-shivani
Skip ORT compare for Gemma qkv/hqkv/qkv_paged/hqkv_paged blocking modes
8f42bc1 
      Extends the existing PR quic#1294 guard to the remaining
      call sites that hit the same issue: test_full_*, test_few_* (main
      loop + qkv_paged block), and test_dummy_*'s hqkv_paged block.

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>
athavale-shivani added a commit to vaibverm/efficient-transformers-blocking-techniques that referenced this pull request Sep 16, 2026
@athavale-shivani
Skip ORT compare for Gemma qkv/hqkv/qkv_paged/hqkv_paged blocking modes
bbe989c 
      Extends the existing PR quic#1294 guard to the remaining
      call sites that hit the same issue: test_full_*, test_few_* (main
      loop + qkv_paged block), and test_dummy_*'s hqkv_paged block.

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>
athavale-shivani added a commit to vaibverm/efficient-transformers-blocking-techniques that referenced this pull request Sep 18, 2026
@athavale-shivani
Skip ORT compare for Gemma qkv/hqkv/qkv_paged/hqkv_paged blocking modes
28a0e19 
      Extends the existing PR quic#1294 guard to the remaining
      call sites that hit the same issue: test_full_*, test_few_* (main
      loop + qkv_paged block), and test_dummy_*'s hqkv_paged block.

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>
quic-rishinr pushed a commit to vaibverm/efficient-transformers-blocking-techniques that referenced this pull request Sep 19, 2026
@athavale-shivani @quic-rishinr
Skip ORT compare for Gemma qkv/hqkv/qkv_paged/hqkv_paged blocking modes
3f823df 
      Extends the existing PR quic#1294 guard to the remaining
      call sites that hit the same issue: test_full_*, test_few_* (main
      loop + qkv_paged block), and test_dummy_*'s hqkv_paged block.

Signed-off-by: Shivani Athavale <athavale@qti.qualcomm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@quic-rishinr quic-rishinr quic-rishinr approved these changes

Assignees

@vbaddi vbaddi

Labels

1.23 Release 1.23 Features enhancement New feature or request

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@vbaddi @mohiso22 @quic-rishinr @ochougul @rmakar-qaic @kdulla