Skip to content

Add MPS fallback for INT8 linear - #107

Draft
ikeyan wants to merge 3 commits into
Comfy-Org:mainfrom
ikeyan:codex/mps-int8-bf16-fallback
Draft

Add MPS fallback for INT8 linear#107
ikeyan wants to merge 3 commits into
Comfy-Org:mainfrom
ikeyan:codex/mps-int8-bf16-fallback

Conversation

@ikeyan

@ikeyan ikeyan commented Aug 11, 2026

Copy link
Copy Markdown

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

@coderabbitai

coderabbitai Bot commented Aug 11, 2026

Copy link
Copy Markdown

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 42f65596-4f3d-4493-920f-9a1bfecf6504

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@github-actions

Copy link
Copy Markdown

🎉 Thank you for your contribution, we really appreciate it! 🎉

Like many open source projects, we require contributors to sign our Contributor License Agreement (CLA). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:

  • Confirm that you own your contribution.
  • Keep the right to reuse your own code.
  • Grant us a copyright license to include and share it within our projects.

CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility.

To sign, please post a new comment on this PR with exactly the following text:


I have read and agree to the Contributor License Agreement


You can retrigger this bot by commenting recheck in this Pull Request. Posted by the CLA Assistant Lite bot.

@seungjulee

Copy link
Copy Markdown

Tested this branch on two Apple Silicon generations. It works on both, including the
floating-point fallback path on pre-Metal-4 hardware.

chip RAM macOS torch ComfyUI comfy-kitchen path exercised result time
M1 Max 64 GB 26.5.1 2.13.0 0.33.0 0.2.31 + this PR float fallback works 105 s
M5 Max 128 GB 26.6.1 2.13.0 0.33.0 0.2.31 + this PR fused Metal 4 works 30 s

Workload: MiniMax H3, minimax_h3_fl2va_pruned_int8_convrot, plain t2v with no LoRA,
256x256 / 22 frames / 6 steps. Both produce valid video with 32 kHz audio.

Baseline without this PR, same machines, same workload:

comfy_kitchen/backends/eager/quantization.py:754 in fast_int8_mm
    return torch._int_mm(lhs, rhs)
NotImplementedError: The operator 'aten::_int_mm' is not currently implemented for the MPS device.

Fails in ~5 s at the first quantized linear. So this PR is the difference between "H3 INT8
does not run at all on Apple Silicon" and "it runs."

To make sure nothing else was supplying an INT8 path, ComfyUI-AppleSilicon-FP8 was moved
out of custom_nodes/ entirely for these runs and the log was asserted to contain zero
references to it. (Worth noting: a dot-prefixed directory is not skipped by ComfyUI
0.33.0, so renaming is not enough to disable a custom node.)

The M1 Max number is probably the interesting one, since it exercises the fallback rather
than the fused kernels: ~3.5x slower than M5 on the same workload. Both are usable, but
if the fp32-vs-bf16 concern raised earlier in #92 is still open, that gap is the shape of it
on real hardware.

Repro script (self-contained; backs up comfy_kitchen, applies this PR, disables the
third-party node, renders, then restores everything):
https://gist.github.com/seungjulee/328fce971c638dce600ebb1465acacd9

Separately, the two environment gotchas that cost us time and may explain other Mac reports:
H3 needs more than 36 GB of unified memory (~37 GB of weights, so a 36 GB Mac OOMs
regardless of resolution), and on some setups ffprobe from a conda/venv on PATH cannot
parse the output mp4 even when the render is perfectly fine.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants