执行摘要

• 一句话：为扩散模型添加逐层NVTX性能标记
• 推荐动作：值得精读，尤其关注 maybe_nvtx_range 的异常安全设计（直接调用 range_push/range_pop 绕过 str.format 陷阱）以及 ComponentResidencyManager 中钩子生命周期与组件执行的集成方式。

功能与动机

扩散子系统缺乏逐层 NVTX 支持，导致 nsys 捕获只显示 CUDA kernel 时间线，没有逻辑结构；SRT LLM 侧已通过 --enable-layerwise-nvtx-marker 提供了类似能力。本 PR 将相同设计引入 multimodal_gen，使用户在运行 LLM 和扩散负载时获得一致的性能分析心智模型。

实现拆解

1. 新建 python/sglang/multimodal_gen/runtime/utils/nvtx_pytorch_hooks.py，实现 maybe_nvtx_range 上下文管理器和 DiffusionNvtxHooks 类，通过 PyTorch forward pre/post hooks 在每个子模块调用时发出 NVTX 范围。
2. 在 ServerArgs 中添加 enable_layerwise_nvtx_marker 配置项。
3. 重构 PipelineStage.__call__ 和 PipelineExecutor 子类，集成 NVTX 门控逻辑（_apply_nvtx_gate、_should_use_stage_nvtx 等），使所有阶段自动参与 NVTX 标记。
4. 在 ComponentResidencyManager 中集成钩子生命周期管理，在组件使用时动态注册/注销钩子，并支持跨阶段隔离和重复模块去重。
5. 新增单元测试 test_nvtx_pytorch_hooks.py，覆盖 maybe_nvtx_range 异常安全、钩子注册、去重、形状收集和 set_enabled 切换等场景。

关键文件：

• python/sglang/multimodal_gen/runtime/utils/nvtx_pytorch_hooks.py（模块 NVTX钩子；类别 source；类型 dependency-wiring；符号 maybe_nvtx_range, DiffusionNvtxHooks, init, register_hooks）: 核心新增文件，实现了 NVTX 标记的基础设施：maybe_nvtx_range 上下文管理器和 DiffusionNvtxHooks 类。
• python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py（模块 执行器；类别 source；类型 core-logic；符号 _component_residency_request, _is_warmup_payload, _should_use_stage_nvtx, _run_stage_with_executor_hooks）: 管道执行器基类，新增 _component_residency_request、_is_warmup_payload、_should_use_stage_nvtx、_run_stage_with_executor_hooks 方法，将 NVTX 门控和范围包裹集成到阶段执行流程。
• python/sglang/multimodal_gen/runtime/managers/memory_managers/component_manager.py（模块 组件管理器；类别 source；类型 core-logic；符号 before_use, remove_nvtx_hooks_for_module, _enable_nvtx_for_use, _disable_active_nvtx）: 组件生命周期管理器，集成 NVTX 钩子的注册/启用/禁用，在 begin_use 和 end_use 时动态管理 DiffusionNvtxHooks 实例，确保钩子与组件使用同步。
• python/sglang/multimodal_gen/test/unit/test_nvtx_pytorch_hooks.py（模块 测试覆盖；类别 test；类型 test-coverage；符号 TestMaybeNvtxRange, test_disabled_returns_noop_context_manager, test_disabled_propagates_exception, test_disabled_does_not_call_nvtx）: 新增单元测试，覆盖 maybe_nvtx_range 的异常安全、DiffusionNvtxHooks 注册/移除/去重/形状收集/启用切换等 12 个用例，mock 方式避免 GPU 依赖。
• python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py（模块 阶段基类；类别 source；类型 core-logic；符号 _apply_nvtx_gate, current_use_nvtx）: PipelineStage 基类，新增 _apply_nvtx_gate 和 current_use_nvtx 属性，在 call 中解析 NVTX 门控，使所有阶段统一参与。

关键符号：maybe_nvtx_range, DiffusionNvtxHooks.register_hooks, DiffusionNvtxHooks.remove_hooks, DiffusionNvtxHooks.set_enabled, PipelineExecutor._run_stage_with_executor_hooks, PipelineStage._apply_nvtx_gate, ComponentResidencyManager._enable_nvtx_for_use

评论区精华

核心讨论围绕将 NVTX 钩子注册从 denoising 阶段提升为公用机制。mickqian 在 denoising.py 上评论建议将 _maybe_register_nvtx_hooks 推广到所有组件，促使作者重构为 ComponentResidencyManager 统一管理。后续测试失败（mock server_args 缺少字段、子类未调用 super().init 导致属性缺失）也被逐步修复，最终通过了 NV CI。

• 将 NVTX 钩子注册提升为通用工具 (design): 作者重构了设计，将钩子注册从 denoising 阶段迁移到 ComponentResidencyManager，通过 _enable_nvtx_for_use 在组件使用时动态注册，去掉了硬编码的 per-stage 钩子声明。

风险与影响

• 风险：NVTX 钩子可能引入额外 CPU 开销，但通过 set_enabled() 默认关闭，仅在主动开启时生效。钩子生命周期与组件管理绑定，但已通过 remove_hooks 和实例跟踪处理了重用和删除场景。执行流重构合并了 before_stage/after_stage 到 _run_stage_with_executor_hooks，可能影响非 NVTX 路径，但通过模块化设计降低了风险。
• 影响：用户可通过 --enable-layerwise-nvtx-marker 开启，默认行为不变。开启后，Nsight Systems 捕获的 trace 会显示 stage、denoising step 和子模块的嵌套范围，帮助定位性能瓶颈。团队获得与 SRT 侧统一的性能分析方案，降低了维护两套 NVTX 工具的成本。
• 风险标记：性能开销可控, 钩子生命周期管理, 执行流重构

关联脉络

• PR #25848 [diffusion] Add CFG gating for denoising: 同属扩散子系统性能优化系列，本 PR 的 NVTX 标记可与 CFG gating 配合进一步分析 denoising 阶段的时序。
• PR #25847 [diffusion] Cache fp32 layernorm params: 扩散子系统性能优化，与 NVTX 标记无关但同属 multimodal_gen 模块，有助于理解整体演进。