【硬核排障 & 猴子补丁 & 幽灵节点】SageAttention GQA 崩溃与 Flash Attention 兼容性修复全记录

【硬核排障 & 猴子补丁 & 幽灵节点】SageAttention GQA 崩溃与 Flash Attention 兼容性修复全记录

环境 : ComfyUI v0.28.0 + Windows + CUDA + SageAttention

问题 : RuntimeError: The size of tensor a (32) must match the size of tensor b (8)

根因 : ComfyUI v0.28.0 全面启用 GQA，SageAttention CUDA kernel 不兼容

修复: 非侵入式幽灵节点动态拦截 + 原生 PyTorch SDPA 回退
Release v0.28.0 · Comfy-Org/ComfyUI · GitHub

---

一、问题现象

1.1 核心报错

RuntimeError: The size of tensor a (32) must match the size of tensor b (8)
    at non-singleton dimension 1

1.2 完整堆栈

File "comfy/ldm/modules/attention.py", line 544, in attention_pytorch
    out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=mask, ...)
File "comfy/ops.py", line 58, in scaled_dot_product_attention
    return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
File "sageattention/core.py", line 152, in sageattn
    out = _C_nn.scaled_dot_product_attention(...)

1.3 触发条件

• ComfyUI 版本: v0.28.0+（v0.27.0 正常）
• 模型: 使用 GQA（Grouped Query Attention）结构的文本编码器，如 T5-XXL、Qwen、LLaMA
• 注意力后端 : SageAttention（sageattention 包已安装）
• 具体维度 : q_heads=32, k_heads=8（或类似倍数关系）
• 受影响的工作流示例: ++utility_z_image_turbo_2k_upscaler.app.json （图片高清放大）++

Image Upscale: Z-Image-Turbo 2K - ComfyUI Workflow

---

二、根因分析

https://github.com/Comfy-Org/ComfyUI/releases/tag/v0.28.0

2.1 ComfyUI v0.28.0 的关键变更

根据 ComfyUI v0.28.0 Release Notes 的 "Performance & Stability" 部分：

GQA attention: GQA on all attention backends; dropped PyTorch 2.4

这意味着：

• 文本编码器阶段 （attention_pytorch 路径）：T5/LLaMA 编码器启用了 GQA
• 采样阶段 （attention_flash 路径）：DiT/UNet 也支持 GQA

2.2 SageAttention 的局限性

SageAttention 的 CUDA kernel 对输入张量有严格限制：

限制项	SageAttention 要求	GQA 实际情况
head_dim	必须是 64/96/128	T5 的 head_dim=64 ✅
num_heads	q/k/v 头数必须相同	q_heads=32, k_heads=8 ❌
attn_mask	不支持复杂 mask	文本编码器有 mask ❌

核心矛盾 : SageAttention kernel 在计算 attention 时，假设 q 和 k 的 shape 在 dim=1（头数维度）上完全匹配，直接做 matmul。GQA 结构中 q_heads ≠ k_heads，导致 RuntimeError。

2.3 为什么 v0.27.0 正常？

【硬核排障 & 猴子补丁 & 幽灵节点】Flash Attention failed, using default SDPA

v0.27.0 的 attention_pytorch 路径在处理文本编码器时，可能：

• 未启用 GQA（skip_reshape=True 时未传入 GQA 参数）
• 或者通过 torch.nn.functional.scaled_dot_product_attention 的原生 GQA 支持自动处理了广播

v0.28.0 显式传入 **gqa_kwargs，暴露了 SageAttention 的不兼容。

---

三、修复方案对比

3.1 方案一：禁用 SageAttention（不推荐）

set COMFYUI_DISABLE_SAGE_ATTENTION=1
python main.py

缺点: 一刀切，所有工作流都损失 SageAttention 的加速收益。

3.2 方案二：降级 ComfyUI（不推荐）

git checkout v0.27.0

缺点: 无法使用 v0.28.0 的新功能和性能优化。

3.3 方案三：非侵入式幽灵节点（推荐 ✅）

核心思想 : 在运行时拦截 torch.nn.functional.scaled_dot_product_attention，检测不兼容条件后回退到原生 PyTorch 实现。

优势:

• ✅ 不修改 ComfyUI 核心代码
• ✅ 不修改虚拟环境包
• ✅ 只在需要时回退，兼容工作流仍享受 SageAttention 加速
• ✅ 零 UI footprint（幽灵节点）
• ✅ 通过安装 1 个插件即可修复

手动安装：

https://github.com/love530love/comfyui-sage-guard

cd ComfyUI/custom_nodes
git clone https://github.com/love530love/comfyui-sage-guard.git

或从 ComfyUI 官方节点商店安装：

ComfyUI注册表

comfy node install comfyui-sage-guard

---

四、实现原理

4.1 劫持点选择

SageAttention 通过 monkey patch 替换 torch.nn.functional.scaled_dot_product_attention 为 sageattention.core.sageattn。我们的劫持需要晚于 SageAttention 的替换，才能覆盖其调用。

调用链:
comfy.ops.scaled_dot_product_attention
    → torch.nn.functional.scaled_dot_product_attention  [已被 Sage 替换为 sageattn]
        → sageattention.core.sageattn
            → _C_nn.scaled_dot_product_attention (CUDA kernel, 崩溃)

劫持后:
comfy.ops.scaled_dot_product_attention
    → SageGuardWrapper.__call__()  [我们的包装器]
        → 检测: 如果 GQA/不兼容 → _fallback_attention (手动实现)
        → 否则 → 原始 sageattn (SageAttention 加速)

4.2 兼容性检测逻辑

def _should_fallback(q, k, v, attn_mask, **kwargs):
    # 1. 非 CUDA 设备
    if q.device.type != 'cuda':
        return True

    # 2. 存在 attn_mask（Sage 对复杂 mask 支持有限）
    if attn_mask is not None:
        return True

    # 3. head_dim 不支持
    head_dim = q.size(-1)
    if head_dim not in {64, 96, 128}:
        return True

    # 4. GQA 检测：q 和 k 的头数不一致
    if q.size(1) != k.size(1):
        return True

    return False

4.3 GQA 手动回退实现

PyTorch 原生 SDPA 支持 GQA 是通过广播 实现的：k/v 的每个头被多个 q 头共享。torch.matmul 不直接支持这种广播，需要手动 repeat_interleave：

def _fallback_attention(q, k, v, attn_mask=None, dropout_p=0.0, 
                        is_causal=False, scale=None, **kwargs):
    q_heads, k_heads = q.size(1), k.size(1)

    # GQA: 将 k/v 重复以匹配 q 的头数
    if q_heads != k_heads:
        repeat_factor = q_heads // k_heads
        k = k.repeat_interleave(repeat_factor, dim=1)
        v = v.repeat_interleave(repeat_factor, dim=1)

    if scale is None:
        scale = q.size(-1) ** -0.5

    scores = torch.matmul(q, k.transpose(-2, -1)) * scale

    if attn_mask is not None:
        scores = scores + attn_mask

    if is_causal:
        causal_mask = torch.triu(
            torch.ones(scores.size(-2), scores.size(-1), 
                      device=scores.device, dtype=torch.bool),
            diagonal=1
        )
        scores = scores.masked_fill(causal_mask, float('-inf'))

    attn = torch.softmax(scores, dim=-1)
    if dropout_p > 0:
        attn = torch.nn.functional.dropout(attn, p=dropout_p, training=True)

    return torch.matmul(attn, v)

数学正确性 : repeat_interleave 在 GQA 中是完全正确的，因为 GQA 的本质就是多个 query 头共享同一个 key/value 头。重复后计算结果与原生 SDPA 的广播实现等价。

4.4 多重替换确保无死角

_guard = SageGuardWrapper()
_guard.patch()  # 替换 torch.nn.functional.scaled_dot_product_attention

# 同步替换 comfy.ops
import comfy.ops
comfy.ops.scaled_dot_product_attention = _guard

# 同步替换 sageattention.core
import sageattention.core as sage_core
sage_core.sageattn = _guard

---

五、踩坑记录

5.1 坑一：替换时机

问题 : 在模块导入时保存 _NATIVE_SDPA = F.scaled_dot_product_attention，但 SageAttention 在之后才完成替换，导致 _NATIVE_SDPA 实际上保存的是 Sage 版本。

解决: 完全放弃依赖外部 SDPA 引用，fallback 使用纯手动实现。

5.2 坑二：torch._C._nn.scaled_dot_product_attention 也被替换

问题: 尝试调用 PyTorch C++ 原生 SDPA 作为 fallback，但发现它也被 SageAttention 劫持了。

解决: 直接用手动实现，不依赖任何外部 SDPA 函数。

5.3 坑三：**kwargs 参数透传

问题 : v0.28.0 新增了 enable_gqa、softmax_scale 等参数，wrapper 如果不接受 **kwargs 会报 TypeError。

解决 : wrapper 的 __call__ 方法必须签名包含 **kwargs，并在所有调用处透传。

5.4 坑四：Flash Attention 的 softmax_scale 参数

问题 : ComfyUI-FixFlashAttnSchema 插件（修复 schema_.has_value() bug）的 monkey patch 不支持 v0.28.0 新增的 softmax_scale 参数。

解决 : 在 flash_attn_wrapper_patched 中添加 **kwargs 吸收新参数。

---

六、相关链接

6.1 修复插件

插件	作用	GitHub
comfyui-sage-guard	SageAttention GQA 兼容性回退	https://github.com/love530love/comfyui-sage-guard
ComfyUI-FixFlashAttnSchema	Flash Attention schema bug 修复 + v0.28.0 兼容	https://github.com/love530love/ComfyUI-FixFlashAttnSchema

6.2 参考资源

• ComfyUI v0.28.0 Release Notes - GQA 全面启用说明
• SageAttention GitHub - SageAttention 官方仓库
• PyTorch SDPA Documentation - 原生 SDPA GQA 支持
• ComfyUI Registry Publishing Guide - 节点发布教程

6.3 相关博客

• 【硬核排障 & 猴子补丁 & 幽灵节点】Flash Attention failed, using default SDPA - 前期 Flash Attention 修复记录

---

七、总结

7.1 核心结论

ComfyUI v0.28.0 的 GQA 全面启用是一个破坏式变更，对依赖 SageAttention 的本地环境造成了兼容性冲击。官方云环境可能使用了不同的 attention backend 或配置，因此未暴露此问题。

7.2 修复哲学

坚持非侵入式修复：

1. 不修改 ComfyUI 核心代码
2. 不修改虚拟环境包
3. 通过幽灵节点（ghost node）在运行时动态拦截
4. 只在必要时回退，最大化保留性能收益

7.3 未来维护

• 跟踪 ComfyUI 版本更新，及时调整兼容性检测逻辑
• 跟踪 SageAttention 版本，如果官方支持 GQA 则可移除部分 fallback
• 保持插件的轻量化和零依赖（除 torch 外）

修复后工作流运行效果：

出图效果（8K）

放大查看效果

日志打印正常

---

本文撰写于 2026-07-18，基于 ComfyUI v0.28.0 + SageAttention + Windows CUDA 环境。