Add cuda kernel support for GGUF inference

[Isotr0py]

What does this PR do?

• Add GGUF CUDA kernel support for GGUF forwarding
• Extend GGUF support to i-matrix quantization
• Only dequantize ops is used currently, because the MMQ/MMVQ implementation is inefficient with diffusers' 3-dimensional batching (it's designed for vLLM's contiguous batching at first)

Test Code

import torch

from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig

ckpt_path = (
    "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
)
transformer = FluxTransformer2DModel.from_single_file(
    ckpt_path,
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.float16),
    torch_dtype=torch.float16,
)
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer,
    torch_dtype=torch.float16,
)
pipe.enable_model_cpu_offload()
prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, generator=torch.manual_seed(0), num_inference_steps=50).images[0]
image.save("flux-gguf.png")

Speed comparison
Native (6.39s/it) vs CUDA kernel (5.32s/it), about 10% speed-up

# Native implementation
Loading pipeline components...:  43%|██████████████████████████████████████▌                                                   | 3/7 [00:00<00:00,  4.62it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:30<00:00, 15.29s/it]
Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:32<00:00,  4.58s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [05:19<00:00,  6.39s/it]

# CUDA Kernel
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:36<00:00, 18.50s/it]
Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:38<00:00,  5.51s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [04:25<00:00,  5.32s/it]

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

This should be fairly safe as long as we get the same values (upto some tolerance) on both the native and kernels' variants.

[Isotr0py]

Yes, the kernels' variants have been tested in vLLM's kernel tests CI compared with numpy implemented dequantize implementation. (https://github.com/vllm-project/vllm/blob/110df74332785ee749af47c5a3eb634d216b8f3b/tests/kernels/quantization/test_gguf.py#L67-L83)

[Isotr0py]

A comparision about torch.compile

Test code

import torch

from diffusers import FluxPipeline, FluxTransformer2DModel, GGUFQuantizationConfig

ckpt_path = (
    "https://huggingface.co/city96/FLUX.1-dev-gguf/blob/main/flux1-dev-Q2_K.gguf"
)
transformer = FluxTransformer2DModel.from_single_file(
    ckpt_path,
    quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
    torch_dtype=torch.bfloat16,
)
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()

pipeline.transformer.compile()

prompt = "A cat holding a sign that says hello world"
for _ in range(2):
    image = pipe(prompt, generator=torch.manual_seed(0)).images[0]
image.save("flux-gguf.png")

Outputs

## main branch
28/28 [05:07<00:00, 10.98s/it]
28/28 [03:03<00:00,  6.54s/it]

## Kernels PR
28/28 [03:27<00:00,  7.40s/it]
28/28 [02:57<00:00,  6.33s/it]

About the warmup time, the kernel PR has a significant low compilation time, but after compilation, both methods have similar speed for generation.

[DN6]

Changes look good to me 👍🏽. Thank you @Isotr0py.

For testing purposes would it be possible to configure whether to use CUDA kernels via an env variable?

And could we add a standalone test to
https://github.com/huggingface/diffusers/blob/15d50f16f2320b669c77eae2034b6612c22bd2ef/tests/quantization/gguf/test_gguf.py

that just compares native forward vs kernel forward using a dummy tensor similar to vLLM?

[DN6]

@Isotr0py Looks like there are numerical differences between native and kernel output. I added a test to check. It produces a slight difference in the generated output

Without kernels

With kernels

The text in the image is different. I've added an option to disable using kernels

[DN6]

Failing test is unrelated.