ComfyUI v0.45.0 将 Qwen3.x 本地文本生成提速至多两倍
Doubling LLM token rates in ComfyUI
ComfyUI 重写文本生成热路径,同一 GPU 上 Qwen3.x 的 token 生成速度最高提升至两倍,模型与输出质量不变。
ComfyUI 官方拆解了本地 LLM 推理的显存带宽瓶颈,并给出 MTP 树状草稿与内核重排的具体做法,可供同类推理优化参考。
ComfyUI 在本地运行多种大语言模型(LLM),用于增强提示词、撰写故事、描述图像,甚至进行基础聊天。最近,我们重写了该文本生成的热路径,重点关注 Qwen3.x 系列 LLM。结果是,在相同 GPU 上使用相同检查点,现在生成 token 的速度最高可提升至两倍,而模型本身及其生成内容的质量均无变化。这也加速了近期最先进图像生成模型所包含的许多提示词增强器。
在 RTX 5090 上,Qwen3.5 9B 从每秒 106 个 token 提升到 219 个 token;在 RTX 5080 上,从每秒 91 个提升到 129 个。Qwen3.8 27B 在 RTX 5090 上描述一张 1536×1536 的图像时,从 55 提升到 122。
Qwen3.8 27B W4A8 在 RTX 5090 上描述一张 1500×1500 的图像并撰写故事:每秒 137 个 token,mtp 于 auto。
本博客讨论 LLM 推理的一些架构与实现细节,以及近期的速度提升。本文涵盖的所有内容均已在 ComfyUI v0.45.0 中提供。只需加载 Qwen3.x 检查点,保持 Text Generate 节点的 mtp 选项为 auto,即可开始使用。
瓶颈在哪里?
文本生成是顺序进行的:每个新 token 都依赖于前一个 token,并且每个 token 都需要完整读取一次模型权重。大规模 LLM 推理通常可以通过将多个推理批处理在一起来节省这种权重读取成本,例如来自多个用户的多个完全独立的聊天。然而,单个 ComfyUI 工作流通常只生成单个序列,因此每一步的批大小都为一。对于量化到四位的 27B 参数模型,例如 Qwen3.8 W4A8,一步大约读取 14 GB 权重来生成一个 token,并且每个权重可能只参与一次数学运算。token 生成的瓶颈取决于你的 GPU 从 VRAM 读取模型的速度。
相比之下,ComfyUI 中用于图像或视频生成的扩散模型在每一步都会处理整个相应的潜在输出,每个权重都作用于完整序列长度。与 LLM 不同,扩散模型受算力瓶颈限制,即取决于你的 GPU 计算速度。这为图像和视频生成领域带来了非常不同的优化问题。
LLM 模型中的大部分权重是线性矩阵乘法。这些由线性内核处理,称为 GEMM/GEMV 内核,是专为矩阵乘法优化的 GPU 内核(程序)。在我们的 LLM 场景中,一个稠密且通常是方形的矩阵会乘以一个向量 M=1(用于单个 token 的状态),或乘以一个非常薄的矩阵,如 M=3(用于小 token 数量的多 token 预测(MTP)验证状态)。这就是“瘦”线性问题。
在近期的这些改动中,针对 VRAM 总线利用率和效率加入了重大优化,纯粹围绕 VRAM 读取速度和 VRAM 内存总线占用时间进行重新优化。
什么是多 token 预测?
Qwen3.x 检查点自带一个多 token 预测(MTP)头:相比主模型,它是一个更小的额外层,利用模型的隐藏状态来预测下一个 token。ComfyUI 利用这一 MTP 特性进行投机解码。MTP 头先草拟出几个 token,然后在下一个 token 的前向传播中,所有候选草稿 token 会被一起投机验证。这一过程的成本大约相当于一个 token,因为该次传播仍然受限于 VRAM 速度。每一个与主模型本会选择的 token 相匹配的 MTP 草稿都会被模型接受,从而省去一次完整的前向传播,进而为每个正确猜中的 MTP token 省去一次完整的模型读取。这就是 MTP 如何能将输出 token 速率提升一倍以上,从而显著提高性能。
在近期的这些改动中,我们既优化了 MTP 的计算内核,又通过新的自适应树 MTP 算法增强了 MTP 算法本身。
优化线性内核
ComfyUI 中的 Qwen3.x 检查点通常是 W4A8(4 位权重、8 位激活)或 W6A8(6 位权重、8 位激活)。在 ComfyUI v0.40.0 中,这些由一条通用 INT8 路径处理:量化激活,运行一个矩阵-向量线性内核,该内核的编写仅匹配 safetensors 文件中存储的 W4A8 VRAM 布局。我们将其替换为一个优化的矩阵-向量(GEMV)内核,其中权重及其分组缩放因子在内存中按内核消费它们的顺序交错排列。这样,每个线程块流式读取一段连续数据,而不必在权重张量和缩放张量之间来回跳转,从而减少 VRAM 寻址并提高批量 VRAM 读取流水线的效率。这种内核专用的 VRAM 重布局在 VRAM 中就地完成,由动态 VRAM 权重加载引擎管理,每次模型加载只需执行一次。此过程不需要额外的 RAM 或 VRAM。这种重布局仅对我们的瘦线性层需要,因为较胖的扩散模型线性层受限于计算,会掩盖上述 VRAM 寻址的延迟。
大多数现代 GPU 都具备张量核心,可用单条 GPU 指令计算一个矩阵乘法运算块。从 30 系列开始的 NVIDIA GPU 针对最小瓦片维度 8 进行了优化,其他矩形维度可以更大。新内核使用这种矩阵乘累加(MMA)指令,权重和输入的排列方式使得上例图中所示的瘦维度 M, 映射到硬件支持的最小瓦片维度 8。这样就能使用单条指令,同时最大限度地减少浪费的未使用计算行(当 M < 8 时)。对于 M=1 的纯单 token 生成场景,其性能与使用非张量(CUDA)核心指令(dp4a)大致相同。然而,考虑到 MTP 中 M > 1 的情况,瓦片化方法会立即变得更快。最终,最多八行的 MTP 猜测几乎变得免费,因为我们的瓦片化器能以与 M=1 相同的速度计算最多 M=8 。
线性内核中新增了序言融合:激活量化现在会吸收其前面的 RMSNorm 或 SwiGLU,因此这些中间张量不会在中间被写入内存,从而最大限度地减少缓存内存流量和内核启动开销。这反过来又最大限度地减少了 VRAM 总线活动之间的小段纯计算间隙,带来小幅额外加速。
关注一下注意力!
Qwen3.5 和 3.8 是混合注意力模型:这类模型使用两种不同的注意力方案,分别应用于不同的 transformer 块。每四层中有三层是带有小型循环状态的 DeltaNet 线性注意力,第四层是带有 KV 缓存的普通分组查询注意力。
对于注意力层,我们添加了一个解码内核,它直接处理单个 token 的短查询,或来自 MTP 的少量草稿 token 及其因果验证行,而无需物化注意力掩码。它将缓存前缀拆分到多个线程块上,因此只有极小的查询会占用 GPU,最后完成的块在同一次启动中合并部分结果。
对于 DeltaNet 层,代价高昂的部分是推测。验证多个 token 的草稿需要每个 token 之后的循环状态,如果草稿被拒绝,则必须恢复该状态。此前,这意味着在每个推测 token 之后为每一层对每个头快照完整的 64 KiB 状态。现在,内核记录小的逐 token 中间值,并推迟状态更新,以便下一次调用只重放已接受的 token,并一次性写入已提交的状态。该内核反过来用少量额外的算术换取大量额外的内存流量,这契合我们的核心瓶颈问题:VRAM 受限但算力有余。
此前,KV 缓存自回归注意力内核与新增的优化 Comfy-Kitchen 注意力方案是分开的。Comfy-Kitchen 注意力以及 KV 缓存支持已在 ComfyUI v0.45 中添加,为博客前面提到的 LLM 特定注意力方案启用 INT8 注意力。这减少了 KV 缓存的 VRAM 占用,同时显著提高了长提示、系统提示或已有上下文的预填充速度。
你可以使用 --use-ck-attention 启动 ComfyUI 以启用 Comfy-Kitchen 注意力。
自适应树 MTP
在 ComfyUI v0.40.0 中,MTP 草稿最初仅限于单条 token 猜测链,对应下一个 2 到 5 个 token 的单一猜测。该链会起草一个 token,将其反馈给模型,然后起草下一个 token,直到为 MTP 设定的深度值。默认情况下,将 MTP 设置为 auto 会根据采样模式或临时成功度量自适应地给出 2 到 5 个 token。这种链式方法的弱点在于,它把所有赌注都押在这单一猜测上,链上更靠后的每个猜测如果出错都可能丢失。如果 MTP 头在深度 1 的首选是错误的,整个步骤只产出一个 token,即使深度 1 的次优猜测是正确的。
链式 MTP 只起草一条路径。树式 MTP 在每一层保留次优 token 作为兄弟节点,一次验证整棵树,并根据命中情况重新规划树的形状。
避免这条链上潜在折点的另一种方法是使用自适应树算法。在每一层,MTP 头的首选会像之前一样,以每次草稿的最佳猜测延续主链,但会保留少量次优猜测作为后续 MTP 验证的兄弟候选。这形成了一个小型草稿树结构,这些额外的候选叶子会被验证,但不会从它们继续草拟。理论上,在任何深度的树结构中,你都可以草拟任意数量的猜测,但每增加一个猜测的收益递减;不过我们的线性内核可以“免费”计算最多 M=8 个。这将猜测数量限制为七个(M=8-1),以匹配我们的计算内核和硬件。验证阶段现在最多携带八行;已提交的 token 加上七个草稿 token,同时必须将每一个视为若干条可能的分支未来。
提示词和输出的文本类型决定了树应该采用什么形状。可预测的输出,例如代码、列表或重复提示词的描述,会奖励深度,这意味着长链会作为整体被接受。创意写作则奖励广度:首选往往错误,但正确的 token 通常在前三名中。自适应 MTP 会持续测量,因此当 MTP 设为自动模式时,它会持续测量并随时重新规划。它按草稿层级保留被接受 token 的排名(是第一选择、第二、第三……还是都没有)的指数衰减计数,并用 CUDA 事件对草稿和验证阶段计时。每 32 步,它会对将七个草稿行布局为宽度链的 64 种方式进行评分,根据条件命中概率估算每步预期接受的 token 数,并除以测得的成本 verify_ms + depth × draft_ms。只有当胜出方案至少好 3% 时它才会切换,以确保形状不会来回震荡;并且由于每种形状都使用相同的八行,为该步捕获的 CUDA 图在切换后仍可继续工作。
当 MTP 值在 Generate Text 节点中设为 auto 时,可以使用自适应树 MTP。
总结
在 ComfyUI v0.40.0 和 v0.45.0 之间,ComfyUI 中的 Qwen3.x 文本生成在相同硬件上显著加快。过去会耗尽内存的较长上下文现在可以顺利容纳,推测解码器现在会根据它正在编写的文本调整自身形状。这些改进可在 NVIDIA 和 AMD、Windows 和 Linux 上运行,除了保持 ComfyUI 为最新版本外,你无需额外步骤。
来源:Comfy Blog · blog.comfy.org