Prhub

#20706 [diffusion] Unify `TeaCacheParams` and `WanTeaCacheParams`

原始 PR 作者 eitanturok 合并时间 2026-03-28 09:51 文件变更 7 提交数 21 评论 13 代码增减 +202 / -112

执行摘要

统一 TeaCache 参数类,移除 Wan 特定实现,支持动态系数回调。

根据PR body,动机是“Unifies WanTeaCacheParams into the standard TeaCacheParams class to make it easier to handle teacache for different models.”,旨在消除Wan模型的特殊参数类,提供更通用的系数选择方式。这是一个独立的变更,与#19957相关。

建议工程师精读此PR,重点关注:

  • TeaCacheParams类的设计变化,特别是get_coefficients方法和coefficients_callback机制,展示了如何通过回调支持动态系数选择。
  • CFG逻辑的简化处理(get_skip_boundaries方法),这对理解扩散模型中的缓存优化和步骤管理有参考价值。
  • 单元测试的设计,验证兼容性和回调优先级,可作为类似重构的测试范例。
讨论亮点

review中主要讨论点在wanvideo.py的CFG处理逻辑:

  • adarshxs询问“could you explain the logic here?”,关注新代码的实现细节。
  • eitanturok详细解释旧代码如何预加倍步骤造成混淆,新实现通过get_skip_boundaries方法直接处理CFG模式,简化了逻辑。
  • adarshxs回复“understood. got confused since the code initially did this for non cfg mode”,表示理解,讨论已解决。
    最终mickqian批准PR,表明争议点已澄清。

实现拆解

实现分为三个核心部分:

  1. 参数类扩展(teacache.py):添加coefficients_callback、start_skipping、end_skipping字段和get_coefficients方法到TeaCacheParams类,移除WanTeaCacheParams类。
  2. 模型适配(wan.py):为Wan模型定义系数回调函数(_wan_1_3b_coefficients和_wan_14b_coefficients),并更新采样参数类使用新的TeaCacheParams。
  3. 运行时逻辑调整(wanvideo.py):修改should_skip_forward_for_cached_states方法,使用get_skip_boundaries处理CFG逻辑,简化步骤边界计算。此外,sampling_params.py扩展了JSON序列化支持callable,test_sampling_params.py添加单元测试验证回调优先级和兼容性。
文件 模块 状态 重要度
python/sglang/multimodal_gen/configs/sample/teacache.py diffusion/config modified 8.0
python/sglang/multimodal_gen/configs/sample/wan.py diffusion/config modified 7.0
python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py diffusion/runtime modified 7.0
python/sglang/multimodal_gen/test/unit/test_sampling_params.py diffusion/test modified 6.0

关键符号

get_coefficients _wan_1_3b_coefficients _wan_14b_coefficients get_skip_boundaries should_skip_forward_for_cached_states

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

CFG 处理逻辑简化 正确性

adarshxs 询问新代码逻辑,eitanturok 解释旧代码如何预加倍步骤造成混淆,新实现通过 get_skip_boundaries 方法简化 CFG 处理。

结论:理解并批准,逻辑已澄清,讨论解决。 · 已解决

风险与影响

技术风险包括:

  • 回调序列化风险:coefficients_callback作为callable,在JSON序列化时可能出错,但sampling_params.py的_json_safe函数已添加支持,通过返回模块限定名来缓解。
  • 步骤边界回归:start_skipping和end_skipping参数的计算变更可能影响teacache跳过行为,尤其是在CFG模式下;但单元测试test_wan_teacache_boundaries_match_legacy_behavior验证了与旧行为的一致性。
  • 系数选择逻辑变更:移除硬编码系数改为回调,可能引入运行时错误,如回调函数返回类型不匹配;但测试覆盖了回调优先级和系数获取。
    总体风险较低,通过测试和代码审查缓解。

影响范围:

  • 用户影响:Teacache参数使用更统一,Wan模型的行为应保持向后兼容,通过测试确保无功能回归。
  • 系统影响:减少类重复(移除WanTeaCacheParams),提高代码复用性和可维护性;新回调机制为其他模型扩展teacache提供灵活框架。
  • 团队影响:简化参数处理逻辑,便于后续开发;CFG逻辑的清晰化有助于减少未来代码混淆。
    影响程度为中,主要影响diffusion模块内配置和运行时逻辑。
回调序列化风险 步骤边界回归 系数选择逻辑变更

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论