[PR #249] Add an mps backend: weight-only INT8 and W4A8/W6A8 linears for Apple GPUs #254

New issue

Open

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

bsakaguchi commented

2026-10-26 22:22:07 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/249
Author: @roblennon
Created: 10/13/2026
Status: 🔄 Open

Base: main ← Head: mps-backend

📝 Commits (9)

1992f38 feat(mps): weight-only int8_linear backend for Apple GPUs
24913d4 feat(mps): weight-only W4A8/W6A8 linear; power-of-two row scaling for fp16
1940aa8 feat(mps): dequantize and requantize W4A8/W6A8 on Apple GPUs
3d26778 feat(mps): range-safe fused prologue and epilogue for the weight-only GEMM
34f5427 feat(mps): unmerged LoRA inside the INT8 and W4A8/W6A8 linears
fd1965c fix(mps): release chunk temporaries before the next decode
fb9d707 fix(mps): return empty contractions before reshaping; declare INT8 out_dtype
650dd6d fix(mps): keep tiny and cancelling terms through the LoRA fast path
43fe176 fix(mps): coerce lora_scale to float before dispatch; list mps in the README matrix

📊 Changes

13 files changed (+3402 additions, -41 deletions)

View changed files

📝 README.md (+70 -39)
📝 comfy_kitchen/__init__.py (+22 -0)
➕ comfy_kitchen/backends/mps/__init__.py (+106 -0)
➕ comfy_kitchen/backends/mps/linear.py (+1090 -0)
➕ comfy_kitchen/backends/mps/w4a8.py (+468 -0)
📝 comfy_kitchen/constraints.py (+38 -0)
📝 comfy_kitchen/registry.py (+4 -1)
📝 comfy_kitchen/tensor/w4a8_int8.py (+10 -1)
📝 tests/test_constraints.py (+42 -0)
➕ tests/test_mps_empty_linear.py (+38 -0)
➕ tests/test_mps_int8_linear.py (+774 -0)
➕ tests/test_mps_linear_lifetime.py (+83 -0)
➕ tests/test_mps_w4a8.py (+657 -0)

📄 Description

What

A new mps backend for Apple GPUs that implements:

int8_linear (INT8, optionally ConvRot).
w4a8_int8_linear (the AsymW4A8Int8Layout formats: 4-bit with or without a Lloyd-Max codebook, 6-bit, symmetric or asymmetric).
dequantize_w4a8_int8_weight and quantize_w4a8_int8_weight, so hosts can merge LoRAs into those weights.

int8_linear and w4a8_int8_linear also accept an unmerged LoRA (lora_down [R, K], lora_up [N, R], lora_scale). These are opt-in parameters: the registry skips a backend that does not declare them, so existing backends and callers are unchanged. The backend registers only when MPS is available and torch.mps.compile_shader exists.

Why

These formats have no fast path on a Mac today. MPS does not implement torch._int_mm, so ComfyUI disables the INT8 formats on Apple GPUs and dequantizes the weights in full on every call, and W4A8/W6A8 cannot run there at all (#92). On a 32 GB Mac that means MiniMax H3 INT8 swaps at 38 to 74 s per step, the 6-bit H3 file does not load, and Krea 2 and Qwen-Image 2.1 INT8 pay for a full bf16 dequantization per layer per call.

Design

Weight-only. Apple GPUs before M5 have no fast integer matrix path. Integer weights stay resident, a Metal kernel decodes them to fp16 in one pass, and the stock MPS fp16 GEMM runs. Activations are not quantized. ConvRot checkpoints store the weight already rotated, so the activation is rotated online (group-wise, cheap) rather than un-rotating the weight each call.
Range safety by construction. A fused prologue divides each activation row by an exact power of two so its max after rotation is at most 1, and applies the rotation in the same pass. The decode divides each weight row the same way. Every dot product is then bounded by K, so the fp16 GEMM cannot overflow for K up to 65504; wider layers compute in fp32. A fused epilogue restores both exponents with integer exponent arithmetic on the raw bits (Apple GPUs flush fp32 subnormals in arithmetic), adds bias and residual in eager's order for each format, and writes bf16 bit-exactly with torch's cast. Power-of-two scaling changes no mantissa, so it is exact at the dtype limits: near-max values, subnormals and zero rows.
fp16, not bf16. On the M1 Max the bf16 GEMM is about 1.75x slower than fp16 (70.0 against 37.8 ms at 4608x6144x6144) because bf16 is emulated, and fp16 was more accurate here (relative error 0.0004 against 0.0029). The normalization above pays for fp16's range.
Unmerged LoRA. The down weight goes through the same prologue and rides as R extra rows of the first decoded weight chunk, so the main GEMM also yields the low-rank activation. An epilogue applies the up weight in fp32 on the simdgroup matrix units. The up weight is carried as mantissa plus exponent, and base and LoRA terms are summed in a common scaled domain, so extreme scales stay finite and a LoRA that cancels the weight gives exactly 0.

Measurements

Mac Studio M1 Max 32 GB, macOS 15.7.9, PyTorch 2.14.1, October 9, 2026. bf16 in and out, median ms per call, relative L2 error against an fp32 linear over the dequantized weight in brackets.

Shape (M, K, N)	dense bf16 `F.linear`	dense fp16 `F.linear`	dequantize per call (today)	mps backend
INT8 4608, 6144, 6144	68.3 (2.3e-3)	38.5 (1.7e-3)	76.4 (2.3e-3)	39.3 (1.7e-3)
INT8 4608, 6144, 16384	183.0	102.7	201.7	104.1
INT8 4608, 16384, 6144	228.1	204.5	242.9	106.1
W4A8 codebook 4608, 16384, 6144			285.7, fp32 scales only	110.4
W6A8 4608, 16384, 6144			299.4, fp32 scales only	110.8

The backend runs 1.9 to 2.3x faster than dequantizing per call, at fp16-dense speed and accuracy. On the wide-input shape it is also about 2x faster than a dense fp16 F.linear, because dense F.linear hits a slow MPS path for transposed weights at large K (reported separately to PyTorch), which the transposed decode avoids. Eager's W4A8/W6A8 dequantize cannot run on MPS with e4m3 scales (no fp8 arithmetic), so for those files there is no path today.

Unmerged LoRA, rank 64, fused overhead per call: +2.0 ms on 4608x6144x6144 (base 38.8 ms) and +3.2 ms on 4608x16384x6144 (base 105.5 ms).

End to end in ComfyUI, seconds per step, measured with an integrated runtime that routes these formats to the backend (that ComfyUI change is a separate PR):

Model	Stock ComfyUI	With the mps backend
Krea 2 Turbo INT8 ConvRot, 512x512	8.7	3.7
Qwen-Image 2.1 INT8 ConvRot, 1024x1024	15.1	8.8
MiniMax H3 INT8, 480x288x22 frames	38 to 74 (swapping)	8.4 to 11.8
MiniMax H3 W6A8, 480x288x22 frames	does not run	6.5

Relationship to #92, #107 and #145

This fixes #92 on Apple GPUs. It is complementary to the two open fallback PRs and touches none of their files:

#145 is a portable eager float-GEMM fallback. It still matters wherever this backend cannot register (older torch without compile_shader) or declines a call, and it can land independently.
#107 adds an MPS float fallback and, in later commits, a Metal 4 kernel that quantizes activations to INT8 for M5-class GPUs. This backend is weight-only by design for M1 to M4. An activation-quantized M5 path like #107's could live inside this backend later, selected by GPU family: per-row INT8 activation quantization fused into the prologue, an int8 x int8 to int32 GEMM through Metal's tensor APIs, and dispatch by GPU family. None of that was attempted here, because there was no M5 hardware to test on.

How tested

tests/test_mps_int8_linear.py and tests/test_mps_w4a8.py compare against eager on the CPU: every format and layout; fp16, bf16 and fp32 inputs; bias and residual order per format; near-max, subnormal and zero rows; exact residual cancellation; K > 65504; strided inputs and storage offsets; empty and 1-D inputs; shader index arithmetic past 2^32; LoRA at rank tails and column tails, including cancelling and extreme-exponent LoRAs.
tests/test_constraints.py covers the opt-in parameter check. tests/test_mps_empty_linear.py covers empty contractions and the output-dtype declaration; tests/test_mps_linear_lifetime.py covers chunk temporaries at M = 1 and M = 3.
Full suite on the rebased head, M1 Max: 1082 passed, 3161 skipped (other backends), 101 s. Ruff clean. There is no macOS CI runner upstream, so these tests skip in CI.

Two deliberate differences from eager are documented and pinned by tests: asymmetric W4A8 computes from the exact stored weight (eager rounds the weight to the input dtype first, so results can differ by that rounding), and dequantized values below the fp32 normal range flush to zero, as eager's own arithmetic does on Apple GPUs. A scale restore that overflows before a cancelling bias near 2e38 returns inf exactly as eager does and is pinned by a parity test.

Limitations

Weight-only: memory is saved and integer weights are decoded at about 4 to 9% over a resident fp16 GEMM, but the matmul itself is fp16. There is no speedup beyond fp16 dense.
quantize_w4a8_int8_weight runs eager on the CPU (MPS has no fp8 arithmetic). Merging a LoRA into 4/6-bit weights is therefore slow (about 6.5 minutes for MiniMax H3) and lossy. The unmerged LoRA arguments exist so hosts can avoid that merge.
Layers with K > 65504 compute in fp32 (slower, exact range).
NVFP4, MXFP8, ConvRot W4A4 and GGUF are separate follow-up series, not part of this PR.
Tested on an M1 Max only.

Hardware tested

Mac Studio M1 Max (32-core GPU), 32 GB, macOS 15.7, PyTorch 2.14.1.

_{🔄 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/249 **Author:** [@roblennon](https://github.com/roblennon) **Created:** 10/13/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `mps-backend` --- ### 📝 Commits (9) - [`1992f38`](https://github.com/Comfy-Org/comfy-kitchen/commit/1992f383aba29dc40c246713d6eaa0cca9df5981) feat(mps): weight-only int8_linear backend for Apple GPUs - [`24913d4`](https://github.com/Comfy-Org/comfy-kitchen/commit/24913d424741923ac96dc3874e655983acd353bb) feat(mps): weight-only W4A8/W6A8 linear; power-of-two row scaling for fp16 - [`1940aa8`](https://github.com/Comfy-Org/comfy-kitchen/commit/1940aa8c901bbeff93d81b2ec11adc0d6951f516) feat(mps): dequantize and requantize W4A8/W6A8 on Apple GPUs - [`3d26778`](https://github.com/Comfy-Org/comfy-kitchen/commit/3d26778cc061f1ec7a899672bb4cdc1848e79604) feat(mps): range-safe fused prologue and epilogue for the weight-only GEMM - [`34f5427`](https://github.com/Comfy-Org/comfy-kitchen/commit/34f5427c8d50bd2fa689ac25d4e08eb1f0f43f19) feat(mps): unmerged LoRA inside the INT8 and W4A8/W6A8 linears - [`fd1965c`](https://github.com/Comfy-Org/comfy-kitchen/commit/fd1965c0c1a2992aae13e4efef623dda8123ac53) fix(mps): release chunk temporaries before the next decode - [`fb9d707`](https://github.com/Comfy-Org/comfy-kitchen/commit/fb9d707b36cc13be896412f03d9febcac347f73d) fix(mps): return empty contractions before reshaping; declare INT8 out_dtype - [`650dd6d`](https://github.com/Comfy-Org/comfy-kitchen/commit/650dd6de3fa5e8f927fa2947a8c230c3b6b689fe) fix(mps): keep tiny and cancelling terms through the LoRA fast path - [`43fe176`](https://github.com/Comfy-Org/comfy-kitchen/commit/43fe176daef525bd3ca0b4b241abe71f2796d9bb) fix(mps): coerce lora_scale to float before dispatch; list mps in the README matrix ### 📊 Changes **13 files changed** (+3402 additions, -41 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+70 -39) 📝 `comfy_kitchen/__init__.py` (+22 -0) ➕ `comfy_kitchen/backends/mps/__init__.py` (+106 -0) ➕ `comfy_kitchen/backends/mps/linear.py` (+1090 -0) ➕ `comfy_kitchen/backends/mps/w4a8.py` (+468 -0) 📝 `comfy_kitchen/constraints.py` (+38 -0) 📝 `comfy_kitchen/registry.py` (+4 -1) 📝 `comfy_kitchen/tensor/w4a8_int8.py` (+10 -1) 📝 `tests/test_constraints.py` (+42 -0) ➕ `tests/test_mps_empty_linear.py` (+38 -0) ➕ `tests/test_mps_int8_linear.py` (+774 -0) ➕ `tests/test_mps_linear_lifetime.py` (+83 -0) ➕ `tests/test_mps_w4a8.py` (+657 -0) </details> ### 📄 Description ### What A new `mps` backend for Apple GPUs that implements: - `int8_linear` (INT8, optionally ConvRot). - `w4a8_int8_linear` (the `AsymW4A8Int8Layout` formats: 4-bit with or without a Lloyd-Max codebook, 6-bit, symmetric or asymmetric). - `dequantize_w4a8_int8_weight` and `quantize_w4a8_int8_weight`, so hosts can merge LoRAs into those weights. `int8_linear` and `w4a8_int8_linear` also accept an unmerged LoRA (`lora_down` [R, K], `lora_up` [N, R], `lora_scale`). These are opt-in parameters: the registry skips a backend that does not declare them, so existing backends and callers are unchanged. The backend registers only when MPS is available and `torch.mps.compile_shader` exists. ### Why These formats have no fast path on a Mac today. MPS does not implement `torch._int_mm`, so ComfyUI disables the INT8 formats on Apple GPUs and dequantizes the weights in full on every call, and W4A8/W6A8 cannot run there at all (#92). On a 32 GB Mac that means MiniMax H3 INT8 swaps at 38 to 74 s per step, the 6-bit H3 file does not load, and Krea 2 and Qwen-Image 2.1 INT8 pay for a full bf16 dequantization per layer per call. ### Design - **Weight-only.** Apple GPUs before M5 have no fast integer matrix path. Integer weights stay resident, a Metal kernel decodes them to fp16 in one pass, and the stock MPS fp16 GEMM runs. Activations are not quantized. ConvRot checkpoints store the weight already rotated, so the activation is rotated online (group-wise, cheap) rather than un-rotating the weight each call. - **Range safety by construction.** A fused prologue divides each activation row by an exact power of two so its max after rotation is at most 1, and applies the rotation in the same pass. The decode divides each weight row the same way. Every dot product is then bounded by K, so the fp16 GEMM cannot overflow for K up to 65504; wider layers compute in fp32. A fused epilogue restores both exponents with integer exponent arithmetic on the raw bits (Apple GPUs flush fp32 subnormals in arithmetic), adds bias and residual in eager's order for each format, and writes bf16 bit-exactly with torch's cast. Power-of-two scaling changes no mantissa, so it is exact at the dtype limits: near-max values, subnormals and zero rows. - **fp16, not bf16.** On the M1 Max the bf16 GEMM is about 1.75x slower than fp16 (70.0 against 37.8 ms at 4608x6144x6144) because bf16 is emulated, and fp16 was more accurate here (relative error 0.0004 against 0.0029). The normalization above pays for fp16's range. - **Unmerged LoRA.** The down weight goes through the same prologue and rides as R extra rows of the first decoded weight chunk, so the main GEMM also yields the low-rank activation. An epilogue applies the up weight in fp32 on the simdgroup matrix units. The up weight is carried as mantissa plus exponent, and base and LoRA terms are summed in a common scaled domain, so extreme scales stay finite and a LoRA that cancels the weight gives exactly 0. ### Measurements Mac Studio M1 Max 32 GB, macOS 15.7.9, PyTorch 2.14.1, October 9, 2026. bf16 in and out, median ms per call, relative L2 error against an fp32 linear over the dequantized weight in brackets. | Shape (M, K, N) | dense bf16 `F.linear` | dense fp16 `F.linear` | dequantize per call (today) | mps backend | |---|---|---|---|---| | INT8 4608, 6144, 6144 | 68.3 (2.3e-3) | 38.5 (1.7e-3) | 76.4 (2.3e-3) | 39.3 (1.7e-3) | | INT8 4608, 6144, 16384 | 183.0 | 102.7 | 201.7 | 104.1 | | INT8 4608, 16384, 6144 | 228.1 | 204.5 | 242.9 | 106.1 | | W4A8 codebook 4608, 16384, 6144 | | | 285.7, fp32 scales only | 110.4 | | W6A8 4608, 16384, 6144 | | | 299.4, fp32 scales only | 110.8 | The backend runs 1.9 to 2.3x faster than dequantizing per call, at fp16-dense speed and accuracy. On the wide-input shape it is also about 2x faster than a dense fp16 `F.linear`, because dense `F.linear` hits a slow MPS path for transposed weights at large K (reported separately to PyTorch), which the transposed decode avoids. Eager's W4A8/W6A8 dequantize cannot run on MPS with e4m3 scales (no fp8 arithmetic), so for those files there is no path today. Unmerged LoRA, rank 64, fused overhead per call: +2.0 ms on 4608x6144x6144 (base 38.8 ms) and +3.2 ms on 4608x16384x6144 (base 105.5 ms). End to end in ComfyUI, seconds per step, measured with an integrated runtime that routes these formats to the backend (that ComfyUI change is a separate PR): | Model | Stock ComfyUI | With the mps backend | |---|---|---| | Krea 2 Turbo INT8 ConvRot, 512x512 | 8.7 | 3.7 | | Qwen-Image 2.1 INT8 ConvRot, 1024x1024 | 15.1 | 8.8 | | MiniMax H3 INT8, 480x288x22 frames | 38 to 74 (swapping) | 8.4 to 11.8 | | MiniMax H3 W6A8, 480x288x22 frames | does not run | 6.5 | ### Relationship to #92, #107 and #145 This fixes #92 on Apple GPUs. It is complementary to the two open fallback PRs and touches none of their files: - #145 is a portable eager float-GEMM fallback. It still matters wherever this backend cannot register (older torch without `compile_shader`) or declines a call, and it can land independently. - #107 adds an MPS float fallback and, in later commits, a Metal 4 kernel that quantizes activations to INT8 for M5-class GPUs. This backend is weight-only by design for M1 to M4. An activation-quantized M5 path like #107's could live inside this backend later, selected by GPU family: per-row INT8 activation quantization fused into the prologue, an int8 x int8 to int32 GEMM through Metal's tensor APIs, and dispatch by GPU family. None of that was attempted here, because there was no M5 hardware to test on. ### How tested - `tests/test_mps_int8_linear.py` and `tests/test_mps_w4a8.py` compare against eager on the CPU: every format and layout; fp16, bf16 and fp32 inputs; bias and residual order per format; near-max, subnormal and zero rows; exact residual cancellation; K > 65504; strided inputs and storage offsets; empty and 1-D inputs; shader index arithmetic past 2^32; LoRA at rank tails and column tails, including cancelling and extreme-exponent LoRAs. - `tests/test_constraints.py` covers the opt-in parameter check. `tests/test_mps_empty_linear.py` covers empty contractions and the output-dtype declaration; `tests/test_mps_linear_lifetime.py` covers chunk temporaries at M = 1 and M = 3. - Full suite on the rebased head, M1 Max: 1082 passed, 3161 skipped (other backends), 101 s. Ruff clean. There is no macOS CI runner upstream, so these tests skip in CI. Two deliberate differences from eager are documented and pinned by tests: asymmetric W4A8 computes from the exact stored weight (eager rounds the weight to the input dtype first, so results can differ by that rounding), and dequantized values below the fp32 normal range flush to zero, as eager's own arithmetic does on Apple GPUs. A scale restore that overflows before a cancelling bias near 2e38 returns inf exactly as eager does and is pinned by a parity test. ### Limitations - Weight-only: memory is saved and integer weights are decoded at about 4 to 9% over a resident fp16 GEMM, but the matmul itself is fp16. There is no speedup beyond fp16 dense. - `quantize_w4a8_int8_weight` runs eager on the CPU (MPS has no fp8 arithmetic). Merging a LoRA into 4/6-bit weights is therefore slow (about 6.5 minutes for MiniMax H3) and lossy. The unmerged LoRA arguments exist so hosts can avoid that merge. - Layers with K > 65504 compute in fp32 (slower, exact range). - NVFP4, MXFP8, ConvRot W4A4 and GGUF are separate follow-up series, not part of this PR. - Tested on an M1 Max only. ### Hardware tested Mac Studio M1 Max (32-core GPU), 32 GB, macOS 15.7, PyTorch 2.14.1. --- <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:22:07 +00:00

bsakaguchi referenced this issue

2026-10-26 22:22:11 +00:00

[PR #254] Vectorize HIP rms_adaln and split-half rms_rope #259