用 ParaAttention 加速 Diffusers 推理First Block Cache、fp8 量化与上下文并行实战指南【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers大型图像与视频生成模型如 FLUX.1-dev、HunyuanVideo凭借出色的生成质量成为业界焦点但其庞大的参数量与计算开销给实时应用和线上部署带来了严峻的推理挑战。本文以 Diffusers 官方文档 docs/source/en/optimization/para_attn.md 为核心系统讲解 ParaAttention 库的三项关键技术——上下文并行Context Parallelism、First Block Cache首块缓存FBCache与fp8 动态量化并结合当前仓库源码src/diffusers/hooks/first_block_cache.py、src/diffusers/models/cache_utils.py深入剖析其底层原理。读完本文你将掌握如何在 NVIDIA L20 GPU 上对 FLUX.1-dev 与 HunyuanVideo 组合运用上述优化手段实现最高约 6.75 倍的端到端推理加速。ParaAttention 是什么ParaAttention 是一个面向扩散模型的推理加速库核心实现了两类优化上下文并行Context Parallelism将长序列图像 token、视频帧 token切分到多张 GPU 上并行计算随 GPU 数量线性扩展吞吐First Block Cache首块缓存复用 transformer 首块的残差输出直接跳过部分去噪步减少计算量。此外ParaAttention 采用组合式设计可以进一步与torch.compile、fp8 动态量化等技术叠加使用获得更大的加速比。本文的基线基准测试baseline默认不做任何优化唯一例外是 HunyuanVideo 为了避免显存溢出OOM而保留了必要的显存管理设置。在 1 张 NVIDIA L20 GPU 上基线测试结果为FLUX.1-dev 生成一张 1024×1024 图像28 步耗时26.36 秒HunyuanVideo 生成 129 帧 720p 视频30 步耗时3675.71 秒。下面的优化方案都以这一基线作为对比参照。[!TIP] 如果想要更快的上下文并行推理在条件允许时优先选用带 NVLink 互联的 NVIDIA A100 或 H100 GPU尤其是在 GPU 数量较多的场景下NVLink 能显著降低跨卡通信开销。First Block Cache用残差相似度跳过冗余去噪步工作原理在扩散模型的迭代去噪过程中相邻时间步之间模型输出的变化往往是平滑、渐进的。First Block Cache 的思想是缓存模型 transformer 各 block 的输出并在后续推理步骤中直接复用从而省去部分前向计算、加速推理。难点在于如何判断何时可以安全复用缓存而不损害生成质量。ParaAttention 给出的方案非常巧妙直接用第一个 transformer block 输出的残差差值来近似整个模型输出的差异。当两次去噪步之间首块输出的残差变化足够小时就认为模型整体输出变化可忽略于是复用之前推理步的残差跳过整个去噪步。据文档所述该方法在 FLUX.1-dev 与 HunyuanVideo 上可带来约 2 倍的速度提升同时保持很好的生成质量。仓库源码中的实现细节First Block Cache 并非黑盒魔法在 Diffusers 仓库中可以找到它完整的钩子hook实现见 src/diffusers/hooks/first_block_cache.pyFirstBlockCacheConfigL32-L48核心参数threshold默认值0.05。阈值越高越倾向于跳过更多层的前向计算、推理越快但生成质量可能下降阈值越低则加速有限。判断依据是首块输出残差与缓存残差之间的absmean绝对均值差值当差值低于阈值时跳过前向。FBCHeadBlockHookL65-L142挂在第一个transformer block 上。每次前向计算首块的真实输出残差并通过_should_compute_remaining_blocks比较当前残差与上一缓存残差若相对差异diff (absmean / prev_hidden_states_absmean)大于阈值则继续计算其余 block否则直接复用缓存的尾块残差拼出最终输出。FBCBlockHookL145-L190挂在其余所有 block含尾块上。当should_compute为 False 时各 block 直接返回输入恒等映射从而彻底跳过计算。apply_first_block_cacheL193-L246遍历模型的named_children()利用 src/diffusers/hooks/_common.py 中定义的_ALL_TRANSFORMER_BLOCK_IDENTIFIERS涵盖blocks、transformer_blocks、single_transformer_blocks、layers、temporal_transformer_blocks等常见命名识别 transformer block 列表随后对头块、中间块、尾块分别挂载对应钩子。也就是说Diffusers 本身也原生集成了 First Block Cache 机制。除了通过para_attn的apply_cache_on_pipe一键启用外还可以直接用 apply_first_block_cache 配合FirstBlockCacheConfig(threshold...)作用到任意支持的 transformer 上模型层面则通过 CacheMixin 提供的enable_cache(config)/disable_cache()统一管理缓存开关同一时刻仅允许启用一种缓存技术如 Pyramid Attention Broadcast、FasterCache、FirstBlockCache 等。仓库的模型测试中也内置了FirstBlockCacheTesterMixin用于验证该缓存的正确性见 tests/models/testing_utils/cache.py。在 FLUX.1-dev 上启用 First Block Cache对 FLUX.1-dev 应用 First Block Cache 只需调用apply_cache_on_pipe。0.08是 FLUX 模型的默认残差差值import time import torch from diffusers import FluxPipeline pipe FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, dtypetorch.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_threshold0.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_steps28, ).images[0] end time.time() print(fTime: {end - begin:.2f}s) print(Saving image to flux.png) image.save(flux.png)residual_diff_thresholdrdt是控制缓存激进程度的关键参数文档给出了不同取值下的对比数据NVIDIA L20单卡| Optimizations | Original | FBCache rdt0.06 | FBCache rdt0.08 | FBCache rdt0.10 | FBCache rdt0.12 | | - | - | - | - | - | - | | Wall Time (s) | 26.36 | 21.83 | 17.01 | 16.00 | 13.78 |可见 rdt 越大跳过的去噪步越多、耗时越短但质量与速度需要权衡。使用默认的 rdt0.08 时推理耗时从基线 26.36 秒降至17.01 秒约 1.55 倍加速同时几乎不损失生成质量。在 HunyuanVideo 上启用 First Block Cache对 HunyuanVideo 同样调用apply_cache_on_pipe其默认残差差值为0.06代码示例中使用的0.6与默认值略有出入实际调参时可从 0.06 起步逐步放大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, subfoldertransformer, dtypetorch.bfloat16, revisionrefs/pr/18, ) pipe HunyuanVideoPipeline.from_pretrained( model_id, transformertransformer, dtypetorch.float16, revisionrefs/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_threshold0.6) pipe.vae.enable_tiling() begin time.time() output pipe( promptA cat walks on the grass, realistic, height720, width1280, num_frames129, num_inference_steps30, ).frames[0] end time.time() print(fTime: {end - begin:.2f}s) print(Saving video to hunyuan_video.mp4) export_to_video(output, hunyuan_video.mp4, fps15)注意此处为视频模型单独启用了pipe.vae.enable_tiling()通过 VAE 分块解码控制显存占用。启用 First Block Cache 后HunyuanVideo 推理耗时从基线 3675.71 秒降至2271.06 秒约 1.62 倍加速且几乎不损失视频质量。fp8 动态量化激活与权重双量化原理与前提fp8 动态量化能够进一步加速推理并降低显存占用。需要强调的是只有激活值activations和权重weights都完成量化才能实际利用 NVIDIA Tensor Cores 的 8 位计算单元。文档推荐使用float8_weight_only仅权重 fp8 量化与float8_dynamic_activation_float8_weight激活动态 fp8 权重 fp8两种模式分别作用于文本编码器text encoder与 transformer 模型。默认采用 per-tensor逐张量量化如果你的 GPU 支持 row-wise逐行量化也可以尝试以获得更好的精度。安装 torchao 与 torch.compile量化能力由 torchao 提供安装命令如下pip3 install -U torch torchao同时建议配合torch.compile使用modemax-autotune-no-cudagraphs或modemax-autotune让编译器自动挑选最优 kernel。首次调用模型时的编译耗时较长但编译完成后收益显著、值得等待。[!TIP] 动态量化会显著改变模型输出的分布因此启用 fp8 后需要把residual_diff_threshold调大一些缓存机制才能继续生效。FLUX.1-devFBCache fp8 torch.compile下面的示例只量化了 transformer如需进一步降低显存也可以同时量化文本编码器import time import torch from diffusers import FluxPipeline pipe FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, dtypetorch.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_threshold0.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, modemax-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_steps28, ).images[0] end time.time() if i 0: print(fWarm up time: {end - begin:.2f}s) else: print(fTime: {end - begin:.2f}s) print(Saving image to flux.png) image.save(flux.png)注意这里采用先跑一次 warm up、再计时的写法以规避 torch.compile 首次编译的时间。文档基准显示fp8 动态量化 torch.compile 组合将推理耗时降至7.56 秒约 3.48 倍加速。HunyuanVideo量化收益有限且需警惕 OOMimport 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, subfoldertransformer, dtypetorch.bfloat16, revisionrefs/pr/18, ) pipe HunyuanVideoPipeline.from_pretrained( model_id, transformertransformer, dtypetorch.float16, revisionrefs/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, modemax-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( promptA cat walks on the grass, realistic, height720, width1280, num_frames129, num_inference_steps1 if i 0 else 30, ).frames[0] end time.time() if i 0: print(fWarm up time: {end - begin:.2f}s) else: print(fTime: {end - begin:.2f}s) print(Saving video to hunyuan_video.mp4) export_to_video(output, hunyuan_video.mp4, fps15)针对视频模型需要注意两点显存限制NVIDIA L20 仅有 48GB 显存HunyuanVideo 在高分辨率、大帧数下激活张量极大编译后若未调用enable_model_cpu_offload很可能触发 OOM。显存低于 80GB 的 GPU 建议降低分辨率与帧数来规避。收益有限大型视频生成模型的计算瓶颈通常是注意力attention计算而非全连接层因此从量化和 torch.compile 中获得的收益不如图像模型明显。这也是下文 HunyuanVideo 的上下文并行组合中默认不开启 fp8 与 torch.compile 的原因。Context Parallelism多 GPU 上下文并行工作原理上下文并行将长上下文图像/视频 token 序列切分到多张 GPU 上协同推理是突破单卡显存与算力瓶颈的核心手段。ParaAttention 的组合式设计允许将上下文并行与 First Block Cache、动态量化自由叠加。在 Diffusers 仓库中与之对应的原生接口是 src/diffusers/hooks/context_parallel.py 中的apply_context_parallel它根据ContextParallelModelPlan在指定子模块上挂载ContextParallelSplitHook输入切分钩子与ContextParallelGatherHook输出聚合钩子实现切分—各卡计算—聚合的完整流程。[!TIP] 如果推理过程需要长期驻留、对外提供服务建议基于torch.multiprocessing自行编写推理进程从而省去每次启动进程、加载模型和重新编译的开销。FLUX.1-devFBCache fp8 torch.compile CP 全组合下面的示例组合了全部四项优化以获得最快的推理速度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, dtypetorch.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_size2, ) parallelize_pipe( pipe, meshmesh, ) parallelize_vae(pipe.vae, meshmesh._flatten()) from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe apply_cache_on_pipe( pipe, residual_diff_threshold0.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, modemax-autotune-no-cudagraphs, ) # Enable memory savings # pipe.enable_model_cpu_offload(gpu_iddist.get_rank()) # pipe.enable_sequential_cpu_offload(gpu_iddist.get_rank()) for i in range(2): begin time.time() image pipe( A cat holding a sign that says hello world, num_inference_steps28, output_typepil if dist.get_rank() 0 else pt, ).images[0] end time.time() if dist.get_rank() 0: if i 0: print(fWarm up time: {end - begin:.2f}s) else: print(fTime: {end - begin:.2f}s) if dist.get_rank() 0: print(Saving image to flux.png) image.save(flux.png) dist.destroy_process_group()几个值得注意的细节先用dist.init_process_group()初始化进程组再通过dist.get_rank()设置当前进程对应的 CUDA 设备init_context_parallel_mesh(pipe.device.type, max_ring_dim_size2)创建并行通信网格max_ring_dim_size2限制环形通信维度的规模parallelize_pipe并行化 transformer 主体parallelize_vae并行化 VAE 解码使用mesh._flatten()展平后的网格torch._inductor.config.reorder_for_compute_comm_overlap True让计算与通信重叠进一步压榨吞吐只有 rank 0 输出 PIL 图像并保存其余 rank 输出pt张量避免重复写盘。将该脚本保存为run_flux.py并用 torchrun 启动--nproc_per_node指定 GPU 数量# Use --nproc_per_node to specify the number of GPUs torchrun --nproc_per_node2 run_flux.py文档基准显示2 张 NVIDIA L20 下推理耗时降至8.20 秒约 3.21 倍加速4 张 L20 下进一步降至3.90 秒约 6.75 倍加速。HunyuanVideoFBCache CP 组合视频模型在量化与编译上的收益有限因此下面只组合 First Block Cache 与上下文并行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, subfoldertransformer, dtypetorch.bfloat16, revisionrefs/pr/18, ) pipe HunyuanVideoPipeline.from_pretrained( model_id, transformertransformer, dtypetorch.float16, revisionrefs/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, meshmesh, ) parallelize_vae(pipe.vae, meshmesh._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, modemax-autotune-no-cudagraphs, # ) # Enable memory savings pipe.vae.enable_tiling() # pipe.enable_model_cpu_offload(gpu_iddist.get_rank()) # pipe.enable_sequential_cpu_offload(gpu_iddist.get_rank()) for i in range(2): begin time.time() output pipe( promptA cat walks on the grass, realistic, height720, width1280, num_frames129, num_inference_steps1 if i 0 else 30, output_typepil if dist.get_rank() 0 else pt, ).frames[0] end time.time() if dist.get_rank() 0: if i 0: print(fWarm up time: {end - begin:.2f}s) else: print(fTime: {end - begin:.2f}s) if dist.get_rank() 0: print(Saving video to hunyuan_video.mp4) export_to_video(output, hunyuan_video.mp4, fps15) dist.destroy_process_group()保存为run_hunyuan_video.py后用 torchrun 启动# Use --nproc_per_node to specify the number of GPUs torchrun --nproc_per_node8 run_hunyuan_video.py文档基准显示8 张 NVIDIA L20 下HunyuanVideo 推理耗时降至649.23 秒约 5.66 倍加速。基准测试汇总FLUX.1-dev1024×102428 步| GPU Type | Number of GPUs | Optimizations | Wall Time (s) | Speedup | | - | - | - | - | - | | NVIDIA L20 | 1 | Baseline | 26.36 | 1.00x | | NVIDIA L20 | 1 | FBCache (rdt0.08) | 17.01 | 1.55x | | NVIDIA L20 | 1 | FP8 DQ | 13.40 | 1.96x | | NVIDIA L20 | 1 | FBCache (rdt0.12) FP8 DQ | 7.56 | 3.48x | | NVIDIA L20 | 2 | FBCache (rdt0.12) FP8 DQ CP | 4.92 | 5.35x | | NVIDIA L20 | 4 | FBCache (rdt0.12) FP8 DQ CP | 3.90 | 6.75x |HunyuanVideo720p129 帧30 步| GPU Type | Number of GPUs | Optimizations | Wall Time (s) | Speedup | | - | - | - | - | - | | NVIDIA L20 | 1 | Baseline | 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 |从两张表可以清晰看出组合优化的收益规律对图像模型 FLUX.1-devFBCache、fp8 动态量化与上下文并行三者可以无损叠加4 卡即可获得约 6.75 倍加速对视频模型 HunyuanVideo加速主要来自 FBCache 与上下文并行且 GPU 数量越多收益越接近线性8 卡达到约 5.66 倍加速。总结与实践建议结合本文与仓库源码可以提炼出以下可复用的实战要点优先开启 First Block Cache它改动最小、收益立竿见影且已原生集成于 Diffusers 的钩子体系中src/diffusers/hooks/first_block_cache.py。调参时记住rdt 越大跳过步数越多从各模型默认值FLUX 为 0.08、HunyuanVideo 为 0.06出发结合质量验收逐步调整。图像模型叠加 fp8 torch.compilequantize_torch.compile(modemax-autotune-no-cudagraphs)是图像模型的黄金组合但动态量化会改变输出分布需同步调大residual_diff_threshold如 0.12。首次运行务必先 warm up 再计时。视频模型靠多卡并行视频生成的瓶颈在注意力计算量化与编译收益有限面对高分辨率、大帧数需求优先用上下文并行扩容并留意 L2048GB等小显存 GPU 的 OOM 风险必要时配合pipe.vae.enable_tiling()与enable_model_cpu_offload。多卡场景用 torchrun 启动脚本内通过dist.init_process_group()parallelize_pipe/parallelize_vae完成并行化再以torchrun --nproc_per_nodeN拉起有 NVLink 的 A100/H100 集群收益更佳。长期服务优先自建常驻进程用torch.multiprocessing编写服务化推理进程可消除反复启动、加载与编译模型的固定开销。按上述路线组合运用你可以在不牺牲生成质量的前提下把 FLUX.1-dev 的单卡推理从 26 秒量级压到 7.5 秒左右把 HunyuanVideo 的 720p 视频生成从一小时量级压到十分钟量级——这正是 ParaAttention 组合式设计带来的实际价值。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 用 ParaAttention 加速 Diffusers 推理：First Block Cache、fp8 量化与上下文并行实战指南
用 ParaAttention 加速 Diffusers 推理First Block Cache、fp8 量化与上下文并行实战指南【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers大型图像与视频生成模型如 FLUX.1-dev、HunyuanVideo凭借出色的生成质量成为业界焦点但其庞大的参数量与计算开销给实时应用和线上部署带来了严峻的推理挑战。本文以 Diffusers 官方文档 docs/source/en/optimization/para_attn.md 为核心系统讲解 ParaAttention 库的三项关键技术——上下文并行Context Parallelism、First Block Cache首块缓存FBCache与fp8 动态量化并结合当前仓库源码src/diffusers/hooks/first_block_cache.py、src/diffusers/models/cache_utils.py深入剖析其底层原理。读完本文你将掌握如何在 NVIDIA L20 GPU 上对 FLUX.1-dev 与 HunyuanVideo 组合运用上述优化手段实现最高约 6.75 倍的端到端推理加速。ParaAttention 是什么ParaAttention 是一个面向扩散模型的推理加速库核心实现了两类优化上下文并行Context Parallelism将长序列图像 token、视频帧 token切分到多张 GPU 上并行计算随 GPU 数量线性扩展吞吐First Block Cache首块缓存复用 transformer 首块的残差输出直接跳过部分去噪步减少计算量。此外ParaAttention 采用组合式设计可以进一步与torch.compile、fp8 动态量化等技术叠加使用获得更大的加速比。本文的基线基准测试baseline默认不做任何优化唯一例外是 HunyuanVideo 为了避免显存溢出OOM而保留了必要的显存管理设置。在 1 张 NVIDIA L20 GPU 上基线测试结果为FLUX.1-dev 生成一张 1024×1024 图像28 步耗时26.36 秒HunyuanVideo 生成 129 帧 720p 视频30 步耗时3675.71 秒。下面的优化方案都以这一基线作为对比参照。[!TIP] 如果想要更快的上下文并行推理在条件允许时优先选用带 NVLink 互联的 NVIDIA A100 或 H100 GPU尤其是在 GPU 数量较多的场景下NVLink 能显著降低跨卡通信开销。First Block Cache用残差相似度跳过冗余去噪步工作原理在扩散模型的迭代去噪过程中相邻时间步之间模型输出的变化往往是平滑、渐进的。First Block Cache 的思想是缓存模型 transformer 各 block 的输出并在后续推理步骤中直接复用从而省去部分前向计算、加速推理。难点在于如何判断何时可以安全复用缓存而不损害生成质量。ParaAttention 给出的方案非常巧妙直接用第一个 transformer block 输出的残差差值来近似整个模型输出的差异。当两次去噪步之间首块输出的残差变化足够小时就认为模型整体输出变化可忽略于是复用之前推理步的残差跳过整个去噪步。据文档所述该方法在 FLUX.1-dev 与 HunyuanVideo 上可带来约 2 倍的速度提升同时保持很好的生成质量。仓库源码中的实现细节First Block Cache 并非黑盒魔法在 Diffusers 仓库中可以找到它完整的钩子hook实现见 src/diffusers/hooks/first_block_cache.pyFirstBlockCacheConfigL32-L48核心参数threshold默认值0.05。阈值越高越倾向于跳过更多层的前向计算、推理越快但生成质量可能下降阈值越低则加速有限。判断依据是首块输出残差与缓存残差之间的absmean绝对均值差值当差值低于阈值时跳过前向。FBCHeadBlockHookL65-L142挂在第一个transformer block 上。每次前向计算首块的真实输出残差并通过_should_compute_remaining_blocks比较当前残差与上一缓存残差若相对差异diff (absmean / prev_hidden_states_absmean)大于阈值则继续计算其余 block否则直接复用缓存的尾块残差拼出最终输出。FBCBlockHookL145-L190挂在其余所有 block含尾块上。当should_compute为 False 时各 block 直接返回输入恒等映射从而彻底跳过计算。apply_first_block_cacheL193-L246遍历模型的named_children()利用 src/diffusers/hooks/_common.py 中定义的_ALL_TRANSFORMER_BLOCK_IDENTIFIERS涵盖blocks、transformer_blocks、single_transformer_blocks、layers、temporal_transformer_blocks等常见命名识别 transformer block 列表随后对头块、中间块、尾块分别挂载对应钩子。也就是说Diffusers 本身也原生集成了 First Block Cache 机制。除了通过para_attn的apply_cache_on_pipe一键启用外还可以直接用 apply_first_block_cache 配合FirstBlockCacheConfig(threshold...)作用到任意支持的 transformer 上模型层面则通过 CacheMixin 提供的enable_cache(config)/disable_cache()统一管理缓存开关同一时刻仅允许启用一种缓存技术如 Pyramid Attention Broadcast、FasterCache、FirstBlockCache 等。仓库的模型测试中也内置了FirstBlockCacheTesterMixin用于验证该缓存的正确性见 tests/models/testing_utils/cache.py。在 FLUX.1-dev 上启用 First Block Cache对 FLUX.1-dev 应用 First Block Cache 只需调用apply_cache_on_pipe。0.08是 FLUX 模型的默认残差差值import time import torch from diffusers import FluxPipeline pipe FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, dtypetorch.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_threshold0.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_steps28, ).images[0] end time.time() print(fTime: {end - begin:.2f}s) print(Saving image to flux.png) image.save(flux.png)residual_diff_thresholdrdt是控制缓存激进程度的关键参数文档给出了不同取值下的对比数据NVIDIA L20单卡| Optimizations | Original | FBCache rdt0.06 | FBCache rdt0.08 | FBCache rdt0.10 | FBCache rdt0.12 | | - | - | - | - | - | - | | Wall Time (s) | 26.36 | 21.83 | 17.01 | 16.00 | 13.78 |可见 rdt 越大跳过的去噪步越多、耗时越短但质量与速度需要权衡。使用默认的 rdt0.08 时推理耗时从基线 26.36 秒降至17.01 秒约 1.55 倍加速同时几乎不损失生成质量。在 HunyuanVideo 上启用 First Block Cache对 HunyuanVideo 同样调用apply_cache_on_pipe其默认残差差值为0.06代码示例中使用的0.6与默认值略有出入实际调参时可从 0.06 起步逐步放大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, subfoldertransformer, dtypetorch.bfloat16, revisionrefs/pr/18, ) pipe HunyuanVideoPipeline.from_pretrained( model_id, transformertransformer, dtypetorch.float16, revisionrefs/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_threshold0.6) pipe.vae.enable_tiling() begin time.time() output pipe( promptA cat walks on the grass, realistic, height720, width1280, num_frames129, num_inference_steps30, ).frames[0] end time.time() print(fTime: {end - begin:.2f}s) print(Saving video to hunyuan_video.mp4) export_to_video(output, hunyuan_video.mp4, fps15)注意此处为视频模型单独启用了pipe.vae.enable_tiling()通过 VAE 分块解码控制显存占用。启用 First Block Cache 后HunyuanVideo 推理耗时从基线 3675.71 秒降至2271.06 秒约 1.62 倍加速且几乎不损失视频质量。fp8 动态量化激活与权重双量化原理与前提fp8 动态量化能够进一步加速推理并降低显存占用。需要强调的是只有激活值activations和权重weights都完成量化才能实际利用 NVIDIA Tensor Cores 的 8 位计算单元。文档推荐使用float8_weight_only仅权重 fp8 量化与float8_dynamic_activation_float8_weight激活动态 fp8 权重 fp8两种模式分别作用于文本编码器text encoder与 transformer 模型。默认采用 per-tensor逐张量量化如果你的 GPU 支持 row-wise逐行量化也可以尝试以获得更好的精度。安装 torchao 与 torch.compile量化能力由 torchao 提供安装命令如下pip3 install -U torch torchao同时建议配合torch.compile使用modemax-autotune-no-cudagraphs或modemax-autotune让编译器自动挑选最优 kernel。首次调用模型时的编译耗时较长但编译完成后收益显著、值得等待。[!TIP] 动态量化会显著改变模型输出的分布因此启用 fp8 后需要把residual_diff_threshold调大一些缓存机制才能继续生效。FLUX.1-devFBCache fp8 torch.compile下面的示例只量化了 transformer如需进一步降低显存也可以同时量化文本编码器import time import torch from diffusers import FluxPipeline pipe FluxPipeline.from_pretrained( black-forest-labs/FLUX.1-dev, dtypetorch.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_threshold0.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, modemax-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_steps28, ).images[0] end time.time() if i 0: print(fWarm up time: {end - begin:.2f}s) else: print(fTime: {end - begin:.2f}s) print(Saving image to flux.png) image.save(flux.png)注意这里采用先跑一次 warm up、再计时的写法以规避 torch.compile 首次编译的时间。文档基准显示fp8 动态量化 torch.compile 组合将推理耗时降至7.56 秒约 3.48 倍加速。HunyuanVideo量化收益有限且需警惕 OOMimport 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, subfoldertransformer, dtypetorch.bfloat16, revisionrefs/pr/18, ) pipe HunyuanVideoPipeline.from_pretrained( model_id, transformertransformer, dtypetorch.float16, revisionrefs/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, modemax-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( promptA cat walks on the grass, realistic, height720, width1280, num_frames129, num_inference_steps1 if i 0 else 30, ).frames[0] end time.time() if i 0: print(fWarm up time: {end - begin:.2f}s) else: print(fTime: {end - begin:.2f}s) print(Saving video to hunyuan_video.mp4) export_to_video(output, hunyuan_video.mp4, fps15)针对视频模型需要注意两点显存限制NVIDIA L20 仅有 48GB 显存HunyuanVideo 在高分辨率、大帧数下激活张量极大编译后若未调用enable_model_cpu_offload很可能触发 OOM。显存低于 80GB 的 GPU 建议降低分辨率与帧数来规避。收益有限大型视频生成模型的计算瓶颈通常是注意力attention计算而非全连接层因此从量化和 torch.compile 中获得的收益不如图像模型明显。这也是下文 HunyuanVideo 的上下文并行组合中默认不开启 fp8 与 torch.compile 的原因。Context Parallelism多 GPU 上下文并行工作原理上下文并行将长上下文图像/视频 token 序列切分到多张 GPU 上协同推理是突破单卡显存与算力瓶颈的核心手段。ParaAttention 的组合式设计允许将上下文并行与 First Block Cache、动态量化自由叠加。在 Diffusers 仓库中与之对应的原生接口是 src/diffusers/hooks/context_parallel.py 中的apply_context_parallel它根据ContextParallelModelPlan在指定子模块上挂载ContextParallelSplitHook输入切分钩子与ContextParallelGatherHook输出聚合钩子实现切分—各卡计算—聚合的完整流程。[!TIP] 如果推理过程需要长期驻留、对外提供服务建议基于torch.multiprocessing自行编写推理进程从而省去每次启动进程、加载模型和重新编译的开销。FLUX.1-devFBCache fp8 torch.compile CP 全组合下面的示例组合了全部四项优化以获得最快的推理速度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, dtypetorch.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_size2, ) parallelize_pipe( pipe, meshmesh, ) parallelize_vae(pipe.vae, meshmesh._flatten()) from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe apply_cache_on_pipe( pipe, residual_diff_threshold0.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, modemax-autotune-no-cudagraphs, ) # Enable memory savings # pipe.enable_model_cpu_offload(gpu_iddist.get_rank()) # pipe.enable_sequential_cpu_offload(gpu_iddist.get_rank()) for i in range(2): begin time.time() image pipe( A cat holding a sign that says hello world, num_inference_steps28, output_typepil if dist.get_rank() 0 else pt, ).images[0] end time.time() if dist.get_rank() 0: if i 0: print(fWarm up time: {end - begin:.2f}s) else: print(fTime: {end - begin:.2f}s) if dist.get_rank() 0: print(Saving image to flux.png) image.save(flux.png) dist.destroy_process_group()几个值得注意的细节先用dist.init_process_group()初始化进程组再通过dist.get_rank()设置当前进程对应的 CUDA 设备init_context_parallel_mesh(pipe.device.type, max_ring_dim_size2)创建并行通信网格max_ring_dim_size2限制环形通信维度的规模parallelize_pipe并行化 transformer 主体parallelize_vae并行化 VAE 解码使用mesh._flatten()展平后的网格torch._inductor.config.reorder_for_compute_comm_overlap True让计算与通信重叠进一步压榨吞吐只有 rank 0 输出 PIL 图像并保存其余 rank 输出pt张量避免重复写盘。将该脚本保存为run_flux.py并用 torchrun 启动--nproc_per_node指定 GPU 数量# Use --nproc_per_node to specify the number of GPUs torchrun --nproc_per_node2 run_flux.py文档基准显示2 张 NVIDIA L20 下推理耗时降至8.20 秒约 3.21 倍加速4 张 L20 下进一步降至3.90 秒约 6.75 倍加速。HunyuanVideoFBCache CP 组合视频模型在量化与编译上的收益有限因此下面只组合 First Block Cache 与上下文并行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, subfoldertransformer, dtypetorch.bfloat16, revisionrefs/pr/18, ) pipe HunyuanVideoPipeline.from_pretrained( model_id, transformertransformer, dtypetorch.float16, revisionrefs/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, meshmesh, ) parallelize_vae(pipe.vae, meshmesh._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, modemax-autotune-no-cudagraphs, # ) # Enable memory savings pipe.vae.enable_tiling() # pipe.enable_model_cpu_offload(gpu_iddist.get_rank()) # pipe.enable_sequential_cpu_offload(gpu_iddist.get_rank()) for i in range(2): begin time.time() output pipe( promptA cat walks on the grass, realistic, height720, width1280, num_frames129, num_inference_steps1 if i 0 else 30, output_typepil if dist.get_rank() 0 else pt, ).frames[0] end time.time() if dist.get_rank() 0: if i 0: print(fWarm up time: {end - begin:.2f}s) else: print(fTime: {end - begin:.2f}s) if dist.get_rank() 0: print(Saving video to hunyuan_video.mp4) export_to_video(output, hunyuan_video.mp4, fps15) dist.destroy_process_group()保存为run_hunyuan_video.py后用 torchrun 启动# Use --nproc_per_node to specify the number of GPUs torchrun --nproc_per_node8 run_hunyuan_video.py文档基准显示8 张 NVIDIA L20 下HunyuanVideo 推理耗时降至649.23 秒约 5.66 倍加速。基准测试汇总FLUX.1-dev1024×102428 步| GPU Type | Number of GPUs | Optimizations | Wall Time (s) | Speedup | | - | - | - | - | - | | NVIDIA L20 | 1 | Baseline | 26.36 | 1.00x | | NVIDIA L20 | 1 | FBCache (rdt0.08) | 17.01 | 1.55x | | NVIDIA L20 | 1 | FP8 DQ | 13.40 | 1.96x | | NVIDIA L20 | 1 | FBCache (rdt0.12) FP8 DQ | 7.56 | 3.48x | | NVIDIA L20 | 2 | FBCache (rdt0.12) FP8 DQ CP | 4.92 | 5.35x | | NVIDIA L20 | 4 | FBCache (rdt0.12) FP8 DQ CP | 3.90 | 6.75x |HunyuanVideo720p129 帧30 步| GPU Type | Number of GPUs | Optimizations | Wall Time (s) | Speedup | | - | - | - | - | - | | NVIDIA L20 | 1 | Baseline | 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 |从两张表可以清晰看出组合优化的收益规律对图像模型 FLUX.1-devFBCache、fp8 动态量化与上下文并行三者可以无损叠加4 卡即可获得约 6.75 倍加速对视频模型 HunyuanVideo加速主要来自 FBCache 与上下文并行且 GPU 数量越多收益越接近线性8 卡达到约 5.66 倍加速。总结与实践建议结合本文与仓库源码可以提炼出以下可复用的实战要点优先开启 First Block Cache它改动最小、收益立竿见影且已原生集成于 Diffusers 的钩子体系中src/diffusers/hooks/first_block_cache.py。调参时记住rdt 越大跳过步数越多从各模型默认值FLUX 为 0.08、HunyuanVideo 为 0.06出发结合质量验收逐步调整。图像模型叠加 fp8 torch.compilequantize_torch.compile(modemax-autotune-no-cudagraphs)是图像模型的黄金组合但动态量化会改变输出分布需同步调大residual_diff_threshold如 0.12。首次运行务必先 warm up 再计时。视频模型靠多卡并行视频生成的瓶颈在注意力计算量化与编译收益有限面对高分辨率、大帧数需求优先用上下文并行扩容并留意 L2048GB等小显存 GPU 的 OOM 风险必要时配合pipe.vae.enable_tiling()与enable_model_cpu_offload。多卡场景用 torchrun 启动脚本内通过dist.init_process_group()parallelize_pipe/parallelize_vae完成并行化再以torchrun --nproc_per_nodeN拉起有 NVLink 的 A100/H100 集群收益更佳。长期服务优先自建常驻进程用torch.multiprocessing编写服务化推理进程可消除反复启动、加载与编译模型的固定开销。按上述路线组合运用你可以在不牺牲生成质量的前提下把 FLUX.1-dev 的单卡推理从 26 秒量级压到 7.5 秒左右把 HunyuanVideo 的 720p 视频生成从一小时量级压到十分钟量级——这正是 ParaAttention 组合式设计带来的实际价值。【免费下载链接】diffusers Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch.项目地址: https://gitcode.com/GitHub_Trending/di/diffusers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考