[PR #107] Add MPS fallback for INT8 linear #148

New issue

Open

opened 2026-10-26 22:20:52 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:20:52 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/107
Author: @ikeyan
Created: 8/11/2026
Status: 🔄 Open

Base: main ← Head: codex/mps-int8-bf16-fallback

📝 Commits (3)

7e189d2 Add MPS fallback for INT8 linear
a9f8906 Add fused Metal 4 INT8 linear for MPS
ae14ebb Tune fused MPS INT8 tiles for large matrices

📊 Changes

3 files changed (+543 additions, -4 deletions)

View changed files

➕ comfy_kitchen/backends/eager/mps_int8.py (+308 -0)
📝 comfy_kitchen/backends/eager/quantization.py (+82 -4)
➕ tests/test_int8_mps.py (+153 -0)

📄 Description

Summary

route eager int8_linear on MPS through a floating-point fallback instead of aten::_int_mm
keep checkpoint weights stored as INT8 and widen only the per-call weight operand to the activation dtype
preserve input activation, ConvRot, scalar/per-output-channel scaling, bias, and output dtype behavior
add MPS regression coverage that verifies the fallback never calls torch._int_mm or torch.int8_mm

Why

PyTorch does not currently implement aten::_int_mm for MPS. The eager backend advertises int8_linear on MPS, so Apple Silicon fails at the first quantized linear layer. This is the failure reported in #92.

Handling the fallback at the int8_linear level allows ComfyUI's normal BF16 execution path to remain on MPS. CPU and CUDA behavior are unchanged.

Validation

PYTHONPATH=. .venv/bin/pytest tests/test_int8_mps.py — 6 passed on Apple M5 Max / MPS
PYTHONPATH=. .venv/bin/pytest tests/test_int8.py tests/test_int8_mps.py — 42 passed, 64 skipped
Ruff lint and formatting checks pass for the changed and new code
MiniMax H3 official workflow completed 20/20 sampling steps plus VAE decode on M5 Max 48 GB with PyTorch 2.12.1; previously it failed at 0/20 with NotImplementedError: aten::_int_mm

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-kitchen/pull/107 **Author:** [@ikeyan](https://github.com/ikeyan) **Created:** 8/11/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `codex/mps-int8-bf16-fallback` --- ### 📝 Commits (3) - [`7e189d2`](https://github.com/Comfy-Org/comfy-kitchen/commit/7e189d2c0e1309c0fa7708dca7447d794123f1e9) Add MPS fallback for INT8 linear - [`a9f8906`](https://github.com/Comfy-Org/comfy-kitchen/commit/a9f89064aab77a3fa0d951b47bf119aee1a40e83) Add fused Metal 4 INT8 linear for MPS - [`ae14ebb`](https://github.com/Comfy-Org/comfy-kitchen/commit/ae14ebb2dbc9d2ecf062286031eda73d02ba9e91) Tune fused MPS INT8 tiles for large matrices ### 📊 Changes **3 files changed** (+543 additions, -4 deletions) <details> <summary>View changed files</summary> ➕ `comfy_kitchen/backends/eager/mps_int8.py` (+308 -0) 📝 `comfy_kitchen/backends/eager/quantization.py` (+82 -4) ➕ `tests/test_int8_mps.py` (+153 -0) </details> ### 📄 Description ## Summary - route eager `int8_linear` on MPS through a floating-point fallback instead of `aten::_int_mm` - keep checkpoint weights stored as INT8 and widen only the per-call weight operand to the activation dtype - preserve input activation, ConvRot, scalar/per-output-channel scaling, bias, and output dtype behavior - add MPS regression coverage that verifies the fallback never calls `torch._int_mm` or `torch.int8_mm` ## Why PyTorch does not currently implement `aten::_int_mm` for MPS. The eager backend advertises `int8_linear` on MPS, so Apple Silicon fails at the first quantized linear layer. This is the failure reported in #92. Handling the fallback at the `int8_linear` level allows ComfyUI's normal BF16 execution path to remain on MPS. CPU and CUDA behavior are unchanged. ## Validation - `PYTHONPATH=. .venv/bin/pytest tests/test_int8_mps.py` — 6 passed on Apple M5 Max / MPS - `PYTHONPATH=. .venv/bin/pytest tests/test_int8.py tests/test_int8_mps.py` — 42 passed, 64 skipped - Ruff lint and formatting checks pass for the changed and new code - MiniMax H3 official workflow completed 20/20 sampling steps plus VAE decode on M5 Max 48 GB with PyTorch 2.12.1; previously it failed at 0/20 with `NotImplementedError: aten::_int_mm` --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 22:20:52 +00:00

bsakaguchi referenced this issue

2026-10-26 22:21:15 +00:00

[PR #149] [MERGED] Revert "Build managed SageAttention wheels for CUDA 13" #176

bsakaguchi referenced this issue

2026-10-26 22:21:15 +00:00

[PR #148] [MERGED] Build managed SageAttention wheels for CUDA 13 #177