Skip to content

[dynamo exporter] Fix dynamic shapes with DynamicCache #1832

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.

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

Merged
merged 6 commits into from
Jun 19, 2025

Conversation

titaiwangms
Copy link
Contributor

@titaiwangms titaiwangms commented May 7, 2025

Describe your changes

transformers introduces "Cache" objects, which is a custom class that contains kv_cache. However, in the current exporter, this becomes an exception of the pre-condition: input should have the same tree spec as the dynamic_shapes. For example, although DynamicCache is now traceable in exporter, it does not make sense to dynamic_shapes. Therefore, we can't use input tree spec to unflatten dynamic_shapes anymore. This PR deletes that, and simply iterates the input tree spec and turn any tuple to list to overcome the format mismatch caused by config.json.

Checklist before requesting a review

  • Add unit tests for this change.
  • Make sure all tests can pass.
  • Update documents if necessary.
  • Lint and apply fixes to your code by running lintrunner -a
  • Is this a user-facing change? If yes, give a description of this change to be included in the release notes.
  • Is this PR including examples changes? If yes, please remember to update example documentation in a follow-up PR.

(Optional) Issue link

@justinchuby
Copy link
Contributor

Feel free to merge this before the IR change

@titaiwangms titaiwangms marked this pull request as ready for review May 8, 2025 20:38
@titaiwangms titaiwangms requested a review from jambayk May 8, 2025 20:39
@titaiwangms titaiwangms requested a review from justinchuby May 8, 2025 22:26
justinchuby
justinchuby previously approved these changes May 9, 2025
@titaiwangms
Copy link
Contributor Author

@jambayk I think CI is being weird. Would you help me merge it?

@xiaoyu-work
Copy link
Contributor

/azp run

Copy link

Azure Pipelines successfully started running 2 pipeline(s).

@xiaoyu-work
Copy link
Contributor

onnxscript fusion test failed: https://github.com/microsoft/Olive/blob/main/test/unit_test/passes/onnx/test_onnxscript_fusion.py#L45. How will this test change onnxscript fusion behavior?

@titaiwangms
Copy link
Contributor Author

onnxscript fusion test failed: https://github.com/microsoft/Olive/blob/main/test/unit_test/passes/onnx/test_onnxscript_fusion.py#L45. How will this test change onnxscript fusion behavior?

Thanks! It looks like TreeSpec is still wrong. I will need another fix on this.

@titaiwangms titaiwangms enabled auto-merge (squash) June 19, 2025 20:06
@titaiwangms titaiwangms merged commit 25b8357 into main Jun 19, 2025
21 checks passed
@titaiwangms titaiwangms deleted the titaiwang/fix_dynamiccache_support branch June 19, 2025 20:56
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.

3 participants