Skip to content

[torch.compile] Custom op matching #24188

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.

Sign up for GitHub

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

Jump to bottom
Closed
wants to merge 13 commits into from
Closed

[torch.compile] Custom op matching #24188

wants to merge 13 commits into from

Conversation

@ProExpertProg
Copy link
Collaborator

@ProExpertProg ProExpertProg commented Sep 3, 2025 •
edited by github-actions bot
Loading

Fixes #23091.

Purpose

Cleanup compilation tests and custom passes and address #23091 and part of #24376. This PR is a prerequisite for matching torch implementations of custom ops in #TODO.

Cleanup:

  • RMSNorm + QuantFP8 fusion:
    • remove custom multi-output matching workaround
    • remove .instance() workaround
    • rename FusionPass to RMSNormQuantFusionPass
  • reduce test_fusion.py running time
  • add cleanup pass that topo-sorts and does DCE
  • add LazyInitPass test utility
  • add env var for pattern match debugging (TODO: docs)
  • fix test_sequence_parallelism.py
  • add assert to test_silu_mul_quant_fusion for easier diagnostics of fp4 failures.

Test Plan

Unit tests all pass.

Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
  • (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

@ProExpertProg ProExpertProg moved this from To triage to In progress in torch.compile integration Sep 3, 2025
This was referenced Sep 3, 2025
Closed
Closed
@ilmarkov ilmarkov mentioned this pull request Sep 5, 2025
Open
@mergify
Copy link

mergify bot commented Sep 6, 2025

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ProExpertProg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Sep 6, 2025
@ProExpertProg ProExpertProg force-pushed the luka/new-custom-op-matching branch from 83ee4d9 to d0f7fdb Compare September 6, 2025 18:55
@mergify mergify bot removed the needs-rebase label Sep 6, 2025
@ProExpertProg ProExpertProg mentioned this pull request Sep 6, 2025
Closed
1 task
@ProExpertProg
pass mgr util
775ca1a 
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
remove old multi_output_match workaround, add new workaround (tests p…
6544936 
…ass)

Signed-off-by: luka <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Improve workaround to properly handle inputs of inputs
9d61fa9 
Signed-off-by: luka <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Remove prints, add debug statement
bb391ce 
Signed-off-by: luka <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Reduce fusion test running time
b0703c6 
Signed-off-by: luka <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Rename FusionPass, remove instance() workaround, other fusion pass cl…
2599dd3 
…eanup

Signed-off-by: luka <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Remove old QuantMultiOutputMatch
80fdf6d 
Signed-off-by: luka <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Remove MultiOutputMatch workaround, add cleanup pass
94aa9e3 
Signed-off-by: luka <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Add LazyInitPass
75c7919 
Signed-off-by: luka <lgovedic@redhat.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Add pattern match debugging to tests
80d53cb 
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Fix seq-par test, add assert to silu-fp4 for clearer error
6dff951 
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg ProExpertProg force-pushed the luka/new-custom-op-matching branch from ea109c9 to 6dff951 Compare September 6, 2025 20:30
@mergify
Copy link

mergify bot commented Sep 6, 2025

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ProExpertProg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Sep 6, 2025
@ProExpertProg
more fusion.py cleanup
f3abbff 
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@ProExpertProg
Cleanup custom_op logging and add warnings for incorrectly specified ops
ce05cd7 
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
@github-project-automation github-project-automation bot moved this from In progress to Done in torch.compile integration Sep 9, 2025
@ProExpertProg ProExpertProg deleted the luka/new-custom-op-matching branch September 9, 2025 22:49
@mergify
Copy link

mergify bot commented Sep 9, 2025

⚠️ The sha of the head commit of this PR conflicts with #24542. Mergify cannot evaluate rules on this PR. ⚠️

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

Reviewers

@simon-mo simon-mo Awaiting requested review from simon-mo simon-mo will be requested when the pull request is marked ready for review simon-mo is a code owner

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon will be requested when the pull request is marked ready for review WoosukKwon is a code owner

@youkaichao youkaichao Awaiting requested review from youkaichao youkaichao will be requested when the pull request is marked ready for review youkaichao is a code owner

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat will be requested when the pull request is marked ready for review robertgshaw2-redhat is a code owner

@mgoin mgoin Awaiting requested review from mgoin mgoin will be requested when the pull request is marked ready for review mgoin is a code owner

@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth will be requested when the pull request is marked ready for review tlrmchlsmth is a code owner

@houseroad houseroad Awaiting requested review from houseroad houseroad will be requested when the pull request is marked ready for review houseroad is a code owner

@hmellor hmellor Awaiting requested review from hmellor hmellor will be requested when the pull request is marked ready for review hmellor is a code owner

@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 will be requested when the pull request is marked ready for review yewentao256 is a code owner

@zou3519 zou3519 Awaiting requested review from zou3519 zou3519 will be requested when the pull request is marked ready for review zou3519 is a code owner

Assignees

No one assigned

Labels

needs-rebase torch.compile

Projects

torch.compile integration
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Bug]: graph.eliminate_dead_code() break the fx graph with enable_fi_allreduce_fusion when TP == 2

1 participant

@ProExpertProg