Skip to content

Add torch.compile support for torch.mm(out_dtype=...) #159026

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
Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Add torch.compile support for torch.mm(out_dtype=...) #159026

wants to merge 7 commits into from

Conversation

yf225
Copy link
Contributor

@yf225 yf225 commented Jul 24, 2025 •
edited
Loading

@yf225 yf225 requested review from jansel, ngimel and PaulZhang12 July 24, 2025 08:53
@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jul 24, 2025 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159026

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 381da56 with merge base ee4c5c7 (image):

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@yf225 yf225 added the topic: not user facing topic category label Jul 24, 2025
@yf225 yf225 force-pushed the mm_out_dtype_compile branch from c18ea69 to e8af708 Compare July 24, 2025 09:15
PaulZhang12
PaulZhang12 approved these changes Jul 24, 2025
View reviewed changes
Copy link
Contributor

@PaulZhang12 PaulZhang12 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

@ngimel
Copy link
Collaborator

ngimel commented Jul 24, 2025

Test failures are real, looks like those functions need to be added to aoti shim?

Skylion007
Skylion007 reviewed Jul 24, 2025
View reviewed changes
torch/_meta_registrations.py
@@ -2421,6 +2421,19 @@ def meta_mm(a, b):
return a.new_empty(N, P)


@register_meta(aten.mm.dtype)
def meta_mm_dtype(a, b, out_dtype):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't there also be check that an and b are on the same device or is that not necessary? (Seems like other metas need that check anyway)

@yf225 yf225 force-pushed the mm_out_dtype_compile branch from e8af708 to 457b296 Compare July 24, 2025 19:30
@yf225 yf225 force-pushed the mm_out_dtype_compile branch from 2a1d938 to 346c820 Compare July 25, 2025 02:49
@ngimel
Copy link
Collaborator

ngimel commented Jul 25, 2025

@desertfire can you help with aoti? Right now we are skipping a test but ideally we'd like to make sure that it's handled correctly

@yf225 yf225 force-pushed the mm_out_dtype_compile branch from 7b44e0c to 381da56 Compare July 25, 2025 20:01
@yf225
Copy link
Contributor Author

yf225 commented Jul 25, 2025 •
edited
Loading

@desertfire for more context: I tried adding "aten.mm.dtype_out": {}, to fallback_ops.py which in turn generates aoti_torch_cuda_mm_dtype_out stub. However, the stub should have been aoti_torch_cuda__mm_dtype_out_cuda to match _mm_dtype_out_cuda ATen C++ function name, hence there is a mismatch error: https://github.com/pytorch/pytorch/actions/runs/16514249751/job/46703501461?pr=159026.

Wonder is there a way to register custom C++ kernel names in the fallback_ops.py ?

desertfire
desertfire reviewed Aug 6, 2025
View reviewed changes
torch/_inductor/kernel/mm.py
@@ -570,6 +570,13 @@ def lazy_register_extern_choice(fn):

aten_mm = ExternKernelChoice(torch.mm, "at::mm_out")

aten_mm_dtype = ExternKernelChoice(
torch.mm,
"at::_mm_dtype_out_cuda",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
"at::_mm_dtype_out_cuda",
"at::mm_dtype_out",

This will make cpp_wrapper/AOTI happy.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And you need to bring back your change to fallback_ops.py.

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

@Skylion007 Skylion007 Skylion007 left review comments

@desertfire desertfire desertfire left review comments

@jansel jansel jansel approved these changes

@ngimel ngimel ngimel approved these changes

@PaulZhang12 PaulZhang12 PaulZhang12 approved these changes

Assignees
No one assigned
Labels
ciflow/inductor module: inductor release notes: inductor (aoti) topic: not user facing topic category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

torch.compile doesn't work with torch.mm(out_dtype=...) variant
6 participants
@yf225 @ngimel @jansel @Skylion007 @desertfire @PaulZhang12