Flash/Sage varlen does not work with torch.compile

快速结论：这个报错通常出现在你对 Diffusers 的 Flux 等模型使用 torch.compile，同时又启用了 Flash/Sage 的 varlen（变长）注意力后端时。优先排查是否用了 varlen API，以及能否切换到 static-length（定长）API 或使用定长序列推理。

适用环境：Issue 中已确认的场景：Python 3.12 虚拟环境（路径含 venv/3.12）、Diffusers 的 transformer_flux.py、PyTorch 2.x 的 torch._dynamo、以及 Flash/Sage varlen 注意力后端。未在 Issue 中确认具体的 PyTorch、CUDA、显卡型号，不做补写。

最快修复方案：暂无确认的一步修复方案。Issue 中被提到的“可优先尝试”的 workaround 是：如果你不使用 attention mask、或只在单一分辨率生成，可以切换到 static-length API；如果生成分辨率集合固定，可用 static-length API 做一次编译预热（warmup），之后复用（可结合 torch MegaCache 与 torch.compile(dynamic=True)）。

注意事项：Issue 回复中明确说明 Flash/Sage 的 varlen 相关支持“非常实验性”，官方文档仍在补充中。上述 workaround 属于建议，并非已验证的一步修复；capture_scalar_outputs / TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1 在 Issue 中只是提出未验证的候选（评论者表示尚未测试），不能当作确定结论。

问题场景

用户在使用 Diffusers 时，尝试在 Flux 模型（diffusers/models/transformers/transformer_flux.py）上用 torch.compile 编译推理流程，同时启用了 Flash/Sage 的 varlen（变长）注意力后端。该问题在尝试对包含 varlen 注意力的 Flux block 做 Dynamo 图捕获时触发，属于“编译 + varlen 注意力”组合场景。

报错原文

UserWarning: Dynamo detected a call to a `functools.lru_cache`-wrapped function. Dynamo ignores the cache wrapper and directly traces the wrapped function.
W0718 20:32:55.420000 1 .../torch/_dynamo/variables/tensor.py:1048] [0/0] Graph break from `Tensor.item()`, consider setting:
    torch._dynamo.config.capture_scalar_outputs = True
or:
    env TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1
to include these operations in the captured graph.
Graph break: from user code at:
  File ".../diffusers/models/transformers/transformer_flux.py", line 733, in forward
    encoder_hidden_states, hidden_states = block(
  File ".../diffusers/models/transformers/transformer_flux.py", line 456, in forward
    attention_outputs = self.attn(
  File ".../diffusers/models/transformers/transformer_flux.py", line 343, in forward
    return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
  File ".../diffusers/models/transformers/transformer_flux.py", line 117, in __call__

原因分析

根据 Issue 中维护者的回复，Flash attention 和 Sage attention 的 varlen API 需要一个 max sequence length 参数。当前实现会在运行过程中通过查看 attention mask（如果提供了）或序列长度来临时推断该值，因此需要调用 .item() 得到一个整数传给 API。而 .item() 会产生 graph break，Dynamo 无法 trace 这类操作，因此在 torch.compile 下会失败或退化为图断裂。这就是 varlen 与 torch.compile 不兼容的直接原因。

环境排查

确认 Python 版本（Issue 中为 3.12）。
确认 Diffusers 版本及是否为包含 Flash/Sage varlen 注意力后端实现的版本。
确认 PyTorch 版本是否带 Dynamo（2.x）。
确认是否启用了 varlen 注意力 API（Flash attention / Sage attention）。
确认是否提供了 attention mask，以及是否在单一或固定分辨率集合下生成。
确认是否使用了 torch.compile，以及是否设置了 dynamic=True。
CUDA 版本、显卡型号在 Issue 中未确认，请按本机实际情况排查。

解决步骤

先判断你是否真的需要 varlen：如果你不使用 attention mask、或只在单一分辨率生成，优先切换到 static-length（定长）注意力 API，避免 .item() 推断路径。
如果生成的分辨率集合是固定的，可以用 static-length API 先做一次编译预热（warmup），之后复用编译结果；Issue 中提到可结合 torch MegaCache，并可尝试 torch.compile(dynamic=True) 参数来减少重编译。
如果必须用 varlen，可考虑注册自己的注意力后端（通过 _AttentionBackendRegistry.register），让实现变成“序列长度已知”的形式——序列长度可以在 transformer 的 forward 之外很容易地算好，从而绕开在 forward 内调用 .item()。此方案在 Issue 中仅为建议，未给出完整实现。
Issue 中有人提出设置 torch._dynamo.config.capture_scalar_outputs = True 或环境变量 TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1，但评论者表示尚未测试，只能作为候选方向尝试，不能视为已验证修复。

验证方法

切换为 static-length API 或完成定长预热后，重新运行带 torch.compile 的推理流程，观察是否不再出现 Graph break from Tensor.item() 及 transformer_flux.py 中的 graph break 堆栈；同时确认生成结果正常、无编译报错。若仍出现相同 graph break，说明仍在走 varlen 的动态推断路径。