Prhub

#34564 [diffusion] Stream and parallelize bit-exact video output saves

原始 PR 作者 mickqian 合并时间 2026-08-12 20:03 文件变更 3 提交数 1 评论 0 代码增减 +295 / -85

执行摘要

流式分块转换并并行保存多路视频输出,显存降约 95%、耗时减半

PR body 明确说明痛点:旧路径在编码前一次性物化整段视频的浮点乘法结果和两套 uint8 布局,临时显存与 pinned-host 内存随视频时长线性增长;num_outputs_per_prompt > 1 时每路独立结果只能串行编码。实测数据驱动了本次改造:1344×768 下五秒转换临时显存峰值约 1.92 GB,十五秒 RGB24 主机缓冲约 1.13 GB,双输出保存约 1.30 s;改造后分别降到约 108 MB、21.7 MB 与 0.69 s。

值得精读,尤其是 _try_save_cuda_videos_direct 的显存/CPU 预算启发式和 _sendfile_all 的部分写处理,是典型的「零拷贝流式 + 资源感知并行」设计。建议关注两点:共享 memfd 缓冲在并行线程间的互斥窗口是否覆盖整个使用周期,以及非 Linux 平台回退路径的 CI 覆盖。若后续推广到更多 diffusion 模型,可考虑把 chunk 预算做成可配置项。

讨论亮点

本 PR 没有公开的 review 评论(review 与 issue 评论均为 0),评审证据主要来自 PR body 的自测:5 秒/15 秒合成视频的 MP4 字节级比对、4×H200 MiniMax-H3 Ref2VA 双输出的 MP4 容器/解码 RGB/解码 PCM 一致性,以及注入 sendfile 失败后的回退与临时文件清理验证。虽然没有评审交锋,但实现中「先算预算再决定是否并行」的做法本身就体现了对 OOM 和 CPU 争用的防御性设计。

实现拆解

  1. 分块转换预算:在 python/sglang/multimodal_gen/runtime/entrypoints/utils.py 中新增常量 _MAX_CUDA_VIDEO_CONVERSION_CHUNK_BYTES = 128 MiB 和辅助函数 _cuda_video_conversion_chunk_frames。它按「每帧临时字节 = 3 × H × W × (元素字节数 + 2)」反推每次能安全处理的帧数,从源头控制 CUDA 转换的临时显存峰值,不再让 frames = (video * 255).clamp_(0, 255).to(torch.uint8) 一次性物化整段视频。
  2. 流式喂给 ffmpeg_try_save_cuda_video_direct 的 ffmpeg 输入从 /proc/self/fd/{fd} 改为 pipe:0,进程由 subprocess.run 改为 Popen;新增 _sendfile_all 把 memfd 中的原始帧分块 os.sendfile 到 ffmpeg stdin,并循环处理部分写入。每块转换、拷入、同步 CUDA 流后立即送走,memfd 缓冲只保留一个 chunk。异常路径会 kill 子进程并读取 stderr 转储,随后统一回退 imageio。
  3. 并行多路保存:新增 _try_save_cuda_videos_direct,前置校验包括 CUDA 张量、3 通道、.mp4 扩展名、多路同设备;显存侧按最大的两个 chunk 临时字节是否超过当前空闲显存 25% 决定是否并行,CPU 侧按两个 x264 自动线程数之和是否超过 os.sched_getaffinity 可用核数决定。满足条件时用 ThreadPoolExecutor(max_workers=2) 并发执行 save_one,否则返回 None 交由串行路径。
  4. 回退与接线save_outputs 对并行结果中 False 的输出逐一路回退到既有 _try_save_cuda_video_direct 串行路径(再失败才走 post_process_sample),测试用 monkeypatch 验证了 [True, False] 场景下失败项确实走串行且不触发帧物化。
  5. MiniMax-H3 校验并行化.../minimax_h3/video_adapter.pyvalidate_final_outputs_sync 把逐路 ffprobe 改为 ThreadPoolExecutor(max_workers=min(4, len(output_paths))) + pool.map,保持输出索引顺序,不一致时仍按原样报错,媒体元数据探测耗时从约 73–82 ms 降到 38–40 ms。
  6. 测试配套python/sglang/multimodal_gen/test/unit/test_output_saving.py 新增 test_multiple_videos_use_parallel_direct_save_with_serial_fallback,覆盖并行入口调用、路径与 fps 参数传递、失败回退;PR 同时做了 5 秒/15 秒合成视频字节级比对和 4×H200 MiniMax-H3 双输出 MP4/解码 RGB/解码 PCM 一致性验证。
文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/entrypoints/utils.py 输出保存 modified 8.65
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/video_adapter.py 媒体校验 modified 6.33
python/sglang/multimodal_gen/test/unit/test_output_saving.py 单元测试 modified 5.39

关键符号

_cuda_video_conversion_chunk_frames _sendfile_all _try_save_cuda_video_direct _try_save_cuda_videos_direct save_one probe_output validate_final_outputs_sync

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

评论区精华

评审以自测与 CI 为主,无公开评论 other

该 PR 的 review 评论数为 0,没有公开的评审交锋;关键证据来自 PR body 中 5 秒与 15 秒合成视频字节级比对、4×H200 MiniMax-H3 Ref2VA 双输出验证以及注入 sendfile 失败的回归用例。

结论:作者用字节级一致性与回退注入测试证明流式路径与旧路径等价,CI 全绿后合入。 · 已解决

风险与影响

  1. 平台与内核依赖:新路径强依赖 os.memfd_createos.sendfile,两者仅在 Linux 可用;_try_save_cuda_video_direct 在缺失时返回 False 走 imageio 回退,非 Linux 不受影响,但回退路径的实际覆盖需要 CI 验证。
  2. OOM 保护是启发式:并行判定只看转换 chunk 的临时字节是否低于空闲显存 25%,没有把 _CudaMemfdVideoBuffer 本身、WAV 文件和 ffmpeg 内部缓冲计入;极端低显存场景下仍可能 OOM。分块预算 128 MiB 也只是按元素大小估算,float32bfloat16 的系数差异会影响实际峰值。
  3. 并行编码的资源竞态:x264 的 -threads_x264_auto_thread_count(height) 自动推导,与推理进程共享 CPU;检查基于 os.sched_getaffinity 的当前可用核数,但同机其他租户的瞬时负载无法感知,可能造成编码与推理互相拖慢。
  4. 共享 memfd 缓冲的并发访问:全局 _cached_cuda_video_buffer_cuda_video_buffer_cache_lock 保护,但若 _acquire_cuda_video_buffer 在返回后即释放锁,两个并行线程可能拿到同一块缓冲并相互覆盖。建议确认锁覆盖整个 copy/sendfile 窗口,或在并行路径为每路分配独立缓冲(这一点从当前源码窗口无法完全确认)。
  5. ffmpeg 管道早退Popen 后若 ffmpeg 因参数或编码错误提前退出,继续向 stdin 写可能触发 BrokenPipeError;代码在 finally 中 kill 与 wait,并用 stderr 临时文件收集原因,但这类错误在真实编码失败时仍可能暴露在回退日志中。

对用户:长视频生成的显存足迹显著下降(1344×768 五秒转换临时显存约 1.92 GB → 108 MB,十五秒 RGB24 主机缓冲约 1.13 GB → 21.7 MB),num_outputs_per_prompt > 1 的保存墙钟时间约减半(1.30 s → 0.69 s),且输出与旧路径字节级一致,属于透明优化。
对系统:并行编码会短暂双倍占用 x264 线程与 CPU 核,已用 affinity 上限约束;多路输出同时 ffprobe 也只存在于 MiniMax-H3 校验阶段。
对团队:直接保存路径从同步 subprocess.run 切换为 Popen + 管道,复杂度上升,后续需要保持回退路径的测试覆盖;共享 memfd 缓冲的线程安全是一个需要跟踪的隐患。

平台依赖 memfd/sendfile 并行编码 CPU 竞态 OOM 保护为启发式 共享 memfd 缓冲并发访问待确认

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论