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

快速结论：该报错出现在 ComfyUI 中使用 HiDream-O1-Image 模型生成的第一采样步骤，原因是 commit b481bc15a 移除了 KV head 扩展，但 HiDream 的自定义 attention 函数（two_pass_attention）未处理 enable_gqa 参数，导致 Q（32 heads）与未扩展的 K/V（8 heads）尺寸不匹配。优先排查 HiDream 相关 attention 代码是否兼容 GQA。

问题场景

用户在 ComfyUI 中使用 HiDream-O1-Image 模型工作流时，每次生成均在第一个采样步骤触发失败。该问题在 commit b481bc15a（对应 PR #14772）之后出现，父版本及更早版本工作正常。即使更新到当前最新 master（5697b9701），问题依旧存在。

报错原文

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

原因分析

PR #14772 修改了 comfy/text_encoders/llama.py，移除了 KV head 的显式扩展逻辑，改为向下游 attention 回调传递 enable_gqa=True 参数。ComfyUI 的所有常规后端均已适配此变化，但 HiDream-O1 模型（位于 comfy/ldm/hidream_o1/model.py）通过 two_pass_attention 函数取代了标准的 optimized_attention。该函数签名是 def two_pass_attention(q, k, v, heads, **kwargs)，接收到的 enable_gqa 参数被 **kwargs 捕获但从未读取或处理。因此，在自回归/因果推理路径中，Q 具有 32 个 head，而 K/V 只有 8 个 head（未扩展），直接送入 comfy.ops.scaled_dot_product_attention 时触发维度 1 尺寸不匹配错误。该问题不依赖具体后端（即使使用 xformers 或 flash-attn，由于 HiDream 绕开了后端分发逻辑，问题依然存在）。

环境排查

Python版本：3.12.10
PyTorch版本：2.7.1+cu128（CUDA 12.8）
GPU：RTX 5090（在 Windows 11 上复现）
ComfyUI 提交号：在 b481bc15a 之后，已验证至 5697b9701
相关 PR：#14772 是引入变更的合并；#14865 是同一 commit 可能影响的另一个问题（flash-attn wrapper 参数数量不匹配）

解决步骤

打开 comfy/ldm/hidream_o1/attention.py 文件。
定位 make_two_pass_attention 函数内部的 two_pass_attention 定义。
在函数起始处插入 GQA KV-head 扩展逻辑：检查 k.shape[1] 是否不等于 H（即 Q 的 head 数），若是则按比例重复扩展 K 和 V 的 head 维度。具体代码：
```
if k.shape[1] != H:
    rep = H // k.shape[1]
    k = k.repeat_interleave(rep, dim=1)
    v = v.repeat_interleave(rep, dim=1)
```
保存文件并重启 ComfyUI。
注意：此修复同时适用于 comfy.ops.scaled_dot_product_attention 路径和 optimized_attention 路径，无需修改其他文件。官方确认此改法可解决问题，并计划以 PR 形式合入。

验证方法

加载 ComfyUI 自带的 HiDream-O1-Image 标准工作流（stock template workflow），执行一次完整生成。若采样步骤正常完成且不再抛出 RuntimeError: The size of tensor a (32) must match the size of tensor b (8) at non-singleton dimension 1，则修复生效。