用异步帧生成流水线把 GPU 利用率从 82% 拉到 99.9%——视频推理加速实战

import torch
import torch.nn as nn
from torch.cuda import CUDAGraph
import time

# ---------- 模拟 VAE 解码器 ----------
class FakeVAEDecoder(nn.Module):
    """模拟一个耗时 ~50ms 的 VAE 解码 kernel"""
    def __init__(self, latent_channels=16, out_channels=3):
        super().__init__()
        self.conv = nn.Conv2d(latent_channels, out_channels, kernel_size=3, padding=1)

    def forward(self, x):
        # 用多次前向模拟真实 VAE 的计算量
        for _ in range(20):
            x = self.conv(x)
        return x

# ---------- 异步帧生成流水线 ----------
class AsyncFramePipeline:
    def __init__(self, vae: nn.Module, num_frames: int, device="cuda"):
        self.vae = vae.to(device).eval()
        self.device = device
        self.num_frames = num_frames

        # 专用 CUDA 流：D2H 拷贝用
        self.d2h_stream = torch.cuda.Stream(device=device)

        # pinned memory 主机缓冲区（D2H 目的地）
        # 假设解码后帧尺寸 3×512×512
        self.host_buffers = [
            torch.empty(3, 512, 512, dtype=torch.float32).pin_memory()
            for _ in range(3)  # 三缓冲：GPU写、D2H搬、CPU读 各占一个
        ]

    def _gpu_decode(self, latent: torch.Tensor) -> torch.Tensor:
        """在默认流上做 VAE 解码"""
        with torch.no_grad():
            return self.vae(latent)

    def _d2h_copy(self, gpu_frame: torch.Tensor, host_buf: torch.Tensor):
        """在专用流上异步 D2H 拷贝"""
        with torch.cuda.stream(self.d2h_stream):
            host_buf.copy_(gpu_frame, non_blocking=True)

    def _host_postprocess(self, host_buf: torch.Tensor) -> torch.Tensor:
        """CPU 端后处理：颜色空间转换、clamp、to uint8 等"""
        frame = host_buf.cpu().clamp(0, 1)
        return (frame * 255).to(torch.uint8)

    def run(self, latents: list[torch.Tensor]):
        """
        latents: 长度 = num_frames, 每个是 (C, H, W) 的 latent tensor
        返回: list of uint8 帧张量
        """
        results = []
        gpu_frames = [None, None, None]  # 三缓冲 GPU 端

        for i in range(self.num_frames):
            # --- 阶段 1: GPU 解码当前帧 (帧 i) ---
            gpu_frames[i % 3] = self._gpu_decode(latents[i].to(self.device))

            # --- 阶段 2: D2H 拷贝上一帧 (帧 i-1, 如果存在) ---
            if i >= 1:
                prev_gpu = gpu_frames[(i - 1) % 3]
                self._d2h_copy(prev_gpu, self.host_buffers[(i - 1) % 3])

            # --- 阶段 3: CPU 处理更早帧 (帧 i-2, 如果存在) ---
            if i >= 2:
                # 等 D2H 流完成，确保数据已到主机
                self.d2h_stream.synchronize()
                processed = self._host_postprocess(self.host_buffers[(i - 2) % 3])
                results.append(processed)

        # 处理流水线尾部残留帧
        self.d2h_stream.synchronize()
        if self.num_frames >= 1:
            # 倒数第二帧的 D2H 已在最后一次循环发起
            if self.num_frames >= 2:
                results.append(self._host_postprocess(self.host_buffers[(self.num_frames - 1) % 3]))
            # 最后一帧：手动 D2H + 处理
            self._d2h_copy(gpu_frames[(self.num_frames - 1) % 3], self.host_buffers[(self.num_frames - 1) % 3])
            self.d2h_stream.synchronize()
            results.append(self._host_postprocess(self.host_buffers[(self.num_frames - 1) % 3]))

        return results

# ---------- 运行 ----------
if __name__ == "__main__":
    torch.cuda.set_device(0)
    vae = FakeVAEDecoder(latent_channels=16, out_channels=3)

    num_frames = 8
    latents = [torch.randn(16, 64, 64) for _ in range(num_frames)]

    pipeline = AsyncFramePipeline(vae, num_frames)

    # Warmup
    _ = pipeline.run(latents[:2])

    torch.cuda.synchronize()
    t0 = time.time()
    frames = pipeline.run(latents)
    torch.cuda.synchronize()
    t1 = time.time()

    print(f"异步流水线: {num_frames} 帧, 总耗时 {t1 - t0:.3f}s")
    print(f"输出帧 shape: {frames[0].shape}, dtype: {frames[0].dtype}")