文生图扩散模型本地专用推理路径设计

一、问题的提出：当模型能跑起来之后，通用栈还剩多少可省？

在主流通用方案中，本地运行文生图扩散模型的典型路线是：深度学习训练框架提供张量与自动求导基础设施，扩散模型库封装采样调度与模型组件，可视化工作流前端负责节点编排与调度，底层再由厂商提供的异构计算运行时承接算子下发。这套组合的目标是通用性——让不同模型结构、不同分辨率、不同硬件后端、不同采样器都能在同一套抽象下工作。

为了保留这种通用性，框架必须付出结构性代价：

• 算子粒度被迫细化。每个算子独立实现、独立注册、独立调度，一个算子计算完成后将结果写回显存，下一个算子再重新读取。对形状多变、组合任意的通用场景这是必要的，但对固定负载则是纯粹的搬运开销。
• 调度权集中在主机侧。CPU 按顺序逐个发起 GPU 任务，GPU 逐个执行，任务之间存在启动间隙与同步点，难以形成连续的执行流水。
• 中间结果难以驻留。算子边界即数据边界，融合需要跨算子改写计算图，通用框架出于正确性与兼容性考虑往往不做激进融合。
• 内存管理偏向保守。为适配未知的模型规模与分辨率组合，显存分配策略通常留有冗余，工作区峰值难以压到理论下限。

而当我们把推理条件逐条固定下来，上述代价中的大部分就不再是”必要的保险”，而变成了”可消除的浪费”：

1. 模型结构与权重在部署期内不变；
2. 采样流程（步数、调度器、引导方式）不变；
3. 目标分辨率只有有限几种，甚至只有一种；
4. 每一步矩阵运算的张量形状可以提前静态推导；
5. 主干去噪网络与解码器的完整链路可以整体驻留在单卡显存内。

当这五个条件同时成立时，一条从计算图与硬件指令集出发重新设计的专用推理路径就成为可行选项。需要强调，这不是把某套已有实现逐行翻译到另一种后端接口，而是重新回答”这条计算链在特定硬件上应当如何被组织”。

二、为什么这类模型适合做专用化实验

选择实验对象时，需要考察模型是否具备”负载稳定 + 结构规整 + 链路可整体驻留”三个特征。以典型的蒸馏式少步文生图模型为例，其特性恰好匹配：

其一，采样步数极少。 蒸馏或对抗蒸馏类模型把传统数十步的采样压缩到个位数步，单次推理的主干网络前向次数随之下降一个量级。这带来的直接后果是：固定开销（权重加载、图构建、内核启动准备）在总时延中的占比显著上升。通用框架擅长摊薄大计算量，但对”次数少、每次都很重”的负载，调度与访存开销更容易暴露。专用化在这里的边际收益最大。

其二，主干为规整的 Transformer 结构。 扩散 Transformer 类主干由重复堆叠的注意力块与前馈块构成，每个块内的矩阵乘、归一化、激活、残差在形状上完全一致。形状规整意味着可以按块做算子融合，把”归一化—投影—激活—再投影”这类链式小算子合并成少量大内核，减少中间张量的显存往返。

其三，文本编码器、主干、解码器三段式链路清晰。 三段之间的数据依赖是单向的，接口张量（文本嵌入、潜变量、像素输出）尺寸固定。这使得整条链路可以被静态编排为一张有向执行图，各段的显存需求可以预先精确核算，从而判断能否整体驻留。

其四，量化后权重体积可控。 采用 6 比特级别的分组量化后，主干权重体积明显下降，为在 24GB 级显存内同时容纳权重、激活工作区与解码器留出空间。这也是”整链路驻留”能够成立的前提之一，否则每步都需在主机与设备间搬运权重，专用化的收益会被抵消。

三、专用路径具体优化了什么

3.1 执行图静态化与任务预编排

通用框架在运行时根据输入动态构建执行序列，每次前向都要走一遍图遍历与算子分派。专用路径在部署阶段就把整条链路固化为静态执行计划：每个内核的入参地址、形状、栅格配置、依赖顺序全部预先确定，运行时只做一次参数表的提交。

// 示意：静态执行计划中的单个节点描述
struct KernelNode {
    KernelId   id;          // 预编译内核标识
    void*      arg_table;   // 预填好的参数表地址
    Dim3       grid;        // 预先计算的栅格配置
    uint32_t   dep_mask;    // 前驱节点位掩码
};
// 运行时仅按下标顺序提交，无形状推导、无算子查找
for (const auto& node : plan.nodes())
    launch(node);

这一改动消除了主机侧的图遍历与分派开销，更重要的是让”下一个内核的入参地址”在提交当前内核时就已经已知，为后续的流水线重叠打下基础。

3.2 算子融合：以块为单位而非以算子为单位

专用路径的融合策略不是逐个小算子尝试合并，而是以 Transformer 块为融合边界做整体设计。典型融合点包括：

• 归一化与投影融合：归一化的统计量计算与随后的线性投影在同一内核内完成，统计量留在寄存器或共享内存，不写回显存；
• 注意力中的缩放与掩码融合：缩放系数作为编译期常量嵌入，避免额外的逐元素内核；
• 前馈网络的双投影融合：上行投影、激活、下行投影合并，激活结果驻留片上；
• 残差与输出投影融合：残差加法与后续投影共享同一份加载数据。

融合的收益不只是减少内核启动次数，更关键的是减少中间张量的显存写入与再读取。在访存带宽成为瓶颈的环节，这一项的收益往往大于计算本身的优化。

3.3 片上存储的显式管理

专用路径可以显式控制数据在寄存器、共享内存、显存三级之间的流动。对固定形状的矩阵乘，可以针对硬件的共享内存容量与波前（wavefront）调度特性，手工设计分块尺寸与数据搬运节奏，使每个计算单元在等待全局访存时始终有可执行的指令。

这里的关键取舍是：分块参数一旦确定就固定下来，不再运行时自适应。通用库为了实现形状无关性，通常保留多套分块模板并在运行时选择；专用路径则把选择结果固化，去掉选择逻辑本身的开销，也去掉为兼容其他形状而保留的冗余分支。

3.4 内核启动与同步的最小化

专用路径在时间轴上把相邻内核组织为连续提交，仅在存在真实数据依赖的位置插入同步。对于不同计算单元之间无依赖的子图（例如某些并行的投影分支），允许其重叠执行。这要求对计算图的依赖关系做精确分析，代价是失去了通用调度器的”自动正确性保证”，收益是去掉了大量保守同步点。

3.5 显存工作区的精确核算

由于所有张量形状静态已知，工作区大小可以在部署阶段精确计算，而非运行时按峰值动态增长。这使得显存分配一次完成、地址固定，既避免了分配器碎片，也让”整链路驻留”的可行性判断有了量化依据。

四、最终得到的结果形态与解读方式

在固定模型、固定分辨率、单卡 24GB 级显存的条件下，一条高度专用化的执行路径可以做到：单张 1024×1024 图像生成的端到端时延进入数十秒量级，显存工作区占用控制在十几 GiB 的范围，为权重与解码器留出余量。

解读这类结果时需要注意三点：

第一，这是特定条件下的端到端数字。 它对应固定的模型、固定的量化方案、固定的分辨率与固定的采样配置。更换其中任一项，执行计划都需要重新生成，性能特征也会变化。专用化的收益与通用性的损失是同一枚硬币的两面。

第二，收益来源应当拆开看。 时延改善大致来自三部分：内核启动与调度开销的消除、中间张量访存的减少、以及计算与搬运的重叠。三者贡献比例随负载特征而变，在访存受限的环节，第二项的权重通常更高。工程上应分别度量，而非只看总数。

第三，显存占用的下降来自”可预测”而非”更节省的算法”。 静态核算让工作区没有为未知形状预留的冗余，这是占用降低的主要来源；算法层面的显存优化（如重计算）在专用路径中反而是可选项，因为形状固定后通常不需要用计算换显存。

五、方法论提炼：什么情况下值得走专用路径

把这次实践抽象出来，可以得到一套判断标准：

维度	适合专用化	应保留通用栈
模型与权重	部署期内固定	频繁更换或多模型共存
采样配置	固定步数与调度	用户可调、组合多样
分辨率	有限枚举	任意尺寸、动态变化
链路驻留	可整体放入单卡	需跨卡切分或卸载
团队能力	可承担内核级维护	希望跟随框架升级

专用路径的本质是把通用性预算花在确定的地方。它不否认通用框架的价值——恰恰相反，通用框架让”模型能跑起来”这件事变得廉价，才使得”跑起来之后再优化”成为有意义的下一步。对于把某个固定模型作为生产环节中稳定组件使用的团队，这条路径提供的是可预测的时延与可精确规划的显存预算；对于需要快速试验多种模型与配置的团队，通用栈仍是更合理的选择。

工程实践上较为稳妥的做法是分层推进：先用通用栈完成正确性验证与质量基线，再把性能热点逐层剥离出来专用化——先做执行图静态化，再做块级融合，最后做片上存储与调度的手工调优。每一层都应保留与通用实现的数值对照，确保优化没有改变输出分布。这样既能拿到专用化的收益，也不至于在一次性推倒重来时失去可回归的基线。