执行摘要

• 一句话：融合因果 Conv3d 的 cat 与 pad 为 Triton 内核，加速 VAE 解码
• 推荐动作：建议团队阅读本 PR，尤其是如何通过 Triton 融合简单操作来消除小开销，以及设计回退和守护条件的方式。对于未来类似性能优化（如卷积前置处理），可复用此模式。另外，建议为 Triton 内核添加单元测试，提高回归覆盖。

功能与动机

PR body 指出视频/图像 VAE 的因果 Conv3d 在用到缓存上下文时，会先做 temporal torch.cat 再做 spatial F.pad，这是 decode-heavy WAN 路径中重复出现的小开销。PR 旨在融合这两个操作，消除额外内核启动和内存读写。

实现拆解

实现分为三步：

1. 新增 Triton 内核：创建 python/sglang/jit_kernel/diffusion/triton/causal_conv3d_pad.py，定义 _fused_cat_pad_5d_kernel（Triton JIT 内核）和封装函数 fused_causal_conv3d_cat_pad。内核根据输出坐标反向推导源坐标，通过 tl.where 在输入 x 和 cache_x 间选择，并应用边界填充。所有 shape/padding 参数作为运行时标量传入，避免重编译。
2. 通用融合入口：在 python/sglang/multimodal_gen/runtime/layers/parallel_conv.py 中添加 _can_fuse_causal_conv3d_cat_pad 检查条件（CUDA、contiguous、形状兼容、padding 对称等），以及 causal_conv3d_cat_pad 函数，其优先使用 Triton 融合路径，否则回退到原有 PyTorch 路径。仅在 CUDA 平台导入 Triton 内核。
3. 替换使用者：修改 wanvae.py 中的 WanCausalConv3d.forward、autoencoder_kl_qwenimage.py 中的 QwenImageCausalConv3d.forward，以及 parallel_conv.py 中的 SpatialParallelCausalConv3d.forward，将原有的 cat+pad 调用替换为 causal_conv3d_cat_pad(x, cache_x, padding)。

未添加专门的单元测试，但 PR 验证了微基准一致性（max diff=0）和 WAN/Qwen 端到端一致性。

关键文件：

• python/sglang/jit_kernel/diffusion/triton/causal_conv3d_pad.py（模块 Triton 内核；类别 source；类型 core-logic；符号 _fused_cat_pad_5d_kernel, fused_causal_conv3d_cat_pad）: 新增 Triton JIT 内核和封装函数，是本次优化的核心。
• python/sglang/multimodal_gen/runtime/layers/parallel_conv.py（模块 并行卷积；类别 source；类型 core-logic；符号 _can_fuse_causal_conv3d_cat_pad, causal_conv3d_cat_pad）: 添加融合入口、守护条件，并修改 SpatialParallelCausalConv3d 的 forward。
• python/sglang/multimodal_gen/runtime/models/vaes/autoencoder_kl_qwenimage.py（模块 VAE 模型；类别 source；类型 data-contract）: QwenImageCausalConv3d 的 forward 改用统一融合入口，减少重复代码。
• python/sglang/multimodal_gen/runtime/models/vaes/wanvae.py（模块 VAE 模型；类别 source；类型 data-contract）: WanCausalConv3d 的 forward 改用统一融合入口，减少重复代码。

关键符号：_fused_cat_pad_5d_kernel, fused_causal_conv3d_cat_pad, _can_fuse_causal_conv3d_cat_pad, causal_conv3d_cat_pad

关键源码片段

python/sglang/jit_kernel/diffusion/triton/causal_conv3d_pad.py

新增 Triton JIT 内核和封装函数，是本次优化的核心。

@triton.jit
def _fused_cat_pad_5d_kernel(
    x_ptr, cache_ptr, out_ptr,
    total, channels, t_size, h_size, w_size,
    cache_t, out_t, out_h, out_w,
    pad_d_left, pad_h_top, pad_w_left,
    block_size: tl.constexpr,
):
    # 一维偏移转五维坐标
    offsets = tl.program_id(0) * block_size + tl.arange(0, block_size)
    mask = offsets < total
    ow = offsets % out_w
    tmp = offsets // out_w
    oh = tmp % out_h
    tmp = tmp // out_h
    out = tmp % out_t
    tmp = tmp // out_t
    oc = tmp % channels
    ob = tmp // channels
    # 反向推导源坐标（减去 padding）
    iw = ow - pad_w_left
    ih = oh - pad_h_top
    src_t = out - pad_d_left
    # 有效性检查：必须在有效范围内
    valid = mask & (iw >= 0) & (iw < w_size) & (ih >= 0) & (ih < h_size) & (src_t >= 0) & (src_t < cache_t + t_size)
    from_cache = src_t < cache_t
    # 夹紧坐标以避免越界
    clamped_iw = tl.minimum(tl.maximum(iw, 0), w_size - 1)
    clamped_ih = tl.minimum(tl.maximum(ih, 0), h_size - 1)
    clamped_x_t = tl.minimum(tl.maximum(src_t - cache_t, 0), t_size - 1)
    clamped_src_t = tl.minimum(tl.maximum(src_t, 0), cache_t - 1)
    # 计算偏移
    x_offsets = (((ob * channels + oc) * t_size + clamped_x_t) * h_size + clamped_ih) * w_size + clamped_iw
    cache_offsets = (((ob * channels + oc) * cache_t + clamped_src_t) * h_size + clamped_ih) * w_size + clamped_iw
    # 加载数据并合并
    x_vals = tl.load(x_ptr + x_offsets, mask=valid & ~from_cache, other=0.0)
    cache_vals = tl.load(cache_ptr + cache_offsets, mask=valid & from_cache, other=0.0)
    tl.store(out_ptr + offsets, tl.where(from_cache, cache_vals, x_vals), mask=mask)
​
def fused_causal_conv3d_cat_pad(x, cache_x, padding):
    # 仅支持对称 padding 和 depth_left >= cache_t
    depth_left = padding[4] - cache_x.shape[2]
    assert depth_left >= 0 and padding[5] == 0
    assert padding[0] == padding[1] and padding[2] == padding[3]
    bsz, channels, t_size, h_size, w_size = x.shape
    cache_t = cache_x.shape[2]
    out = torch.empty(bsz, channels, t_size + cache_t + depth_left,
                       h_size + padding[2]*2, w_size + padding[0]*2,
                       device=x.device, dtype=x.dtype)
    grid = (triton.cdiv(out.numel(), 256),)
    _fused_cat_pad_5d_kernel[grid](x, cache_x, out, out.numel(), channels,
        t_size, h_size, w_size, cache_t, out.shape[2], out.shape[3], out.shape[4],
        depth_left, padding[2], padding[0], 256)
    return out

python/sglang/multimodal_gen/runtime/layers/parallel_conv.py

添加融合入口、守护条件，并修改 SpatialParallelCausalConv3d 的 forward。

def _can_fuse_causal_conv3d_cat_pad(x, cache_x, padding):
    # 检查是否满足融合条件：CUDA 设备、contiguous、5D、形状兼容、对称 padding 等
    if cache_x is None or fused_causal_conv3d_cat_pad is None:
        return False
    if not x.is_cuda or not x.is_contiguous() or not cache_x.is_contiguous():
        return False
    if x.dim() != 5 or cache_x.dim() != 5 or x.dtype != cache_x.dtype:
        return False
    if x.shape[0] != cache_x.shape[0] or x.shape[1] != cache_x.shape[1]:
        return False
    if x.shape[3:] != cache_x.shape[3:]:
        return False
    width_left, width_right, height_top, height_bottom, depth_left, depth_right = padding
    if width_left != width_right or height_top != height_bottom or depth_right != 0:
        return False
    if depth_left < cache_x.shape[2]:
        return False
    return bool(width_left or height_top)
​
def causal_conv3d_cat_pad(x, cache_x, padding):
    # 统一入口：优先使用 Triton 融合，否则回退到 cat+F.pad
    if cache_x is not None and padding[4] > 0:
        if cache_x.device != x.device:
            cache_x = cache_x.to(x.device)
        if _can_fuse_causal_conv3d_cat_pad(x, cache_x, padding):
            return fused_causal_conv3d_cat_pad(x, cache_x, padding)
        x = torch.cat([cache_x, x], dim=2)
        padding[4] -= cache_x.shape[2]
    if any(padding):
        x = F.pad(x, padding)
    return x

评论区精华

来自 gemini-code-assist[bot] 的评审要点：

• 避免 tl.constexpr 动态参数：使用 tl.constexpr 会导致 Triton 为每种 shape 组合重编译，造成延迟尖峰。建议改为运行时标量参数，并夹紧坐标防止越界。
  • 作者回应：已在提交 15dd1fe 中修复，使用运行时参数，并通过 tl.minimum(tl.maximum(...)) 夹紧坐标（因环境 Triton 版本不支持 tl.clamp）。

• 导入用 try-except 包裹：建议用 try-except ImportError 捕获 Triton 缺失，避免导入时崩溃。

  • 作者拒绝：CUDA 扩散执行依赖 Triton，吞掉 ImportError 会隐藏无效安装，使回退意外触发；Triton 缺失应视为环境错误。

  • tl.constexpr 动态参数导致重编译 (performance): 作者在提交 15dd1fe 中已修复，改为运行时参数，并使用 tl.minimum/tl.maximum 夹紧坐标。

  • Triton 导入是否应使用 try-except (design): 作者拒绝：CUDA 扩散执行依赖 Triton，吞掉 ImportError 会隐藏无效安装，使回退意外触发；Triton 缺失应视为环境错误。

风险与影响

• 风险：
  1. 新 Triton 内核正确性：尽管微基准验证了 max diff=0，但缺乏自动化单元测试；若未来 CUDA 或 Triton 版本变动，可能引入数值差异。
  2. 平台限制：融合路径仅支持 CUDA；非 CUDA 设备（如 ROCm、NPU、CPU）将始终走回退 PyTorch 路径，性能不受益但正确性不受影响。
  3. 回退路径风险：回退路径仅在被融合检查拒绝时触发；若因导入失败导致 fused_causal_conv3d_cat_pad 为 None（当前不会，因为严格导入），则会回退；但若内核在运行时崩溃（非法内存访问），将直接抛出异常而非回退，因为未捕获内核异常。
  4. 假设条件：内核假设 padding 对称（width_left==width_right, height_top==height_bottom, depth_right==0）且 depth_left >= cache_t，若不满足则断言失败或静默错误。当前所有调用点满足这些条件。
  5. 性能回归：对于极小 shape（如 T=1, C=128），融合内核微基准有时略慢（如 0.0293ms->0.0292ms），但差别很小，不影响整体。
     - 影响：对用户：WAN 和 Qwen Image VAE 解码端到端性能提升约 5%（WAN）及解码阶段加速（decode ms 从 483.96 降至 466.25），Qwen Image 解码时间从 44.97ms 降至 40.24ms。数值严格一致，用户无需修改代码。系统上，新增 Triton 内核缓存，首次运行会有编译开销（已通过运行时参数避免重编译）。团队影响低，改动集中在 4 个文件，回退机制保证安全。
     - 风险标记：无单元测试覆盖, 仅 CUDA 平台, 内核崩溃无回退, 静默断言风险

关联脉络

• 暂无明显关联 PR