执行摘要

• 一句话：Cosmos3 解码后处理 GPU 化
• 推荐动作：该 PR 是聚焦且高效的性能优化，设计简洁、收益明确。值得阅读的关键点：如何用静态方法拆分 GPU 后处理和 CPU 后处理，以及如何通过控制流避免不必要的 CPU 拷贝。建议合并，并在未来补充单元测试以确保后处理稳定性。

功能与动机

Cosmos3 解码阶段原通过 diffusers.video_processor 将 GPU tensor 转换为 CPU float32 numpy 数组，此操作为不必要的 CPU 拷贝，尤其在 guardrails 关闭时浪费显存带宽和主存带宽。PR 作者通过基准测试确认该后处理是解码阶段的主要瓶颈，优化后解码阶段耗时降低约 20%。

实现拆解

1. 移除 VideoProcessor 初始化：在 __init__ 中删除对 diffusers.video_processor 的导入和 self.video_processor 的创建，该工具类不再需要。
2. 新增 _postprocess_tensor 静态方法：对解码后的 GPU tensor 执行 decoded * 0.5 + 0.5 去归一化、clamp(0, 1) 截断、.float() 类型转换，全程在 GPU 上完成，返回 tensor。
3. 新增 _postprocess_video_np 静态方法：在需要 numpy 输出（guardrails 路径）时，将 GPU tensor 调整为 [B, H, W, C] 格式并拷贝到 CPU，返回 numpy 数组。图像模式压榨时间维度，视频模式保留时间维度。
4. 修改 forward 控制流：无条件调用 _postprocess_tensor 得到 GPU tensor；仅当 guardrails 启用时才调用 _postprocess_video_np 转 numpy 并调用安全检测；非 guardrails 且非图像模式下，仅记录后处理 tensor 形状日志。

关键文件：

• python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py（模块 扩散引擎；类别 source；类型 data-contract；符号 _postprocess_tensor, _postprocess_video_np）: 唯一修改的文件，包含所有变更：删除 VideoProcessor 初始化、新增两个后处理静态方法、修改 forward 控制流。

关键符号：_postprocess_tensor, _postprocess_video_np

关键源码片段

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py

唯一修改的文件，包含所有变更：删除 VideoProcessor 初始化、新增两个后处理静态方法、修改 forward 控制流。

# 新增：GPU tensor 后处理，全程在 GPU 上完成
# 将解码后的像素值从 [-1, 1] 映射到 [0, 1] 并截断、转换类型
@staticmethod
def _postprocess_tensor(decoded: torch.Tensor) -> torch.Tensor:
    return (decoded * 0.5 + 0.5).clamp(0, 1).float()
​
# 新增：仅在需要 numpy 时（guardrails 路径）才拷贝到 CPU
# 调整维度顺序为 [B, H, W, C]（图像）或 [B, T, H, W, C]（视频）
@staticmethod
def _postprocess_video_np(video: torch.Tensor, is_image_gen: bool) -> np.ndarray:
    if is_image_gen:
        # squeeze 时间维度，permute 为 BHWC
        return video.squeeze(2).permute(0, 2, 3, 1).cpu().numpy()
    # 视频保留时间维度，permute 为 BTHWC
    return video.permute(0, 2, 3, 4, 1).cpu().numpy()
​
# forward 方法中关键变更：
# 之前：无条件调用 self.video_processor.postprocess/postprocess_video 输出 numpy
# 现在：先调用 _postprocess_tensor 保持 GPU tensor，仅 guardrails 路径才转 numpy
output = self._postprocess_tensor(decoded)
if self._guardrails and batch.use_guardrails is not False:
    # ... 导入 check_video_safety
    output = self._postprocess_video_np(output, is_image_gen) # 延迟 CPU 拷贝
    # 安全检测……
elif not is_image_gen:
    self.log_info(f"Postprocessed video tensor shape: {output.shape}")

评论区精华

未发现实质性 review 讨论。仅有的两条评论来自 gemini-code-assist[bot]（配额超限警告）和作者自动触发的 CI 重跑指令。

• 暂无高价值评论线程

风险与影响

• 风险：
  • 功能回归风险：原有路径在 guardrails 关闭时输出 numpy 数组，新路径输出 GPU tensor。下游消费方（如输出保存）必须能够处理 tensor 输入。PR 描述声称下游已支持 tensor 到 uint8 的转换，但未提供确认证据。实际测试显示 PNG 输出保存成功，但风险较低。
  • 安全检测路径：guardrails 路径仍正确转换为 numpy，安全检测逻辑不受影响。
  • 性能风险：无，优化方向正确且基准测试验证了收益。
  • 缺少测试覆盖：本次变更未包含测试文件，建议在后续 PR 中补充单元测试覆盖 _postprocess_tensor 和 _postprocess_video_np 的正确性。
• 影响：
  • 用户：Cosmos3 推理延迟降低约 3%，解码阶段加速约 20%，用户体验提升。
  • 系统：减少 GPU→CPU 拷贝和 CPU 内存占用，提升显存带宽利用率。
  • 团队：代码量精简 27 行（+14/-13），移除对 diffusers.video_processor 的依赖，降低了外部库耦合。
  • 影响范围：仅影响 Cosmos3 模型，其他扩散模型不受影响。
  • 风险标记：缺少测试覆盖, 输出类型变更（numpy→tensor）

关联脉络

• PR #26947 [diffusion] Speed up PNG image output saving: 同属扩散引擎性能优化系列，加速 PNG 输出保存，与本次后处理 GPU 化相辅相成。