[PR #275] Add ck.input_act with Metal swiglu/situ for Apple GPUs (depends on #270), based on main @ `f8b7e97` (v0.2.40) #277

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/275
Author: @LXD-8
Created: 10/23/2026
Status: 🔄 Open

Base: main ← Head: mps-input-act

📝 Commits (10+)

15dd533 Add Metal RoPE and AdaLN kernels for Apple GPUs (mps backend)
79e38b2 mps: decline on a shader compile failure, and on a missing primary scale
95f6f76 tests: count the shader compile attempts in the fallback test
e7f0778 tests: cover adaln and split-half in the mps fallback regressions
0152510 README: list mps first in the backend priority order
8a41306 README: describe the backend order as runtime-dependent, mps first when available
01df9fb README: add mps to the backend list
a9a2f37 Docs: paired rope entries launch one Metal kernel per tensor
e40980a README: the 1024 head-dim limit applies to RoPE only
bd62ecc README: list every case the mps call rules send to eager

📊 Changes

13 files changed (+1174 additions, -47 deletions)

View changed files

📝 README.md (+71 -47)
📝 comfy_kitchen/__init__.py (+38 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+9 -0)
➕ comfy_kitchen/backends/eager/input_act.py (+37 -0)
➕ comfy_kitchen/backends/mps/__init__.py (+118 -0)
➕ comfy_kitchen/backends/mps/activations.py (+27 -0)
➕ comfy_kitchen/backends/mps/kernels.py (+377 -0)
➕ comfy_kitchen/backends/mps/norm.py (+22 -0)
➕ comfy_kitchen/backends/mps/rope.py (+142 -0)
➕ tests/test_input_act.py (+32 -0)
➕ tests/test_mps_adaln.py (+51 -0)
➕ tests/test_mps_input_act.py (+57 -0)
➕ tests/test_mps_rope.py (+193 -0)

📄 Description

Based on main @ f8b7e97 (v0.2.40). Depends on #270: this branch carries #270's commits, and only the last commit (Add ck.input_act ...) is new here. I'll rebase it once #270 lands.

What

ck.input_act(x, input_act, act_weight=None, act_eps=0.0): the input activations the quantized linears take as input_act (gelu_tanh, swiglu, situ, rms_norm), as a standalone, registry-dispatched op.

eager runs the existing apply_input_act unchanged, so results are bit-identical to what the quantized backends already use as their fallback.
mps runs the gated pairs (swiglu, situ) as one Metal kernel, in float32, float16 and bfloat16, with strided rows as long as the last dim is contiguous. Every other activation, dtype or layout is declined by the call rule and runs on eager.
None / "none" returns x itself without entering the custom op (a custom op may not return its input).
A fake implementation is registered (the gated pairs halve the last dim), so it works under torch.compile(fullgraph=True).

Why

The quantized linears fuse these activations into their quantizer. An unquantized linear still has to apply act(x) on its own, and today a host can only do that eagerly. In ComfyUI that is comfy/ops.py's INPUT_ACT_EAGER, which the fused MLPs (MOSS-Image's situ, Qwen-Image 2.1's swiglu) go through whenever their weights are not quantized. On MPS, situ in particular is a chain of about ten elementwise ops over the full [tokens, 2 * hidden] activation. A dispatched op lets a backend take that over, the same way the rope and norm entry points work.

The ComfyUI side is a one-line change in _eager_input_act (call ck.input_act instead of INPUT_ACT_EAGER outside training). I'll send it to ComfyUI once this is released, rather than as a version-probing branch now.

Measurements

MacBook Pro M5 Max (40-core GPU), 128 GB, macOS 27.2. bf16, median ms per call, eager backend against this backend:

Call	eager	mps	speedup
`situ`, x (1, 4141, 16384) (MOSS-Image)	4.74	0.63	7.6x
`swiglu`, x (1, 4139, 24576) (Qwen-Image 2.1)	0.99	0.59	1.7x

That is PyTorch 2.14.1; nightly 2.16.0.dev20261022 gives the same (7.5x, 1.6x). swiglu gains less because eager already does it in two ops.

End to end in ComfyUI 0.42.0 (PyTorch nightly, --use-pytorch-cross-attention), seconds per 1024x1024 image. Both columns run #270; the right one adds this PR plus the one-line ComfyUI change above:

Model	#270	#270 + this PR
MOSS-Image SFT bf16, 10 steps, CFG 2.5	24.8	21.8
Qwen-Image 2.1 Turbo bf16, 8 steps	12.5	12.4

The MOSS gain is situ; for Qwen the difference is within run-to-run noise.

I have not measured M1 to M4.

Verification

tests/test_input_act.py (runs everywhere): the eager op equals apply_input_act bit for bit for all four activations, None / "none" return the input, and torch.compile(fullgraph=True) gives the right shape and values.
tests/test_mps_input_act.py (skipped off MPS): swiglu and situ across four shapes and three dtypes, strided rows, and the calls that must stay on eager. Same criterion as #270: no further from an fp32 CPU reference than eager running on MPS. Swapping the two halves of swiglu in the kernel fails 15 of them.
Passes on PyTorch 2.14.1 and on nightly 2.16.0.dev20261022.
Full suite on this Mac: 973 passed / 3568 skipped / 0 failed (main: 770 / 3568 / 0).
ruff check clean on the changed files.

_{🔄 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/275 **Author:** [@LXD-8](https://github.com/LXD-8) **Created:** 10/23/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `mps-input-act` --- ### 📝 Commits (10+) - [`15dd533`](https://github.com/Comfy-Org/comfy-kitchen/commit/15dd533275add241957262a01ebfc230cd9c60ef) Add Metal RoPE and AdaLN kernels for Apple GPUs (mps backend) - [`79e38b2`](https://github.com/Comfy-Org/comfy-kitchen/commit/79e38b20c0ca94b0172196b8ffbf0c41577ff0da) mps: decline on a shader compile failure, and on a missing primary scale - [`95f6f76`](https://github.com/Comfy-Org/comfy-kitchen/commit/95f6f767e143ad9a74fd2c28ebd7ea602defd775) tests: count the shader compile attempts in the fallback test - [`e7f0778`](https://github.com/Comfy-Org/comfy-kitchen/commit/e7f07782227ac7c487d4bb09c5c8a1f1cf4ce0d8) tests: cover adaln and split-half in the mps fallback regressions - [`0152510`](https://github.com/Comfy-Org/comfy-kitchen/commit/015251007c95170db499ccf9f5627d2cf8d9fd31) README: list mps first in the backend priority order - [`8a41306`](https://github.com/Comfy-Org/comfy-kitchen/commit/8a413068f90e5cf7dec09179b9cfbbf42f3843e1) README: describe the backend order as runtime-dependent, mps first when available - [`01df9fb`](https://github.com/Comfy-Org/comfy-kitchen/commit/01df9fb86120c169ef07ea211e943a09d2ff9bf9) README: add mps to the backend list - [`a9a2f37`](https://github.com/Comfy-Org/comfy-kitchen/commit/a9a2f376a62094c6b2effbeb81efeb896b8cd300) Docs: paired rope entries launch one Metal kernel per tensor - [`e40980a`](https://github.com/Comfy-Org/comfy-kitchen/commit/e40980a1310cf2f6a35df18a61ac85ac7d9f0efc) README: the 1024 head-dim limit applies to RoPE only - [`bd62ecc`](https://github.com/Comfy-Org/comfy-kitchen/commit/bd62ecc4613730e80cc8cf70eb9ee7eb7c22f21a) README: list every case the mps call rules send to eager ### 📊 Changes **13 files changed** (+1174 additions, -47 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+71 -47) 📝 `comfy_kitchen/__init__.py` (+38 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+9 -0) ➕ `comfy_kitchen/backends/eager/input_act.py` (+37 -0) ➕ `comfy_kitchen/backends/mps/__init__.py` (+118 -0) ➕ `comfy_kitchen/backends/mps/activations.py` (+27 -0) ➕ `comfy_kitchen/backends/mps/kernels.py` (+377 -0) ➕ `comfy_kitchen/backends/mps/norm.py` (+22 -0) ➕ `comfy_kitchen/backends/mps/rope.py` (+142 -0) ➕ `tests/test_input_act.py` (+32 -0) ➕ `tests/test_mps_adaln.py` (+51 -0) ➕ `tests/test_mps_input_act.py` (+57 -0) ➕ `tests/test_mps_rope.py` (+193 -0) </details> ### 📄 Description Based on `main` @ `f8b7e97` (v0.2.40). **Depends on #270**: this branch carries #270's commits, and only the last commit (`Add ck.input_act ...`) is new here. I'll rebase it once #270 lands. ### What `ck.input_act(x, input_act, act_weight=None, act_eps=0.0)`: the input activations the quantized linears take as `input_act` (`gelu_tanh`, `swiglu`, `situ`, `rms_norm`), as a standalone, registry-dispatched op. - **eager** runs the existing `apply_input_act` unchanged, so results are bit-identical to what the quantized backends already use as their fallback. - **mps** runs the gated pairs (`swiglu`, `situ`) as one Metal kernel, in float32, float16 and bfloat16, with strided rows as long as the last dim is contiguous. Every other activation, dtype or layout is declined by the call rule and runs on eager. - `None` / `"none"` returns `x` itself without entering the custom op (a custom op may not return its input). - A fake implementation is registered (the gated pairs halve the last dim), so it works under `torch.compile(fullgraph=True)`. ### Why The quantized linears fuse these activations into their quantizer. An unquantized linear still has to apply `act(x)` on its own, and today a host can only do that eagerly. In ComfyUI that is `comfy/ops.py`'s `INPUT_ACT_EAGER`, which the fused MLPs (MOSS-Image's `situ`, Qwen-Image 2.1's `swiglu`) go through whenever their weights are not quantized. On MPS, `situ` in particular is a chain of about ten elementwise ops over the full `[tokens, 2 * hidden]` activation. A dispatched op lets a backend take that over, the same way the rope and norm entry points work. The ComfyUI side is a one-line change in `_eager_input_act` (call `ck.input_act` instead of `INPUT_ACT_EAGER` outside training). I'll send it to ComfyUI once this is released, rather than as a version-probing branch now. ### Measurements MacBook Pro M5 Max (40-core GPU), 128 GB, macOS 27.2. bf16, median ms per call, eager backend against this backend: | Call | eager | mps | speedup | |---|---|---|---| | `situ`, x (1, 4141, 16384) (MOSS-Image) | 4.74 | 0.63 | 7.6x | | `swiglu`, x (1, 4139, 24576) (Qwen-Image 2.1) | 0.99 | 0.59 | 1.7x | That is PyTorch 2.14.1; nightly 2.16.0.dev20261022 gives the same (7.5x, 1.6x). `swiglu` gains less because eager already does it in two ops. End to end in ComfyUI 0.42.0 (PyTorch nightly, `--use-pytorch-cross-attention`), seconds per 1024x1024 image. Both columns run #270; the right one adds this PR plus the one-line ComfyUI change above: | Model | #270 | #270 + this PR | |---|---|---| | MOSS-Image SFT bf16, 10 steps, CFG 2.5 | 24.8 | 21.8 | | Qwen-Image 2.1 Turbo bf16, 8 steps | 12.5 | 12.4 | The MOSS gain is `situ`; for Qwen the difference is within run-to-run noise. I have not measured M1 to M4. ### Verification - `tests/test_input_act.py` (runs everywhere): the eager op equals `apply_input_act` bit for bit for all four activations, `None` / `"none"` return the input, and `torch.compile(fullgraph=True)` gives the right shape and values. - `tests/test_mps_input_act.py` (skipped off MPS): `swiglu` and `situ` across four shapes and three dtypes, strided rows, and the calls that must stay on eager. Same criterion as #270: no further from an fp32 CPU reference than eager running on MPS. Swapping the two halves of `swiglu` in the kernel fails 15 of them. - Passes on PyTorch 2.14.1 and on nightly 2.16.0.dev20261022. - Full suite on this Mac: 973 passed / 3568 skipped / 0 failed (`main`: 770 / 3568 / 0). - `ruff check` clean on the changed files. --- <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:23 +00:00

bsakaguchi referenced this issue

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

[PR #277] Add FP16 tensor-core fallback for int8_linear on devices without INT8 tensor cores #280