Introduce lowbit quantized linear MPS kernels #954

manuelcandales · 2024-09-26T16:30:52Z

Summary:
The following is the directory structure of the submitted code under torchao

experimental/
├── kernels/
│   └── mps/
│       ├── codegen/
│       │   └── gen_metal_shader_lib.py
│       ├── metal/
│       │   └── (metal shaders)
│       ├── src/
│       │   └── (tensor agnostic mps kernel implementations)
│       ├── test/
│       │   └── (directly test mps kernel implementations)
│       └── metal.yaml
└── ops/
    └── mps/
        ├── register.mm
        ├── setup.py
        └── test/
            └── (test torch custom ops)

Differential Revision: D63342895

pytorch-bot · 2024-09-26T16:30:56Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 6661971 with merge base 0f6bae5 ():
💚 Looks good so far! There are no failures yet. 💚

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

facebook-github-bot · 2024-09-26T16:31:01Z

This pull request was exported from Phabricator. Differential Revision: D63342895

Summary: The following is the directory structure of the submitted code under torchao ``` experimental/ ├── kernels/ │ └── mps/ │ ├── metal/ │ │ └── (metal shaders) │ ├── src/ │ │ └── (tensor agnostic mps kernel implementations) │ └── test/ │ │ └── (directly test mps kernel implementations) └── ops/ └── mps/ ├── register.mm ├── setup.py └── test/ └── (test torch custom ops) ``` Differential Revision: D63342895

facebook-github-bot · 2024-09-26T18:31:35Z

This pull request was exported from Phabricator. Differential Revision: D63342895

Summary: The following is the directory structure of the submitted code under torchao ``` experimental/ ├── kernels/ │ └── mps/ │ ├── metal/ │ │ └── (metal shaders) │ ├── src/ │ │ └── (tensor agnostic mps kernel implementations) │ └── test/ │ │ └── (directly test mps kernel implementations) └── ops/ └── mps/ ├── register.mm ├── setup.py └── test/ └── (test torch custom ops) ``` Differential Revision: D63342895

facebook-github-bot · 2024-09-26T19:56:00Z

This pull request was exported from Phabricator. Differential Revision: D63342895

kimishpatel

Left quite a few comments. Lets address those. At high level

Figure out plan to how we gonna read .metal file. Likely similar to what torch does but see what MLX does too
Better test cases

torchao/experimental/kernels/mps/metal/divbit.metal

torchao/experimental/ops/mps/test/test_lowbit.py

kimishpatel

Left quite a few comments. Lets address those. At high level

Figure out plan to how we gonna read .metal file. Likely similar to what torch does but see what MLX does too
Better test cases

Summary: The following is the directory structure of the submitted code under torchao ``` experimental/ ├── kernels/ │ └── mps/ │ ├── metal/ │ │ └── (metal shaders) │ ├── src/ │ │ └── (tensor agnostic mps kernel implementations) │ └── test/ │ │ └── (directly test mps kernel implementations) └── ops/ └── mps/ ├── register.mm ├── setup.py └── test/ └── (test torch custom ops) ``` Differential Revision: D63342895

facebook-github-bot · 2024-10-09T23:38:58Z

This pull request was exported from Phabricator. Differential Revision: D63342895

Summary: The following is the directory structure of the submitted code under torchao ``` experimental/ ├── kernels/ │ └── mps/ │ ├── metal/ │ │ └── (metal shaders) │ ├── src/ │ │ └── (tensor agnostic mps kernel implementations) │ └── test/ │ │ └── (directly test mps kernel implementations) └── ops/ └── mps/ ├── register.mm ├── setup.py └── test/ └── (test torch custom ops) ``` Differential Revision: D63342895

facebook-github-bot · 2024-10-10T12:32:58Z

This pull request was exported from Phabricator. Differential Revision: D63342895

torchao/experimental/kernels/mps/src/OperationUtils.h

malfet · 2024-10-10T14:15:41Z

torchao/experimental/kernels/mps/src/OperationUtils.h

+    const std::string& source) {
+  NSError* error = nil;
+  MTLCompileOptions* options = [MTLCompileOptions new];
+  [options setLanguageVersion:MTLLanguageVersion3_1];


You probably want to disable fast math here as well...

torchao/experimental/kernels/mps/src/dispatch.h

malfet · 2024-10-10T14:20:44Z

torchao/experimental/ops/mps/setup.py

+    ext_modules=[
+        CppExtension(
+            name="torchao_mps_ops",
+            sources=["register.mm"],


I think you either need to package your kernels into the wheel, or using bin2c just convert them to headers. Or convert them into a metallib

Summary: The following is the directory structure of the submitted code under torchao ``` experimental/ ├── kernels/ │ └── mps/ │ ├── codegen/ │ │ └── gen_metal_shader_lib.py │ ├── metal/ │ │ └── (metal shaders) │ ├── src/ │ │ └── (tensor agnostic mps kernel implementations) │ ├── test/ │ │ └── (directly test mps kernel implementations) │ └── metal.yaml └── ops/ └── mps/ ├── register.mm ├── setup.py └── test/ └── (test torch custom ops) ``` Reviewed By: malfet Differential Revision: D63342895

facebook-github-bot · 2024-10-10T14:45:21Z

This pull request was exported from Phabricator. Differential Revision: D63342895

Summary: The following is the directory structure of the submitted code under torchao ``` experimental/ ├── kernels/ │ └── mps/ │ ├── codegen/ │ │ └── gen_metal_shader_lib.py │ ├── metal/ │ │ └── (metal shaders) │ ├── src/ │ │ └── (tensor agnostic mps kernel implementations) │ ├── test/ │ │ └── (directly test mps kernel implementations) │ └── metal.yaml └── ops/ └── mps/ ├── register.mm ├── setup.py └── test/ └── (test torch custom ops) ``` Reviewed By: malfet Differential Revision: D63342895

facebook-github-bot · 2024-10-10T15:45:08Z

This pull request was exported from Phabricator. Differential Revision: D63342895

Differential Revision: D63342895 Pull Request resolved: #954

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 26, 2024

facebook-github-bot added the fb-exported label Sep 26, 2024

manuelcandales force-pushed the export-D63342895 branch from 9ad395e to e06802b Compare September 26, 2024 18:31

manuelcandales force-pushed the export-D63342895 branch from e06802b to e783565 Compare September 26, 2024 19:55

bghira mentioned this pull request Sep 28, 2024

mps low-bit kernels from torchao huggingface/optimum-quanto#322

Closed

kimishpatel requested changes Oct 2, 2024

View reviewed changes

manuelcandales force-pushed the export-D63342895 branch from e783565 to b4943fe Compare October 9, 2024 23:38

manuelcandales force-pushed the export-D63342895 branch from b4943fe to 6c16c82 Compare October 10, 2024 12:31

malfet approved these changes Oct 10, 2024

View reviewed changes

manuelcandales force-pushed the export-D63342895 branch from 6c16c82 to 730561d Compare October 10, 2024 14:44

manuelcandales force-pushed the export-D63342895 branch from 730561d to 6661971 Compare October 10, 2024 15:45

facebook-github-bot merged commit ad184da into pytorch:main Oct 10, 2024
18 of 19 checks passed

jainapurva pushed a commit that referenced this pull request Oct 15, 2024

Introduce lowbit quantized linear MPS kernels

2761ed5

Differential Revision: D63342895 Pull Request resolved: #954

yanbing-j pushed a commit to yanbing-j/ao that referenced this pull request Dec 9, 2024

README.md: Add 3.1 to the table (pytorch#954)

60ed230

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Introduce lowbit quantized linear MPS kernels #954

Introduce lowbit quantized linear MPS kernels #954

manuelcandales commented Sep 26, 2024 •

edited

Loading

pytorch-bot bot commented Sep 26, 2024 •

edited

Loading

facebook-github-bot commented Sep 26, 2024

facebook-github-bot commented Sep 26, 2024

facebook-github-bot commented Sep 26, 2024

kimishpatel left a comment

kimishpatel left a comment

facebook-github-bot commented Oct 9, 2024

facebook-github-bot commented Oct 10, 2024

malfet Oct 10, 2024

malfet Oct 10, 2024

facebook-github-bot commented Oct 10, 2024

facebook-github-bot commented Oct 10, 2024

Introduce lowbit quantized linear MPS kernels #954

Introduce lowbit quantized linear MPS kernels #954

Conversation

manuelcandales commented Sep 26, 2024 • edited Loading

pytorch-bot bot commented Sep 26, 2024 • edited Loading

🔗 Helpful Links

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

✅ No Failures

facebook-github-bot commented Sep 26, 2024

facebook-github-bot commented Sep 26, 2024

facebook-github-bot commented Sep 26, 2024

kimishpatel left a comment

Choose a reason for hiding this comment

kimishpatel left a comment

Choose a reason for hiding this comment

facebook-github-bot commented Oct 9, 2024

facebook-github-bot commented Oct 10, 2024

malfet Oct 10, 2024

Choose a reason for hiding this comment

malfet Oct 10, 2024

Choose a reason for hiding this comment

facebook-github-bot commented Oct 10, 2024

facebook-github-bot commented Oct 10, 2024

manuelcandales commented Sep 26, 2024 •

edited

Loading

pytorch-bot bot commented Sep 26, 2024 •

edited

Loading