Ship the TorchAO kernels in the wheel - #22187
Merged
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22187
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
shoumikhin
force-pushed
the
torchao-kernels-wheel
branch
from
August 26, 2026 22:07
4c1a7d1 to
a4be612
Compare
shoumikhin
force-pushed
the
torchao-kernels-wheel
branch
from
August 26, 2026 22:28
a4be612 to
dd5f0a5
Compare
metascroy
approved these changes
Aug 26, 2026
shoumikhin
force-pushed
the
torchao-kernels-wheel
branch
from
August 26, 2026 23:15
dd5f0a5 to
865b049
Compare
The Apple framework build ships the TorchAO low-bit quantized kernels for every slice it produces, and the wheel never builds them at all. So a model quantized with those kernels runs from a Swift application and not from the wheel, on the same machine. Turn them on for the wheel on Apple Silicon, which is the only architecture they build for, and ship them as a linkable component the way the other kernel sets already are. The kernels are declared as a static archive in the TorchAO submodule, so rather than change a third-party target they are wrapped into a shared library the same way the quantized kernels are. That wrapper names the thread pool explicitly, because the kernels call into it and the wheel ships it as its own library, so this resolves against that one copy instead of bundling a second, and it carries the same runtime search path the other shipped libraries do so it can find them from wherever the package is installed. Nothing changes off Apple Silicon, and nothing changes for a build that is not producing a wheel: without the shared option the kernels stay a static archive as before. Also add the component to the release checks and to the C++ documentation, so the new library has an expected name, has its single owner asserted, and appears in the component table a consumer reads. The checks also assert the Python extension depends on every shipped library, and these kernels are excluded from that the same way the quantized ones are: torchao registers its operators itself at export time, so the extension has no reason to link them. Test Plan: Configured the wheel preset on macOS arm64 and confirmed the option is on and the wrapper target exists, then built it. The first attempt failed to link, which is what led to naming the thread pool: the kernels reference it for their parallel loops, and the archive alone does not carry it. With that dependency the library builds and records a dependency on the shipped thread pool rather than absorbing a copy of it. Verified the built library is a shared library with the unversioned name the packaging entry looks for, and that the symbol the new ownership row uses is real, visible to the tool that check runs, and defined by exactly one shipped library. Corrected after CI: the first version failed the release check asserting the Python extension depends on every shipped library, because it does not link these kernels and has no reason to. Excluded them alongside the quantized kernels, which are excluded for the same reason. Also corrected after review: the library was missing the runtime search path every other shipped library carries, so it could not have loaded from an installed wheel at all. Verified by rebuilding and loading it from a layout with only that path present, and by confirming it fails without it.
shoumikhin
force-pushed
the
torchao-kernels-wheel
branch
from
August 26, 2026 23:44
865b049 to
ca328cb
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Apple framework build ships the TorchAO low-bit quantized kernels for every
slice it produces, and the wheel never builds them at all. So a model quantized
with those kernels runs from a Swift application and not from the wheel, on the
same machine.
Turn them on for the wheel on Apple Silicon, which is the only architecture they
build for, and ship them as a linkable component the way the other kernel sets
already are.
The kernels are declared as a static archive in the TorchAO submodule, so rather
than change a third-party target they are wrapped into a shared library the same
way the quantized kernels are. That wrapper names the thread pool explicitly,
because the kernels call into it and the wheel ships it as its own library, so this
resolves against that one copy instead of bundling a second, and it carries the same
runtime search path the other shipped libraries do so it can find them from wherever
the package is installed.
Nothing changes off Apple Silicon, and nothing changes for a build that is not
producing a wheel: without the shared option the kernels stay a static archive as
before.
Also add the component to the release checks and to the C++ documentation, so the new
library has an expected name, has its single owner asserted, and appears in the
component table a consumer reads. The checks also assert the Python extension depends
on every shipped library, and these kernels are excluded from that the same way the
quantized ones are: torchao registers its operators itself at export time, so the
extension has no reason to link them.
Test Plan:
Configured the wheel preset on macOS arm64 and confirmed the option is on and the
wrapper target exists, then built it.
The first attempt failed to link, which is what led to naming the thread pool: the
kernels reference it for their parallel loops, and the archive alone does not carry
it. With that dependency the library builds and records a dependency on the shipped
thread pool rather than absorbing a copy of it.
Verified the built library is a shared library with the unversioned name the
packaging entry looks for, and that the symbol the new ownership row uses is real,
visible to the tool that check runs, and defined by exactly one shipped library.
Corrected after CI: the first version failed the release check asserting the Python
extension depends on every shipped library, because it does not link these kernels and
has no reason to. Excluded them alongside the quantized kernels, which are excluded for
the same reason.
Also corrected after review: the library was missing the runtime search path every
other shipped library carries, so it could not have loaded from an installed wheel at
all. Verified by rebuilding and loading it from a layout with only that path present,
and by confirming it fails without it.