Prhub

#21955 [diffusion] chore: fix stage profiler for multi-stage denoising

原始 PR 作者 mickqian 合并时间 2026-04-03 01:16 文件变更 5 提交数 2 评论 2 代码增减 +15 / -8

执行摘要

修复多阶段去噪场景下性能分析器记录步骤时序错误的问题。

PR标题和提交信息表明需要修复多阶段去噪(multi-stage denoising)场景下的stage profiler问题。从代码变更看,原逻辑依赖stage_name是否包含'denoising_step_'字符串来判断是否为去噪步骤,但在多阶段去噪场景下,某些阶段的名称可能不以此前缀开头,导致步骤执行时间未被正确记录到metrics中。

该PR值得精读,展示了性能分析工具如何适配复杂场景(多阶段去噪)的设计决策。重点关注:

1) StageProfiler如何通过record_as_step标志解耦阶段名称约定和业务逻辑。
2) 从index-based到顺序记录的简化设计。
3) 同步逻辑(SGLANG_DIFFUSION_SYNC_STAGE_PROFILING)与步骤记录的关联。

讨论亮点

review中gemini-code-assist[bot]指出_should_record_as_step方法从使用'denoising_step_' in self.stage_name改为self.stage_name.startswith('denoising_step_'),这是一个行为变更:如果存在包含但不以该字符串开头的自定义阶段名称,将不再被自动归类为步骤。但作者未回复此评论,从最终代码看仍采用了startswith方案,可能认为这是可接受的精确化改进。

实现拆解

实现方案分为两部分:

1) 在四个去噪阶段文件(denoising.py、denoising_av.py、denoising_dmd.py、mova.py)的forward方法调用StageProfiler时,显式传入record_as_step=True参数。
2) 在perf_logger.py中重构StageProfiler:新增record_as_step实例变量和_should_record_as_step方法,将原record_steps方法重命名为record_step并移除index参数,修改同步和记录逻辑以使用新的判断条件。

文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/utils/perf_logger.py multimodal_gen/runtime/utils modified 8.0
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py multimodal_gen/pipelines_core modified 6.0
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py multimodal_gen/pipelines_core modified 6.0

关键符号

StageProfiler.__init__ StageProfiler._should_record_as_step StageProfiler.__exit__ PerfMetrics.record_step

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

评论区精华

阶段名称匹配逻辑从包含改为前缀匹配的行为变更 正确性

gemini-code-assist[bot] 指出 startswith 替换 in 可能影响历史自定义阶段名称的自动归类

结论:作者未回应,代码采用 startswith 方案,可能认为这是可接受的精确化 · 已解决

风险与影响

主要风险:

1) 行为变更风险:从包含匹配改为前缀匹配,若存在历史自定义阶段名称包含但不以'denoising_step_'开头,将不再被自动记录为步骤,可能导致性能数据缺失。
2) 兼容性风险:record_steps方法签名变更(移除index参数),但调用方已同步更新,影响范围可控。
3) 逻辑正确性:新增record_as_step参数需确保在所有多阶段去噪场景正确传递,否则仍可能漏记录。

对用户:修复后多阶段去噪的性能分析数据将更准确,有助于调试和优化。对系统:性能监控功能更完善,但变更仅影响分析记录,不影响核心推理逻辑。对团队:统一了步骤记录方式,简化了后续维护,但需注意行为变更可能影响现有自定义阶段名称的监控。

行为变更风险 依赖参数传递正确性

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论