Comfy Kitchen

Comfy Kitchen is a fast kernel library for diffusion-model inference that provides optimized low-precision compute across multiple backends. It implements FP8, NVFP4, and MXFP8 quantization, dequantization, and scaled matrix-multiply kernels, plus rotary position embedding (RoPE) operations, with eager, CUDA, and Triton implementations selected based on hardware capability.

At its core is QuantizedTensor, a torch.Tensor subclass that transparently intercepts PyTorch operations and dispatches them to optimized quantized kernels when available. This lets diffusion pipelines run in reduced precision on modern NVIDIA GPUs (Ada and Blackwell) with minimal code changes, improving throughput and memory efficiency.

What You Get

• FP8, NVFP4, and MXFP8 quantize/dequantize kernels for diffusion inference
• Scaled matrix-multiply kernels for quantized tensor-core formats
• A QuantizedTensor torch.Tensor subclass that dispatches ops to optimized kernels
• Layout classes mapping formats to Ada and Blackwell tensor-core requirements
• Multiple backends (eager, CUDA, Triton) selected by capability

Common Use Cases

• Accelerating diffusion-model inference with FP8/FP4 quantization
• Reducing GPU memory usage for large image and video generation pipelines
• Adding tensor-core quantized matmul to PyTorch models with minimal changes

Under The Hood

Architecture - The library is built around a backend-capabilities matrix: each operation (quantize/dequantize/scaled_mm/apply_rope) has eager, CUDA, and Triton implementations, and a QuantizedTensor subclass intercepts PyTorch dispatch to route ops to the optimal kernel; layout classes encode the tensor-core format and its SM-version hardware requirement.

Tech Stack - It is a PyTorch extension in Python with CUDA and Triton kernels, targeting FP8 (E4M3), NVFP4 (E2M1), and MXFP8 formats on Ada (SM 8.9) and Blackwell (SM 10.0) GPUs.

Code Quality - The project is actively developed with a rapidly growing user base and a clearly documented capability matrix, though as a young, fast-moving library some backends do not yet cover every operation.

API Design - The QuantizedTensor abstraction keeps the public surface ergonomic: users wrap tensors and continue calling standard PyTorch ops, while explicit quantize/scaled_mm functions are available for fine-grained control.