mirror of
https://github.com/Comfy-Org/comfy-kitchen.git
synced 2026-11-04 23:37:42 +00:00
Labels
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
No due date set.
Dependencies
No dependencies set.
Reference
Comfy-Org/comfy-kitchen#273
Loading…
Add table
Add a link
Reference in a new issue
No description provided.
Delete branch "%!s()"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
📋 Pull Request Information
Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/270
Author: @LXD-8
Created: 10/22/2026
Status: 🔄 Open
Base:
main← Head:mps-rope-adaln📝 Commits (10+)
15dd533Add Metal RoPE and AdaLN kernels for Apple GPUs (mps backend)79e38b2mps: decline on a shader compile failure, and on a missing primary scale95f6f76tests: count the shader compile attempts in the fallback teste7f0778tests: cover adaln and split-half in the mps fallback regressions0152510README: list mps first in the backend priority order8a41306README: describe the backend order as runtime-dependent, mps first when available01df9fbREADME: add mps to the backend lista9a2f37Docs: paired rope entries launch one Metal kernel per tensore40980aREADME: the 1024 head-dim limit applies to RoPE onlybd62eccREADME: list every case the mps call rules send to eager📊 Changes
8 files changed (+905 additions, -47 deletions)
View changed files
📝
README.md(+64 -47)📝
comfy_kitchen/__init__.py(+9 -0)➕
comfy_kitchen/backends/mps/__init__.py(+110 -0)➕
comfy_kitchen/backends/mps/kernels.py(+314 -0)➕
comfy_kitchen/backends/mps/norm.py(+22 -0)➕
comfy_kitchen/backends/mps/rope.py(+142 -0)➕
tests/test_mps_adaln.py(+51 -0)➕
tests/test_mps_rope.py(+193 -0)📄 Description
Based on
main@f8b7e97(v0.2.40).What
An
mpsbackend that runs, on Apple GPUs, every RoPE entry point and both AdaLN variants as one Metal kernel per tensor (a single launch per call, two for the paired(xq, xk)entries):apply_rope,apply_rope1,apply_rope_split_half,apply_rope_split_half1rms_rope,rms_rope1,rms_rope_split_half(includingrot_dim),rms_rope_split_half1adaln,rms_adalnfloat32, float16 and bfloat16; any even head dim up to 1024. Same signatures as eager.
Why
On MPS these ops run on eager today, as chains of small ops (a norm, dtype casts, broadcasts, multiply-adds), each its own launch and full pass over memory. They sit on every transformer block of the DiTs that call comfy-kitchen, so on a Mac they are a visible share of each step.
Design
norm/split_half/rot_dimas flags: one simdgroup per row, each lane owns the same pairs in every pass, so the in-place forms write intoxsafely. The adaln kernel does one threadgroup per row, two-pass mean/variance (or RMS) in fp32.(B, H, N, D), broadcastfreqs_cis/scale/shift) need no.contiguous().F.rms_normon that dtype does, so results differ from eager only by rounding (not bit-identical).trim_rope_freqsandcheck_rope_inplaceare the shared helpers eager uses.rope_unsupported/adaln_unsupportedare the single source of truth for what the kernels take, and the call rules use them. Other devices, dtypes (including non-fp32freqs_cis), a non-contiguous last dim, mixed-dtypescale/shiftor more than four non-mergeable leading dims fall through to eager, so registering this backend never changes which calls succeed.torch.mps.compile_shaderexists.Relationship to #249
#249 adds an
mpsbackend for the INT8 / W4A8 linears. This PR adds different functions to a backend of the same name and registers it the same way:comfy_kitchen/__init__.pyblock is the one from #249, line for line, except that its comment does not say "weight-only".backends/mps/__init__.pyhas the same_build_constraints()/_register()shape.kernels.py,rope.py,norm.py) do not overlap with #249's files.Whichever lands first, the other only needs to merge two capability dicts and the README matrix column. I'm happy to rebase onto #249 instead if you prefer that order. The two would also compose: #249's
int8_linearruns the same blocks whose rope and norm calls this backend takes.Measurements
MacBook Pro M5 Max (40-core GPU), 128 GB, macOS 27.2. bf16, median ms per call, eager backend against this backend:
adaln, x (1, 4139, 4096)rms_adaln, x (1, 4141, 3072)rms_ropeq+k, (1, 4096, 32, 128)rms_rope_split_halfq+k, packed qkv (1, 4141, 24, 128)apply_ropeq+k, (1, 4096, 24, 128)That is PyTorch 2.14.1. PyTorch nightly 2.16.0.dev20261022 gives the same picture (6.6x to 19.9x).
End to end in ComfyUI 0.42.0 (PyTorch nightly,
--use-pytorch-cross-attention), seconds per 1024x1024 image, this branch againstmain, nothing else changed:mainI only have an M5 Max. I have not measured M1 to M4. The kernels use plain simdgroup reductions and no M5-only features, so I expect them to run on any Apple GPU that
compile_shadersupports, but I can't vouch for the speedups there.Verification
tests/test_mps_rope.pyandtests/test_mps_adaln.py: 166 tests, skipped off MPS. Every variant across five layouts, four head dims and three dtypes, plus the in-place forms, partialrot_dim, a packed-qkv in-place call that must leavevuntouched, and the cases that must fall back to eager (including a missing scale and a shader that fails to compile). The reference is eager in fp32 on the CPU. A result passes when it is no further from that reference than eager running on MPS in the same dtype (with a small floor where both are exact). For therms_ropevariants the kernel is in fact closer to fp32 than eager on MPS, whose bf16rms_normis less precise than the CPU's (the kernel matches eager on the CPU).rms_adaln) fails 123 and 15 of these tests respectively.main770 passed / 3568 skipped / 0 failed; this branch 936 / 3568 / 0, so nothing that passed before fails now.ruff checkclean on the changed files.🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.