Prhub

#22374 [diffusion] fix: fix cache dit refresh none mask

原始 PR 作者 mickqian 合并时间 2026-04-09 11:58 文件变更 3 提交数 5 评论 3 代码增减 +296 / -28

执行摘要

修复缓存 DIT 刷新时 mask 为 None 的 bug,并优化分布式预热图像处理。

从代码变更推断,修复当scm_presetNone时,steps_computation_mask不应调用cache_dit.steps_mask函数,否则可能导致TypeError或逻辑错误。PR body未提供详细描述,但标题和变更表明这是一个针对扩散模型缓存集成的bugfix。

建议开发者和维护者精读cache_dit_integration.py的变更,理解缓存mask处理的条件逻辑设计;同时关注scheduler.py中的分布式预热优化,这对于多GPU或集群部署有参考价值。

讨论亮点

Review评论为空,表明PR在无讨论情况下由作者合并,变更可能被认为较为直接或紧急。

实现拆解

实现分为三部分:

1) 在cache_dit_integration.py中,修改refresh_context_on_transformerrefresh_context_on_dual_transformer函数,仅在scm_preset不为None时调用cache_dit.steps_mask生成mask,否则设为None;
2) 在scheduler.py中,重构prepare_server_warmup_reqs方法,提取_prepare_shared_warmup_image_path私有方法,使用广播同步预热图像路径,优化分布式环境下的预热过程;
3) 新增单元测试文件test_cache_dit_integration.py,通过模拟依赖提供测试覆盖,验证缓存DIT逻辑的正确性。

文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py diffusion/cache modified 8.0
python/sglang/multimodal_gen/runtime/managers/scheduler.py scheduling modified 6.0
python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py test added 4.0

关键符号

refresh_context_on_transformer refresh_context_on_dual_transformer prepare_server_warmup_reqs _prepare_shared_warmup_image_path

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

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

风险包括:

1) 缓存DIT逻辑修改可能影响扩散模型推理的正确性,尤其是当scm_preset为None时的边缘情况;
2) 分布式预热改动引入新的同步逻辑(如广播操作),增加复杂性和潜在死锁风险;
3) 新增测试覆盖了模拟环境,但需确保与实际生产环境的兼容性和覆盖率。

影响范围:

1) 用户:修复bug后提升扩散模型服务的稳定性和可靠性;
2) 系统:优化预热过程,减少多进程冲突,提高分布式部署性能;
3) 团队:代码更健壮,新增测试增强维护性,但改动涉及核心缓存和调度模块,需注意回归测试。

缓存逻辑变更 分布式同步风险

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论