Newshot ComfyUI 会在本地运行各种大语言模型（LLM），用于增强提示词、写故事、描述图像，甚至基础聊天。最近，我们重写了该文本生成的热路径，重点针对 Qwen3.x 系列 LLM。结果是，在相同 GPU 上，相同 checkpoint 现在生成 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 图像并写一个故事：mtp 设为 auto 时，每秒 137 个 token。
这篇博客讨论 LLM 推理的一些架构与实现细节，以及近期的加速。这里涵盖的所有内容都已在 ComfyUI v0.45.0 中提供。只需加载一个 Qwen3.x checkpoint，让 Text Generate 节点的 mtp 选项保持 auto，就可以开始了。
瓶颈是什么？
生成文本是顺序进行的：每个新 token 都依赖前一个 token，并且每个 token 都需要完整读取一次模型权重。大规模 LLM 推理通常可以通过将多个推理批处理在一起来节省这种权重读取成本，例如多个用户发起的多个完全独立的聊天。然而，单个 ComfyUI 工作流通常只生成一个序列，因此每一步的 batch 都是 1。对于一个量化到四位的 27B 参数模型，例如 Qwen3.8 W4A8，一步大约要读取 14 GB 权重来生成一个 token，并且每个权重可能只参与一次数学运算。token 生成的瓶颈取决于你的 GPU 从 VRAM 读取模型的速度。
相比之下，ComfyUI 中用于图像或视频生成的扩散模型在每一步都会处理完整的相应 latent 输出，每个权重都会作用于整个序列长度。与 LLM 不同，扩散模型受限于算力，也就是你的 GPU 计算的速度。这给图像和视频生成领域带来了非常不同的优化问题。
LLM 模型中的大部分权重都是线性矩阵乘法。这些由线性 kernel 处理，也就是 GEMM/GEMV kernel，它们是针对矩阵乘法优化的 GPU kernel（程序）。在我们的 LLM 场景中，一个稠密的、通常是方形的矩阵，会乘以一个向量 M=1（用于单个 token 的状态），或乘以一个非常薄的矩阵，例如 M=3（用于小 token 数的 Multi-Token Prediction（MTP）验证状态）。这就是“瘦长”的线性问题。
在这些近期改动中，我们针对 VRAM 总线利用率和效率加入了显著优化，纯粹围绕 VRAM 读取速度和 VRAM 内存总线占空比进行重新优化。
什么是 Multi-token Prediction ？
Qwen3.x checkpoint 自带一个多 token 预测（MTP）头：与主模型相比，这是一个更小的额外层，利用模型的隐藏状态来猜测下一个 token。ComfyUI 使用这个 MTP 特性进行投机解码。MTP 头会起草几个 token，然后在下一个 token 的前向传递中，所有候选草稿 token 会一起被投机验证。这一过程的成本约等于一个 token，因为该次传递仍受 VRAM 速度瓶颈限制。每个与主模型本来会选择的结果匹配的 MTP 草稿都会被模型接受，从而节省一次完整前向传递，进而为每个正确猜中的 MTP token 节省一次完整模型读取。这就是 MTP 如何能将输出 token 速率提高一倍以上，从而显著提升性能。
在最近的这些改动中，我们既优化了 MTP 的计算 kernel，也通过新的自适应树 MTP 算法增强了 MTP 算法本身。
优化线性 Kernel
ComfyUI 中的 Qwen3.x checkpoint 通常是 W4A8（4-bit 权重，8-bit 激活）或 W6A8（6-bit 权重，8-bit 激活）。在 ComfyUI v0.40.0 中，这些由通用 INT8 路径处理：量化激活，运行一个矩阵-向量线性 kernel，该 kernel 编写为只匹配存储在 safetensors 文件中的 W4A8 VRAM 布局。我们将其替换为优化的矩阵-向量（GEMV）kernel，其中权重及其每组缩放因子在内存中按 kernel 消费它们的顺序交错排列。这样，每个线程块流式读取连续的一段数据，而不必在权重张量和 scale 张量之间跳转，从而减少 VRAM 寻址并提高批量 VRAM 读取流水线的效率。这种特定于 kernel 的 VRAM 重排是在 VRAM 中原地完成的，由动态 VRAM 权重加载引擎管理，每次模型加载只需执行一次。此过程不需要额外 RAM 或 VRAM。只有我们的窄线性层需要这种重排，因为更宽的 diffusion 模型线性层受计算瓶颈限制，并掩盖了上述 VRAM 寻址的延迟。
大多数现代 GPU 都配备 Tensor Core，可用单条 GPU 指令计算一块矩阵乘法运算。从 30 系列开始的 NVIDIA GPU 针对最小 tile 维度 8 进行优化，其他矩形维度可以更大。新 kernel 使用这种 matrix-multiply-accumulate（MMA）指令，并安排权重和输入，使得如上图示例所示的窄维度 M 映射到硬件支持的最小 tile 维度 8。这样可以使用单条指令，同时尽量减少浪费的未使用计算行（当 M < 8 时）。对于 M=1 的纯单 token 生成场景，其性能与使用非 Tensor（CUDA）Core 指令（dp4a）大致相同。然而，在考虑 MTP 且 M > 1 时，分块方法会立即变得更快。最终，最多八行的 MTP 猜测几乎变得免费，因为我们的分块器可以以与 M=1 相同的速度计算最多 M=8。
线性 kernel 新增了 prologue 融合：激活量化现在会吸收其前面的 RMSNorm 或 SwiGLU，因此这些中间张量不会在中间写入内存，从而最小化缓存内存流量和 kernel 启动开销。这反过来最小化了 VRAM 总线活动之间的小段纯计算间隙，带来少量额外加速。
关注 attention！
Qwen3.5 和 3.8 是混合 attention 模型：这类模型使用两种不同的 attention 方案，分别应用于不同的 transformer 块。每四层中有三层是带小型递归状态的 DeltaNet 线性 attention，第四层是带 KV cache 的普通 grouped-query attention。
文章标题：在 ComfyUI 中将 LLM token 速率翻倍
正文： 针对注意力层，我们加入了一个 decode kernel，它直接处理单个 token 的短查询，或来自 MTP 的少量 draft token 及其因果验证行，而无需物化注意力掩码。它把缓存前缀拆分到许多线程块上，因此只有很小的查询会填满 GPU，最后一个完成的块会在同一次 launch 内合并部分结果。
针对 DeltaNet 层，代价高昂的部分是投机。验证由多个 token 组成的 draft 需要每个 token 之后的循环状态，而且如果某个 draft 被拒绝，状态必须被恢复。此前，这意味着在每个投机 token 之后、每一层都要对完整状态做快照，每个 head 64 KiB。现在 kernel 记录每 token 的小型中间量，并推迟状态更新，使下一次调用只重放被接受的 token，并一次性写入已提交的状态。该 kernel 反过来用少量额外算术换取大量额外的内存流量，这正好契合我们的核心瓶颈问题：VRAM 受限而算力有余。
此前，KV-cache 自回归注意力 kernel 与新增的优化版 Comfy-Kitchen 注意力方案是分开的。Comfy-Kitchen 注意力以及 KV Cache 支持已在 ComfyUI v0.45 中加入，为博客前文所述的 LLM 专用注意力方案启用 INT8 注意力。这降低了 KV cache 的 VRAM 占用，同时也显著提升了长 prompt、system prompt 或已有上下文的 pre-fill 速度。你可以用 --use-ck-attention 启动 ComfyUI 来启用 Comfy-Kitchen 注意力。
自适应树 MTP
在 ComfyUI v0.40.0 中，MTP draft 最初被限制为单条 token 猜测链，对应对接下来 2 到 5 个 token 的单一猜测。这条链会 draft 一个 token，把它反馈给模型，然后再 draft 下一个 token，直到达到为 MTP 设定的深度值。默认情况下，把 MTP 设为 auto 会给出 2 到 5 个 token，并根据采样模式或临时的成功率测量进行自适应。这种链式方法的弱点在于，它把所有赌注都押在这单一猜测上，链上往下的每一次猜测只要出错就可能全部丢失。如果 MTP head 在深度 1 的首选是错的，整个步骤只产出一个 token，即使深度 1 的次优猜测是对的。
链式 MTP 只 draft 一条路径。树式 MTP 在每一层保留次优 token 作为兄弟节点，在一次遍历中验证整棵树，并根据命中情况重新规划树的形状。
避免链中这一潜在扭结的另一种方法是使用自适应树算法。在每一层，MTP head 的首选会像以前一样以每次 draft 的最佳猜测延续主链，但少量次优猜测会被保留为兄弟候选，用于随后的 MTP 验证。这形成一个小型 draft 树结构，这些额外的候选叶节点会被验证，但不会从它们继续 draft。理论上，在任何深度树结构中，你都可以 draft 任意数量的猜测，每增加一个猜测的收益递减，不过我们的线性 kernel 可以“免费”计算最多 M=8。这把猜测数量限制为七（M=8-1），以匹配我们的计算 kernel 和硬件。验证遍历现在最多携带八行；已提交的 token 加上七个 draft token，同时必须把每一个都当作若干可能分支的未来来处理。
提示词和输出文本的类型决定了树应当采用何种形态。可预测的输出（例如代码、列表，或重复提示词的描述）适合深度，也就是长链会被整体接受。创意写作则适合广度：首选往往错误，但正确的 token 通常在前三候选之内。自适应 MTP 会持续测量，因此当 MTP 设为 auto 模式时，它会边运行边测量并重新规划。它会按 draft 层级分别维护一个指数衰减计数，记录被接受的 token 处在哪个排名——是第一选择、第二、第三……还是都不在——并用 CUDA events 为 draft 阶段和 verify 阶段计时。每 32 步，它会对把七行 draft 排成一条宽度链的 64 种布局方式全部打分，依据条件命中概率估算每步的期望接受 token 数，再除以实测成本 verify_ms + depth × draft_ms。只有当胜出布局至少好 3% 时才会切换，以保证形态不会来回抖动；并且由于每种形态都使用同样的八行，为这一步捕获的 CUDA graphs 在切换后依然有效。
在 Generate Text 节点中将 MTP 值设为 auto 时，即可使用自适应树 MTP。
总结
在 ComfyUI v0.40.0 到 v0.45.0 之间，ComfyUI 中的 Qwen3.x 文本生成在相同硬件上显著变快。过去会内存不足的更长上下文现在可以顺利容纳，推测解码器现在也会根据它正在书写的文本自行调整形态。这些改进可在 NVIDIA 与 AMD、Windows 与 Linux 上运行，除了保持 ComfyUI 为最新版本之外，你无需额外步骤。喜欢本文的工作？我们正在招聘希望让我们所有产品中的开源生成式 AI 更快的工程师！查看开放职位 →