Flash/Sage varlen does not work with torch.compile

快速结论：当你在 Diffusers 的 Flux 注意力处理器中启用 Flash/Sage varlen（变长）注意力并叠加 torch.compile 时，varlen 实现为在运行时推导最大序列长度会触发 Tensor.item()，被 Dynamo 判定为 Graph break，从而无法完整编译。优先排查是否使用了 varlen 注意力接口、是否传入 attention mask，以及能否改用静态长度（static-length）接口。

适用环境：Issue 中确认的环境为：Diffusers（涉及 transformer_flux.py 的注意力处理）、PyTorch 2.x（torch._dynamo / torch.compile）、Python 3.12 虚拟环境、在 InferenceSh GPU 应用上运行。其余 CUDA、显卡型号、具体 Diffusers 版本号在 Issue 中未给出确认信息。

最快修复方案：暂无确认的一步修复方案。Issue 中维护者给出的是“可优先尝试”的绕过手段：如果不需要 attention mask 或只在单一分辨率下生成，改用静态长度注意力接口；使用编译器预热 + torch MegaCache（可配合 dynamic=True）避免重复编译。

注意事项：维护者明确表示该特性“非常实验性”，相关文档与默认开关仍在推进中（曾提及可能在 2.9 默认开启，但 Issue 中并非已验证发布结论）。Issue 最终以 stale 自动关闭，未在讨论串里确认修复落地。

问题场景

用户在 Diffusers 中试用 PR 引入的 Flash/Sage varlen 注意力接口（setter 函数与上下文写法），运行 Flux 模型（diffusers/models/transformers/transformer_flux.py 中的 FluxAttnProcessor / 注意力处理器调用链）时，同时启用 torch.compile，期望整个前向可被编译，但在推理开始后立即出现 Dynamo 相关的 Graph break 与报错日志，编译无法按预期完成。

报错原文

UserWarning: Dynamo detected a call to a `functools.lru_cache`-wrapped function. Dynamo ignores the cache wrapper and directly traces the wrapped function.

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__
...

原因分析

最可能的原因是：Flash attention 与 Sage attention 的 varlen API 需要一个 max_seq_len 之类的整数参数，而当前实现是在前向过程中“即时”推断该值——通过查看 attention mask（若提供）或序列长度来推导。要在运行时把张量转成可传给 API 的整数值，就必须调用 .item()，而 .item() 会导致 Dynamo 发生 Graph break，因此 varlen 路径无法被 torch.compile 完整捕获。

此外日志中的 functools.lru_cache 提示是 Dynamo 对被缓存函数的追踪行为说明，属于伴随现象，不是本次编译失败的核心原因。

环境排查

确认正在使用的注意力处理器是否为 Flash/Sage 的 varlen 接口，而非 static-length 接口。
确认是否传入了 attention mask；若传入，varlen 路径会依赖 mask 推导序列长度。
确认是否只在一个固定分辨率（或固定分辨率集合）下生成。
确认 PyTorch 版本及 torch.compile 的调用参数（是否使用 dynamic=True）。
确认 Diffusers 版本是否包含该 varlen 注意力实现。
若参考日志中的 scalar 捕获提示，确认是否设置过 torch._dynamo.config.capture_scalar_outputs 或环境变量 TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1（Issue 中仅作为建议提出，作者表示尚未测试）。

解决步骤

先确认自己确实走的是 varlen 路径：检查 Flux 注意力处理器的配置与调用，确认使用的是 Flash/Sage 的变长接口。
若不需要 attention mask，或只在单一分辨率下生成：改用静态长度（static-length）注意力 API，绕开运行时的 .item() 推导。
若需要覆盖固定的一组分辨率：先用静态长度 API 做编译器预热（compiler warmup），再配合 torch MegaCache 复用已编译结果，避免反复重编译；可同时尝试在 torch.compile 中使用 dynamic=True 参数。
若以上都不适用：通过 _AttentionBackendRegistry.register 注册自定义注意力后端，把序列长度做成“已知”的——在 transformer 前向之外先算好序列长度，再传入实现，从而避免在前向内触发 .item()。
可优先尝试（未验证）：按日志提示设置 torch._dynamo.config.capture_scalar_outputs = True 或环境变量 TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1，看是否能减少 Graph break。Issue 中维护者引用该建议后，作者回复“尚未测试”。

验证方法

重新运行相同的 Flux 推理 + torch.compile 组合，观察是否仍出现 Graph break from `Tensor.item()` 以及指向 transformer_flux.py 的 user code 堆栈；若 Graph break 消失、编译图能完整生成且推理结果正常，则说明该路径已规避问题。若改用静态长度接口，还应确认注意力输出与未编译时的结果一致。