[PR #145] Add a float-GEMM fallback for int8_linear on devices without aten::_int_mm #173

Open
opened 2026-10-26 22:21:11 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/145
Author: @vqt123
Created: 8/30/2026
Status: 🔄 Open

Base: main ← Head: int8-mps-fallback


📝 Commits (3)

  • 71b99e9 Add a float-GEMM fallback for int8_linear on devices without aten::_int_mm
  • c754eb5 Chunk the fallback's float32 output scaling; keep vector inputs rank-preserving
  • 668cad7 Cast the fallback weight one output-channel slice at a time

📊 Changes

2 files changed (+295 additions, -0 deletions)

View changed files

📝 comfy_kitchen/backends/eager/quantization.py (+81 -0)
➕ tests/test_int8_fallback.py (+214 -0)

📄 Description

Problem

A fresh Comfy Desktop install on a Mac crashes out of the box on the recommended template. Install the Desktop app on Apple Silicon, open the featured MiniMax H3 text-to-video template, download its ~43 GB of models, press Run — and the run dies at sampling step 0:

NotImplementedError: The operator 'aten::_int_mm' is not currently implemented for the MPS device.

That first-run experience is this bug. H3's diffusion model ships INT8 weights (int8_convrot), and the eager int8_linear dispatches to torch._int_mm, which has no MPS kernel (pytorch/pytorch#141287) — so every INT8 checkpoint fails on every Mac, with the flagship template as the most visible case (#92; Comfy-Org/ComfyUI#15967 and Comfy-Org/ComfyUI#15133 are the same crash reported against ComfyUI). The kernel is reached through two doors and both crash: the QuantizedTensor dispatch handlers (F.linear(x, qt) / mm / addmm) and direct calls with a fused input activation (ComfyUI's linear_input_act, H3's MLP down-projection).

Fix

A float-GEMM fallback inside the eager int8_linear, engaged only when the device lacks aten::_int_mm — CUDA/CPU are trusted, other device types are probed once, so a torch that grows an MPS kernel is picked up with no code change.

  • ConvRot rotates the activations, exactly as the INT8 path does; the stored weight is already in the rotated basis, so it is never rotated or dequantized.
  • The INT8 weight values are cast to the activation dtype for the GEMM (exact: |q| ≤ 127).
  • The weight scale (scalar or per-output-channel) is applied to the output in float32.

Weights stay INT8 in memory; the only weight-sized temporary is the per-call cast to the activation dtype. Activations are not quantized on this path, so the result is slightly more accurate than the INT8 path, not less.

Verification

  • tests/test_int8_fallback.py (12 tests): both doors with torch._int_mm monkey-patched to assert the fallback never touches it; convrot × swiglu combinations; a float32 test pinning the rotate-activations math to dequantize-then-GEMM within 1e-3; the CPU native path asserted untouched; MPS-device tests (skipped off-Mac).
  • Isolated (M4 Pro 48 GB, MPS, 8192×4096×4096, ConvRot, per-channel scale): mean |diff| vs dequantize-then-GEMM 0.28 % of |y|; both are 0.9 % from the unquantized weights — the int8 quantization error itself. ~51 ms/call vs 43 ms for a weight-dequant variant, with no float32 weight-sized temporaries.
  • End to end (M4 Pro 48 GB): the MiniMax H3 T2V template went from an immediate crash to completing — 608×352 at 39 and 124 frames with this exact implementation; 864×480×124 with an earlier weight-dequant variant of the same guard.

Relationship to #107

@ikeyan's #107 takes the same fallback approach and has been stalled on the CLA since Aug 13; @seungjulee verified that branch on two Apple Silicon generations. This PR is an independent implementation (that diff was deliberately not read) so the CLA can be signed cleanly. Credit to @namikazi25 for the exact diagnosis in #92, and to @ikeyan and @seungjulee for getting there first — happy to defer to #107 if it revives. A native Metal INT8 GEMM remains the right long-term fix; this makes INT8 checkpoints work today.

Fixes #92.

🤖 Generated with Claude Code

https://claude.ai/code/session_0112gckHYByJuzZb1teQtpKk


🔄 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/145 **Author:** [@vqt123](https://github.com/vqt123) **Created:** 8/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `int8-mps-fallback` --- ### 📝 Commits (3) - [`71b99e9`](https://github.com/Comfy-Org/comfy-kitchen/commit/71b99e91635656b54e680f7ef7d59b87af4fd6d8) Add a float-GEMM fallback for int8_linear on devices without aten::_int_mm - [`c754eb5`](https://github.com/Comfy-Org/comfy-kitchen/commit/c754eb5eaa0eb919e0780a8a561c695e4c2a78be) Chunk the fallback's float32 output scaling; keep vector inputs rank-preserving - [`668cad7`](https://github.com/Comfy-Org/comfy-kitchen/commit/668cad7807d06474f8ac65a422fd36bd91a4ab6a) Cast the fallback weight one output-channel slice at a time ### 📊 Changes **2 files changed** (+295 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/eager/quantization.py` (+81 -0) ➕ `tests/test_int8_fallback.py` (+214 -0) </details> ### 📄 Description ## Problem **A fresh Comfy Desktop install on a Mac crashes out of the box on the recommended template.** Install the Desktop app on Apple Silicon, open the featured MiniMax H3 text-to-video template, download its ~43 GB of models, press Run — and the run dies at sampling step 0: ``` NotImplementedError: The operator 'aten::_int_mm' is not currently implemented for the MPS device. ``` That first-run experience is this bug. H3's diffusion model ships INT8 weights (`int8_convrot`), and the eager `int8_linear` dispatches to `torch._int_mm`, which has no MPS kernel (pytorch/pytorch#141287) — so *every* INT8 checkpoint fails on every Mac, with the flagship template as the most visible case (#92; Comfy-Org/ComfyUI#15967 and Comfy-Org/ComfyUI#15133 are the same crash reported against ComfyUI). The kernel is reached through two doors and both crash: the `QuantizedTensor` dispatch handlers (`F.linear(x, qt)` / `mm` / `addmm`) and direct calls with a fused input activation (ComfyUI's `linear_input_act`, H3's MLP down-projection). ## Fix A float-GEMM fallback inside the eager `int8_linear`, engaged only when the device lacks `aten::_int_mm` — CUDA/CPU are trusted, other device types are probed once, so a torch that grows an MPS kernel is picked up with no code change. - ConvRot rotates the **activations**, exactly as the INT8 path does; the stored weight is already in the rotated basis, so it is never rotated or dequantized. - The INT8 weight values are cast to the activation dtype for the GEMM (exact: |q| ≤ 127). - The weight scale (scalar or per-output-channel) is applied to the output in float32. Weights stay INT8 in memory; the only weight-sized temporary is the per-call cast to the activation dtype. Activations are not quantized on this path, so the result is slightly *more* accurate than the INT8 path, not less. ## Verification - `tests/test_int8_fallback.py` (12 tests): both doors with `torch._int_mm` monkey-patched to assert the fallback never touches it; convrot × swiglu combinations; a float32 test pinning the rotate-activations math to dequantize-then-GEMM within 1e-3; the CPU native path asserted untouched; MPS-device tests (skipped off-Mac). - Isolated (M4 Pro 48 GB, MPS, 8192×4096×4096, ConvRot, per-channel scale): mean |diff| vs dequantize-then-GEMM 0.28 % of |y|; both are 0.9 % from the unquantized weights — the int8 quantization error itself. ~51 ms/call vs 43 ms for a weight-dequant variant, with no float32 weight-sized temporaries. - End to end (M4 Pro 48 GB): the MiniMax H3 T2V template went from an immediate crash to completing — 608×352 at 39 and 124 frames with this exact implementation; 864×480×124 with an earlier weight-dequant variant of the same guard. ## Relationship to #107 @ikeyan's #107 takes the same fallback approach and has been stalled on the CLA since Aug 13; @seungjulee verified that branch on two Apple Silicon generations. This PR is an independent implementation (that diff was deliberately not read) so the CLA can be signed cleanly. Credit to @namikazi25 for the exact diagnosis in #92, and to @ikeyan and @seungjulee for getting there first — happy to defer to #107 if it revives. A native Metal INT8 GEMM remains the right long-term fix; this makes INT8 checkpoints work today. Fixes #92. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_0112gckHYByJuzZb1teQtpKk --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
Comfy-Org/comfy-kitchen#173
No description provided.