# ParaAttention

    

    

大型图像和视频生成模型（例如 [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) 和 [HunyuanVideo](https://huggingface.co/tencent/HunyuanVideo)）由于体积庞大，在实时应用和部署中可能带来推理方面的挑战。

[ParaAttention](https://github.com/chengzeyi/ParaAttention) 是一个实现了 **context parallelism** 和 **first block cache** 的库，并且可以与其他技术（torch.compile、fp8 动态量化）结合使用，以加速推理。

本指南将展示如何在 NVIDIA L20 GPU 上对 FLUX.1-dev 和 HunyuanVideo 应用 ParaAttention。我们的基线基准测试未应用任何优化，但 HunyuanVideo 除外，以避免内存溢出错误。

我们的基线基准测试表明，FLUX.1-dev 能够在 28 步内以 26.36 秒生成 1024x1024 分辨率的图像，而 HunyuanVideo 能够在 30 步内以 3675.71 秒生成 720p 分辨率的 129 帧视频。

> [!TIP]
> 若需借助上下文并行实现更快的推理速度，建议（在可用的情况下）使用支持 NVLink 的 NVIDIA A100 或 H100 GPU，尤其是在 GPU 数量较多的场景下。

## 第一块缓存

缓存模型中 transformers 块的输出并在后续推理步骤中复用它们，可以降低计算成本并加快推理速度。

然而，很难确定何时复用缓存，以确保生成的图像或视频的质量。ParaAttention 直接使用 **第一个 Transformer 块输出的残差差异** 来近似模型输出之间的差异。当差异足够小时，就会复用之前推理步骤的残差差异。换句话说，即跳过去噪步骤。

这使 FLUX.1-dev 和 HunyuanVideo 的推理速度提升了 2 倍，且质量非常好。

    
    AdaCache 的工作原理，First Block Cache 是它的一个变体

要在 FLUX.1-dev 上应用第一块缓存，请按照下方所示调用 `apply_cache_on_pipe`。0.08 是 FLUX 模型的默认残差差异值。

```python
import time
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe, residual_diff_threshold=0.08)

# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

begin = time.time()
image = pipe(
    "A cat holding a sign that says hello world",
    num_inference_steps=28,
).images[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")

print("Saving image to flux.png")
image.save("flux.png")
```

| 优化 | 原始 | FBCache rdt=0.06 | FBCache rdt=0.08 | FBCache rdt=0.10 | FBCache rdt=0.12 |
| - | - | - | - | - | - |
| 预览 | ![Original](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-original.png) | ![FBCache rdt=0.06](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.06.png) | ![FBCache rdt=0.08](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.08.png) | ![FBCache rdt=0.10](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.10.png) | ![FBCache rdt=0.12](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-fbc-0.12.png) |
| 实际耗时（秒） | 26.36 | 21.83 | 17.01 | 16.00 | 13.78 |

First Block Cache 将推理速度从基线降低至 17.01 秒，即速度提升了 1.55 倍，同时几乎未造成质量损失。

要在 HunyuanVideo 上应用 First Block Cache，请按照下方所示使用 `apply_cache_on_pipe`。0.06 是 HunyuanVideo 模型的默认残差差异值。

```python
import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe, residual_diff_threshold=0.6)

pipe.vae.enable_tiling()

begin = time.time()
output = pipe(
    prompt="A cat walks on the grass, realistic",
    height=720,
    width=1280,
    num_frames=129,
    num_inference_steps=30,
).frames[0]
end = time.time()
print(f"Time: {end - begin:.2f}s")

print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
```

  
  您的浏览器不支持 video 标签。

 不使用 FBCache 的 HunyuanVideo

  
  您的浏览器不支持 video 标签。

 HunyuanVideo 与 FBCache

First Block Cache 将推理速度从基线的 2271.06 秒缩短，实现了 1.62 倍的加速，同时几乎未造成质量损失。

## fp8 量化

使用动态量化的 fp8 可进一步加速推理并降低内存占用。若要使用 8 位 [NVIDIA Tensor Cores](https://www.nvidia.com/en-us/data-center/tensor-cores/)，激活值和权重都必须进行量化。

使用 `float8_weight_only` 和 `float8_dynamic_activation_float8_weight` 对文本编码器和 transformer 模型进行量化。

默认的量化方法是按张量量化，但如果您的 GPU 支持按行量化，您也可以尝试使用它以获得更高的精度。

使用以下命令安装 [torchao](https://github.com/pytorch/ao/tree/main)。

```bash
pip3 install -U torch torchao
```

使用 [torch.compile](https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) 并设置 `mode="max-autotune-no-cudagraphs"` 或 `mode="max-autotune"` 可以选择性能最佳的 kernel。如果模型是首次被调用，编译过程可能会花费较长时间，但一旦模型完成编译，其收益是值得的。

此示例仅对 transformer 模型进行量化，但您也可以对文本编码器进行量化，以进一步降低内存占用。

> [!TIP]
> 动态量化可能会显著改变模型输出的分布，因此你需要将 `residual_diff_threshold` 更改为更大的值，使其生效。

```python
import time
import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(
    pipe,
    residual_diff_threshold=0.12,  # Use a larger value to make the cache take effect
)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

for i in range(2):
    begin = time.time()
    image = pipe(
        "A cat holding a sign that says hello world",
        num_inference_steps=28,
    ).images[0]
    end = time.time()
    if i == 0:
        print(f"Warm up time: {end - begin:.2f}s")
    else:
        print(f"Time: {end - begin:.2f}s")

print("Saving image to flux.png")
image.save("flux.png")
```

fp8 动态量化和 torch.compile 将推理速度降低至 7.56 秒，相较于基线快 3.48 倍。

```python
import time
import torch
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
pipe.vae.enable_tiling()
# pipe.enable_model_cpu_offload()
# pipe.enable_sequential_cpu_offload()

for i in range(2):
    begin = time.time()
    output = pipe(
        prompt="A cat walks on the grass, realistic",
        height=720,
        width=1280,
        num_frames=129,
        num_inference_steps=1 if i == 0 else 30,
    ).frames[0]
    end = time.time()
    if i == 0:
        print(f"Warm up time: {end - begin:.2f}s")
    else:
        print(f"Time: {end - begin:.2f}s")

print("Saving video to hunyuan_video.mp4")
export_to_video(output, "hunyuan_video.mp4", fps=15)
```

NVIDIA L20 GPU 仅有 48GB 显存，在编译后，如果未调用 `enable_model_cpu_offload`，可能会遇到内存溢出（OOM）错误，因为 HunyuanVideo 在高分辨率和大量帧数下运行时会产生非常大的激活张量。对于显存小于 80GB 的 GPU，可以尝试降低分辨率和帧数以避免 OOM 错误。

大型视频生成模型通常受限于注意力计算，而非全连接层。这些模型从量化和 torch.compile 中获得的收益并不显著。

## 上下文并行

Context Parallelism 对推理进行并行化，并可通过多块 GPU 进行扩展。ParaAttention 的组合式设计允许您将 Context Parallelism 与 First Block Cache 以及动态量化相结合。

> [!TIP]
> 有关如何使用多 GPU 扩展推理的详细说明和示例，请参阅 [ParaAttention](https://github.com/chengzeyi/ParaAttention/tree/main) 仓库。

如果推理过程需要持久化并可对外提供服务，建议使用 [torch.multiprocessing](https://pytorch.org/docs/stable/multiprocessing.html) 编写自定义的推理处理器。这样可以消除启动进程以及加载和重新编译模型所带来的开销。

下面的代码示例结合了 First Block Cache、fp8 动态量化、torch.compile 和 Context Parallelism，以实现最快的推理速度。

```python
import time
import torch
import torch.distributed as dist
from diffusers import FluxPipeline

dist.init_process_group()

torch.cuda.set_device(dist.get_rank())

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae

mesh = init_context_parallel_mesh(
    pipe.device.type,
    max_ring_dim_size=2,
)
parallelize_pipe(
    pipe,
    mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(
    pipe,
    residual_diff_threshold=0.12,  # Use a larger value to make the cache take effect
)

from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only

quantize_(pipe.text_encoder, float8_weight_only())
quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
torch._inductor.config.reorder_for_compute_comm_overlap = True
pipe.transformer = torch.compile(
   pipe.transformer, mode="max-autotune-no-cudagraphs",
)

# Enable memory savings
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())

for i in range(2):
    begin = time.time()
    image = pipe(
        "A cat holding a sign that says hello world",
        num_inference_steps=28,
        output_type="pil" if dist.get_rank() == 0 else "pt",
    ).images[0]
    end = time.time()
    if dist.get_rank() == 0:
        if i == 0:
            print(f"Warm up time: {end - begin:.2f}s")
        else:
            print(f"Time: {end - begin:.2f}s")

if dist.get_rank() == 0:
    print("Saving image to flux.png")
    image.save("flux.png")

dist.destroy_process_group()
```

将其保存为 `run_flux.py`，并使用 [torchrun](https://pytorch.org/docs/stable/elastic/run.html) 启动。

```bash
# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=2 run_flux.py
```

与基线相比，使用 2 块 NVIDIA L20 GPU 时，推理速度降至 8.20 秒，即快 3.21 倍。使用 4 块 L20 时，推理速度为 3.90 秒，即快 6.75 倍。

下面的代码示例结合了 First Block Cache 和 Context Parallelism，以实现最快的推理速度。

```python
import time
import torch
import torch.distributed as dist
from diffusers import HunyuanVideoPipeline, HunyuanVideoTransformer3DModel
from diffusers.utils import export_to_video

dist.init_process_group()

torch.cuda.set_device(dist.get_rank())

model_id = "tencent/HunyuanVideo"
transformer = HunyuanVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    dtype=torch.bfloat16,
    revision="refs/pr/18",
)
pipe = HunyuanVideoPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.float16,
    revision="refs/pr/18",
).to("cuda")  # or "mps", "xpu", "cpu"

from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae

mesh = init_context_parallel_mesh(
    pipe.device.type,
)
parallelize_pipe(
    pipe,
    mesh=mesh,
)
parallelize_vae(pipe.vae, mesh=mesh._flatten())

from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe

apply_cache_on_pipe(pipe)

# from torchao.quantization import quantize_, float8_dynamic_activation_float8_weight, float8_weight_only
#
# torch._inductor.config.reorder_for_compute_comm_overlap = True
#
# quantize_(pipe.text_encoder, float8_weight_only())
# quantize_(pipe.transformer, float8_dynamic_activation_float8_weight())
# pipe.transformer = torch.compile(
#    pipe.transformer, mode="max-autotune-no-cudagraphs",
# )

# Enable memory savings
pipe.vae.enable_tiling()
# pipe.enable_model_cpu_offload(gpu_id=dist.get_rank())
# pipe.enable_sequential_cpu_offload(gpu_id=dist.get_rank())

for i in range(2):
    begin = time.time()
    output = pipe(
        prompt="A cat walks on the grass, realistic",
        height=720,
        width=1280,
        num_frames=129,
        num_inference_steps=1 if i == 0 else 30,
        output_type="pil" if dist.get_rank() == 0 else "pt",
    ).frames[0]
    end = time.time()
    if dist.get_rank() == 0:
        if i == 0:
            print(f"Warm up time: {end - begin:.2f}s")
        else:
            print(f"Time: {end - begin:.2f}s")

if dist.get_rank() == 0:
    print("Saving video to hunyuan_video.mp4")
    export_to_video(output, "hunyuan_video.mp4", fps=15)

dist.destroy_process_group()
```

将其保存为 `run_hunyuan_video.py`，并使用 [torchrun](https://pytorch.org/docs/stable/elastic/run.html) 启动。

```bash
# Use --nproc_per_node to specify the number of GPUs
torchrun --nproc_per_node=8 run_hunyuan_video.py
```

与基线相比，使用 8 块 NVIDIA L20 GPU 时，推理速度降低至 649.23 秒，即快了 5.66 倍。

## 基准测试

| GPU 类型 | GPU 数量 | 优化 | 实际耗时（秒） | 加速 |
| - | - | - | - | - |
| NVIDIA L20 | 1 | 基线 | 26.36 | 1.00x |
| NVIDIA L20 | 1 | FBCache (rdt=0.08) | 17.01 | 1.55x |
| NVIDIA L20 | 1 | FP8 DQ | 13.40 | 1.96x |
| NVIDIA L20 | 1 | FBCache (rdt=0.12) + FP8 DQ | 7.56 | 3.48x |
| NVIDIA L20 | 2 | FBCache (rdt=0.12) + FP8 DQ + CP | 4.92 | 5.35x |
| NVIDIA L20 | 4 | FBCache (rdt=0.12) + FP8 DQ + CP | 3.90 | 6.75x |

| GPU 类型 | GPU 数量 | 优化 | 实际耗时（秒） | 加速 |
| - | - | - | - | - |
| NVIDIA L20 | 1 | 基线 | 3675.71 | 1.00x |
| NVIDIA L20 | 1 | FBCache | 2271.06 | 1.62x |
| NVIDIA L20 | 2 | FBCache + CP | 1132.90 | 3.24x |
| NVIDIA L20 | 4 | FBCache + CP | 718.15 | 5.12x |
| NVIDIA L20 | 8 | FBCache + CP | 649.23 | 5.66x |