Comfy_kitchen patch for Apple Silicon to dequantize FP8 on CPU instead of GPU (not supported as of June 10, 2026)

apply_mps_fp8_patch.py

#!/usr/bin/env python3
"""
Re-apply the Apple Silicon (MPS) FP8 patches to comfy_kitchen.

comfy_kitchen ships FP8-quantized weights that PyTorch's MPS backend cannot
cast or hold, so FP8 models (e.g. Ideogram 4) crash on Apple Silicon. This
script patches two functions:

  1. backends/eager/quantization.py : dequantize_per_tensor_fp8
       Do the fp8 -> bf16 dequant on CPU (the only place MPS allows the cast),
       preserving the per-tensor scale, then move the bf16 result back to MPS.
       -> makes the model RUN correctly.

  2. tensor/base.py : QuantizedTensor.dequantize
       Cache each weight's dequantized bf16 result on the instance, so the slow
       CPU round-trip happens once per weight instead of once per weight per
       sampling step. Guards the _version read for inference-mode tensors.
       -> makes it RUN FAST.

Both edits live in site-packages and are wiped by a `comfy_kitchen` upgrade.
Run this script (with the ComfyUI venv's python) after any update:

    /Users/roberttaylor/Documents/ComfyUI/.venv/bin/python3 apply_mps_fp8_patch.py

It is idempotent: already-patched files are detected and skipped. A timestamped
.bak copy is written next to each file before it is modified.
"""
from __future__ import annotations

import datetime
import importlib.util
import pathlib
import sys

# ---------------------------------------------------------------------------
# Patch definitions: (relative path, marker, original snippet, patched snippet)
#   marker   - substring present ONLY in the patched version (idempotency check)
#   original - exact unpatched text shipped by comfy_kitchen
#   patched  - exact replacement text
# ---------------------------------------------------------------------------

QUANT_ORIGINAL = '''def dequantize_per_tensor_fp8(
    x: torch.Tensor, scale: torch.Tensor, output_type: torch.dtype = torch.bfloat16
) -> torch.Tensor:
    dq_tensor = x.to(dtype=output_type) * scale.to(dtype=output_type)
    return dq_tensor'''

QUANT_PATCHED = '''def dequantize_per_tensor_fp8(
    x: torch.Tensor, scale: torch.Tensor, output_type: torch.dtype = torch.bfloat16
) -> torch.Tensor:
    if x.device.type == "mps":
        # MPS cannot cast or hold fp8 dtypes. Move the raw fp8 bytes to CPU
        # (device-only copy, no cast), do the whole dequant on CPU, then send
        # the finished bf16/fp16 result back to the original device.
        x_cpu = x.to(device="cpu")
        scale_cpu = scale.to(device="cpu", dtype=output_type)
        dq_tensor = (x_cpu.to(dtype=output_type) * scale_cpu).to(x.device)
    else:
        dq_tensor = x.to(dtype=output_type) * scale.to(dtype=output_type)
    return dq_tensor'''

QUANT_MARKER = 'MPS cannot cast or hold fp8 dtypes'

BASE_ORIGINAL = '''    def dequantize(self) -> torch.Tensor:
        # Ensure qdata is contiguous - backends may not handle non-contiguous views
        # (e.g., after transpose/view operations)
        qdata = self._qdata.contiguous() if not self._qdata.is_contiguous() else self._qdata

        # Check if this is a logically transposed tensor (e.g., NVFP4 with deferred transpose)
        is_transposed = getattr(self._params, "transposed", False)

        if is_transposed:
            physical_shape = (self._params.orig_shape[1], self._params.orig_shape[0])
            full = self.layout_cls.dequantize(qdata, self._params)
            if full.shape[:2] != physical_shape:
                slices = tuple(slice(0, s) for s in physical_shape)
                full = full[slices]
            return full.t()

        full = self.layout_cls.dequantize(qdata, self._params)
        orig = self._params.orig_shape
        if full.shape != orig:
            slices = tuple(slice(0, s) for s in orig)
            return full[slices]
        return full'''

BASE_PATCHED = '''    def dequantize(self) -> torch.Tensor:
        # MPS has no fp8 cast, so dequant round-trips through CPU (very slow) and
        # was being recomputed every forward pass. Model weights are static during
        # inference, so cache the dequantized result on the instance and reuse it.
        # Keyed on qdata storage pointer + version, so any in-place mutation of the
        # quantized data (offload/streaming) invalidates the cache automatically.
        use_cache = self._qdata.device.type == "mps"
        if use_cache:
            # Inference-mode tensors don't track a version counter, so guard the
            # read. When unavailable the data is immutable anyway, so data_ptr is
            # a stable key on its own.
            try:
                version = self._qdata._version
            except (RuntimeError, AttributeError):
                version = None
            cache_key = (self._qdata.data_ptr(), version)
            cached = getattr(self, "_dq_cache", None)
            if cached is not None and cached[0] == cache_key:
                return cached[1]

        # Ensure qdata is contiguous - backends may not handle non-contiguous views
        # (e.g., after transpose/view operations)
        qdata = self._qdata.contiguous() if not self._qdata.is_contiguous() else self._qdata

        # Check if this is a logically transposed tensor (e.g., NVFP4 with deferred transpose)
        is_transposed = getattr(self._params, "transposed", False)

        if is_transposed:
            physical_shape = (self._params.orig_shape[1], self._params.orig_shape[0])
            full = self.layout_cls.dequantize(qdata, self._params)
            if full.shape[:2] != physical_shape:
                slices = tuple(slice(0, s) for s in physical_shape)
                full = full[slices]
            result = full.t()
        else:
            full = self.layout_cls.dequantize(qdata, self._params)
            orig = self._params.orig_shape
            if full.shape != orig:
                slices = tuple(slice(0, s) for s in orig)
                result = full[slices]
            else:
                result = full

        if use_cache:
            self._dq_cache = (cache_key, result)
        return result'''

BASE_MARKER = '_dq_cache'

PATCHES = [
    ("backends/eager/quantization.py", QUANT_MARKER, QUANT_ORIGINAL, QUANT_PATCHED),
    ("tensor/base.py", BASE_MARKER, BASE_ORIGINAL, BASE_PATCHED),
]


def find_comfy_kitchen() -> pathlib.Path:
    spec = importlib.util.find_spec("comfy_kitchen")
    if spec and spec.submodule_search_locations:
        return pathlib.Path(list(spec.submodule_search_locations)[0])
    # Fallback to the known venv location.
    fallback = pathlib.Path(
        "/Users/roberttaylor/Documents/ComfyUI/.venv/lib/python3.12/"
        "site-packages/comfy_kitchen"
    )
    if fallback.is_dir():
        return fallback
    sys.exit(
        "ERROR: could not locate the comfy_kitchen package. Run this script with "
        "the ComfyUI venv's python:\n"
        "  /Users/roberttaylor/Documents/ComfyUI/.venv/bin/python3 "
        + pathlib.Path(__file__).name
    )


def apply_patch(path: pathlib.Path, marker: str, original: str, patched: str) -> str:
    if not path.is_file():
        return f"MISSING  {path}  (file not found — comfy_kitchen layout changed?)"

    text = path.read_text()

    if marker in text:
        return f"SKIP     {path.name}  (already patched)"

    if original not in text:
        return (
            f"FAIL     {path.name}  (original code not found — comfy_kitchen "
            "changed upstream; re-derive the patch manually)"
        )

    stamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    backup = path.with_suffix(path.suffix + f".bak-{stamp}")
    backup.write_text(text)
    path.write_text(text.replace(original, patched, 1))
    return f"PATCHED  {path.name}  (backup: {backup.name})"


def main() -> None:
    root = find_comfy_kitchen()
    print(f"comfy_kitchen: {root}\n")
    results = [apply_patch(root / rel, m, o, p) for rel, m, o, p in PATCHES]
    for line in results:
        print(" ", line)
    print()
    if any(r.startswith("FAIL") or r.startswith("MISSING") for r in results):
        print("One or more patches did NOT apply — see above. Restart ComfyUI only "
              "after resolving.")
        sys.exit(1)
    print("Done. Restart ComfyUI for changes to take effect.")


if __name__ == "__main__":
    main()