Skip to content

feat(onnxslim_pass):Add split_to_slice pass test and update qkv_split…#43

Open
co-seven wants to merge 1 commit into
spacemit-com:mainfrom
co-seven:main
Open

feat(onnxslim_pass):Add split_to_slice pass test and update qkv_split…#43
co-seven wants to merge 1 commit into
spacemit-com:mainfrom
co-seven:main

Conversation

@co-seven

Copy link
Copy Markdown
Contributor

… pass

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 extends xslim.onnxslim_pass with two new ONNX graph rewrites: (1) a Split -> Slice rewrite when only a single split output is actually used, and (2) an additional QKV MatMul fusion pattern that recognizes a MatMul -> Add -> Reshape -> Transpose -> Split -> Squeeze layout and rewrites it into three independent Q/K/V branches. It also adds/updates unit tests to validate the new fusion and rewrite behaviors.

Changes:

  • Add a new fusion pass (split_to_slice) that rewrites eligible Split nodes into a single Slice.
  • Add a new QKV fusion matcher for MatMul/Add/Reshape/Transpose/Split/Squeeze patterns (contiguous Q/K/V weight layout).
  • Add tests covering the new QKV fusion and the Split -> Slice rewrite/keep behavior.

Reviewed changes

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

File Description
tests/test_onnxslim_pass.py Adds models + assertions to validate the new QKV fusion pattern and Split -> Slice rewrite behavior.
src/xslim/onnxslim_pass/split_to_slice.py Introduces the SingleConsumerSplitToSlicePatternMatcher and rewrite logic.
src/xslim/onnxslim_pass/qkv_split.py Adds a new matcher that fuses the reshape/transpose/split/squeeze QKV pattern into 3 branches.
src/xslim/onnxslim_pass/__init__.py Ensures the new split_to_slice pass is imported/registered.

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants