Prhub

#34085 [diffusion] Clean up kernels and shared fast paths

原始 PR 作者 BBuf 合并时间 2026-08-09 00:37 文件变更 50 提交数 5 评论 1 代码增减 +1349 / -1367

diffusion jit-kernel refactor test performance

执行摘要

集中 diffusion 共享快路径，统一内核与质量门控

PR body 明确指出：近期扩散性能工作在并行推进中留下大量重复，"That left duplicated launch plumbing, repeated numerical helpers, model-local fallback logic, and a few overly narrow or dead components"，并且"those copies could drift in supported shapes, fallback behavior, or precision contracts"。因此本 PR 的目标是让每个共享关注点只保留一份实现，并使后端选择显式化；lossless 路径继续要求 reference 等价或 bit-exact，非 bit-exact 融合仍限定在 quality=high。

值得精读。本 PR 展示了扩散内核在性能优化之后的收敛模式：共享数值原语保证 bit-exact 口径、QualityGatedFusion 以 all-or-nothing 方式管理非 bit-exact 融合、模型侧不再各写 fallback。建议重点阅读 quality_gate.py、native_bf16_rmsnorm.py 和 denoising.py 的表驱动挂载逻辑；如果团队要新增扩散模型或新融合内核，这套"模型无关命名 + 位精确契约 + 静态守卫"的组织方式可以直接借鉴。

讨论亮点

本 PR 没有来自其他维护者的 review 评论，唯一的评论是作者 BBuf 贴出的 CI 链接。讨论内容以 PR body 中的验证数据为主：H200 上 parent-vs-PR A/B 共 19 个内核用例 × 5 次运行，每次输出 hash 都与父提交一致，中位延迟差异在 -0.90% 到 +0.43% 之间；post-rebase 最终 SHA 上 2700 个相关测试通过、1 个跳过；生产形状 A/B 中 LN+modulate、residual-gate、Ulysses relayout、adaLN modulation 均保持 bit-exact。

实现拆解

  1. 共享数值原语下沉：新增 python/sglang/kernels/ops/diffusion/triton/numerics.py，集中 round_bf16_to_fp32、mul_rn_f32（inline asm 阻断 FMA 收缩）、div_rn_f32、rsqrt_approx_f32、cuda_rsqrtf，让 rmsnorm_scale_shift_bitexact.py、layernorm_modulate.py 等复用，消除各内核之间数值口径漂移的可能。
  2. BF16 原生 RMSNorm 通用化：把 Z-Image 模块里重复的 _rmsnorm_scale_kernel、_rmsnorm_tanh_residual_kernel 及启动逻辑抽到新增 native_bf16_rmsnorm.py，公开 rmsnorm_scale、rmsnorm_tanh_residual；同时补强守卫：三者设备一致、全 bf16、shape[-1] <= 8192、非空、weight 形状与连续、行 stride 扁平连续。zimage_native_norm.py 保留 QK RMSNorm 专用实现并同步收紧 dtype/device/numel 守卫。
  3. 模型侧快路径入口统一：FLUX、FLUX.2、GLM-Image、ERNIE-Image、LTX-2 各自维护的 _*_residual_gate_add（以及 FLUX 的 _flux_modulate）连同 _DISABLED 全局标志全部删除，改为统一调用 residual_gate_add / modulate_scale_shift，把 dtype 守卫、异常回退、torch.compile 下抛错等策略收口到内核层一处；fused_linear_gelu 的开关判断也从直接读写 _sgl_fused_gelu_enabled 改为 fused_gelu_active()。
  4. 质量门控集中管理：新增 quality_gate.py 的 QualityGatedFusion 类（mark/metadata/is_enabled/iter_sites/mount/unmount），denoising 阶段用 _QUALITY_FUSION_HANDLERS 表驱动 _maybe_toggle_quality_fusions，对 linear+GELU、LN+modulate、gate-RMSNorm 三个家族按 batch 的 quality 做 all-or-nothing 挂载/卸载。
  5. 测试与维护配套：test_zimage_native_norm.py 改名为 test_native_bf16_rmsnorm.py 并扩展拒绝用例；新增 test_quality_gate.py、test_ulysses_qkv.py、test_scale_shift.py；删除针对已移除包装的 test_ernie_residual_gate_add.py；对 fused LN+modulate 增加 torch.compile(fullgraph=True) 覆盖。5 个提交还包括对齐 modulation JIT kernel namespace、修复 quality-gate 测试入口、清理死融合助手等收尾工作。
文件 模块 状态 重要度
python/sglang/multimodal_gen/runtime/models/dits/flux.py 扩散模型 modified 8.17
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py 扩散模型 modified 7.5
python/sglang/kernels/ops/diffusion/triton/native_bf16_rmsnorm.py 内核层 added 7.69
python/sglang/kernels/ops/diffusion/triton/zimage_native_norm.py 内核层 modified 7.26
python/sglang/kernels/ops/diffusion/quality_gate.py 门控协议 added 7.22
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py 去噪管线 modified 7.15
python/sglang/kernels/ops/diffusion/triton/numerics.py 数值原语 added 7.09
test/registered/kernels/ops/diffusion/test_native_bf16_rmsnorm.py 测试 renamed 7.01

关键符号

rmsnorm_scale rmsnorm_tanh_residual QualityGatedFusion.mount QualityGatedFusion.unmount QualityGatedFusion.is_enabled residual_gate_add modulate_scale_shift fused_gelu_active _maybe_toggle_quality_fusions round_bf16_to_fp32 mul_rn_f32 zimage_qk_rmsnorm_native

关键源码片段

python/sglang/kernels/ops/diffusion/triton/native_bf16_rmsnorm.py infrastructure

新增的通用 BF16 原生 RMSNorm 融合模块，将 Z-Image 专用的两个 Triton 内核与守卫逻辑抽取为模型无关实现，是本次重构的基础设施核心。

# python/sglang/kernels/ops/diffusion/triton/native_bf16_rmsnorm.py
# 从 Z-Image 专用模块抽出的通用 BF16 原生 RMSNorm 融合，供多个扩散模型共享。
import torch
import triton
import triton.language as tl
​
MAX_HIDDEN_SIZE = 8192
​
​
@triton.jit
def _rmsnorm_scale_kernel(
    y_ptr, x_ptr, weight_ptr, scale_ptr,
    x_row_stride, scale_row_stride, seq_len,
    dim: tl.constexpr, eps: tl.constexpr, block_dim: tl.constexpr,
):
    # 每个 program 处理一行；rstd 的求值顺序与 aten bf16 链保持一致，
    # 每次乘 / 规约都落回 bf16，从而与 eager 的两次舍入逐位一致。
    row = tl.program_id(0)
    offsets = tl.arange(0, block_dim)
    mask = offsets < dim
​
    x = tl.load(x_ptr + row * x_row_stride + offsets, mask=mask, other=0.0)
    square = (x * x).to(tl.bfloat16)
    mean_square = (tl.sum(square, axis=0) / dim).to(tl.bfloat16)
    rstd = tl.rsqrt((mean_square + eps).to(tl.bfloat16).to(tl.float32)).to(tl.bfloat16)
​
    batch = row // seq_len
    weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0)
    scale = tl.load(scale_ptr + batch * scale_row_stride + offsets, mask=mask, other=0.0)
    y = (((x * rstd).to(tl.bfloat16) * weight).to(tl.bfloat16) * scale).to(tl.bfloat16)
    tl.store(y_ptr + row * dim + offsets, y, mask=mask)
​
​
def _flat_row_stride(x: torch.Tensor) -> int | None:
    # 行 stride 必须严格形成“扁平行”布局，否则拒绝快路径，
    # 避免把非连续视图按连续行处理导致错位。
    if x.dim() < 2 or x.stride(-1) != 1:
        return None
    row_stride = x.stride(-2)
    expected_stride = row_stride * x.shape[-2]
    for dim in range(x.dim() - 3, -1, -1):
        if x.stride(dim) != expected_stride:
            return None
        expected_stride *= x.shape[dim]
    return row_stride
​
​
def rmsnorm_scale(
    x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor, eps: float,
) -> torch.Tensor | None:
    # 统一守卫：同一 CUDA 设备、全部 bf16、隐藏维不超过 8192、
    # 非空输入、weight 形状匹配且连续；任一不满足都返回 None 让调用方回退 eager。
    if not _can_use_operand(x, weight, scale):
        return None
​
    dim = x.shape[-1]
    x_rows = x.numel() // dim
    scale_rows = scale.numel() // dim
    if x_rows % scale_rows != 0:
        return None
​
    x_row_stride = _flat_row_stride(x)
    scale_row_stride = _flat_row_stride(scale)
    if x_row_stride is None or scale_row_stride is None:
        return None
​
    out = torch.empty_like(x, memory_format=torch.contiguous_format)
    with torch.get_device_module().device(x.device):
        _rmsnorm_scale_kernel[(x_rows,)](
            out.reshape(-1, dim),
            x,
            weight,
            scale,
            x_row_stride,
            scale_row_stride,
            x_rows // scale_rows,
            dim,
            eps,
            block_dim=triton.next_power_of_2(dim),
            num_warps=8,
        )
    return out
python/sglang/kernels/ops/diffusion/quality_gate.py infrastructure

新增 `QualityGatedFusion` 协议类，统一定义 fusion site 的标记、启用状态与 all-or-nothing 挂载语义，是整个 quality=high 门控机制的核心抽象。

# python/sglang/kernels/ops/diffusion/quality_gate.py
class QualityGatedFusion:
    """单个融合家族（family）的挂载协议。
​
    marker_attr 标记 site，enabled_attr 是普通模块属性，forward 在
    torch.compile 下可直接读取它，而不依赖本 Python 对象。
    """
​
    __slots__ = ("enabled_attr", "marker_attr", "name")
​
    def __init__(self, *, name: str, marker_attr: str, enabled_attr: str) -> None:
        self.name = name
        self.marker_attr = marker_attr
        self.enabled_attr = enabled_attr
​
    def mark(self, module, metadata=True) -> None:
        # 标记 site 并默认关闭，避免未挂载时误走非 bit-exact 路径。
        setattr(module, self.marker_attr, metadata)
        setattr(module, self.enabled_attr, False)
​
    def is_enabled(self, module) -> bool:
        return bool(getattr(module, self.enabled_attr, False))
​
    def iter_sites(self, root):
        # marker_attr 是 site 的“身份证”，遍历模块树即可收集整个家族。
        for module in root.modules():
            if hasattr(module, self.marker_attr):
                yield module
​
    def mount(self, root, *, reject_reason=None, logger=None) -> bool:
        # All-or-nothing：任一 site 不满足静态守卫就整个家族回退 reference，
        # 避免同一家族部分开启导致数值口径不一致。
        sites = list(self.iter_sites(root))
        if not sites:
            return False
​
        if reject_reason is not None:
            for site in sites:
                reason = reject_reason(site)
                if reason is None:
                    continue
                self._set_enabled(sites, False)
                if logger is not None:
                    logger.info(
                        "%s: %s site failed static guards (%s); keeping the "
                        "whole model on the reference path",
                        self.name,
                        type(site).__name__,
                        reason,
                    )
                return False
​
        self._set_enabled(sites, True)
        return True
​
    def unmount(self, root) -> None:
        self._set_enabled(self.iter_sites(root), False)
​
    def _set_enabled(self, sites, enabled: bool) -> None:
        for site in sites:
            setattr(site, self.enabled_attr, enabled)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py core-logic

将质量门控挂载逻辑重构为 `_QUALITY_FUSION_HANDLERS` 表驱动，三个融合家族统一遍历，日志从手写 if/else 收敛为集合驱动输出。

# python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
# 三个 quality=high 融合家族统一注册为处理器表，新增家族只需追加一项。
_QUALITY_FUSION_HANDLERS: tuple[
    tuple[str, Callable[[nn.Module], bool], Callable[[nn.Module], None]], ...
] = (
    (
        "fused linear+GELU (cublasLt epilogue)",
        mount_fused_linear_gelu,
        unmount_fused_linear_gelu,
    ),
    (
        "fused LN+modulate (affine folding)",
        mount_fused_ln_modulate,
        unmount_fused_ln_modulate,
    ),
    (
        "fused gate RMSNorm (BF16-native Triton)",
        mount_fused_gate_rmsnorm,
        unmount_fused_gate_rmsnorm,
    ),
)
​
​
def _maybe_toggle_quality_fusions(self, batch: Req) -> None:
    # quality="high" 挂载全部允许的融合，否则全部卸载；
    # quality 参与动态 batch 签名，因此 batch 内统一，进程级切换是安全的。
    want = getattr(batch.sampling_params, "quality", "lossless") == "high"
    if want == self._quality_fusions_mounted:
        return
    mounted_fusions: set[str] = set()
    for transformer in filter(None, [self.transformer, self.transformer_2]):
        for description, mount, unmount in _QUALITY_FUSION_HANDLERS:
            if want:
                if mount(transformer):
                    mounted_fusions.add(description)
            else:
                unmount(transformer)
    self._quality_fusions_mounted = want
    for description in sorted(mounted_fusions):
        logger.info("Mounted %s for quality=high", description)

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论，后续有更多讨论时会体现在这里。

风险与影响

  1. 行为对齐差异：统一入口后，原先各模型包装的守卫可能不完全一致。例如 LTX-2 旧版 _ltx2_residual_gate_add 没有 half-dtype 限制，而 FLUX/GLM/ERNIE 旧版都有；若共享 residual_gate_add 以其中一方为准，对另一方是静默行为变化（性能或数值路径）。需确认共享实现覆盖了所有原守卫。
  2. 平台兼容性：numerics.py 中的正确舍入依赖 PTX inline asm（mul.rn.f32、div.rn.f32），在非 NVIDIA 平台（AMD、NPU、Apple Silicon）上不可编译或不可用，依赖调用方的平台/设备守卫；H200 上的验证无法覆盖这些平台。
  3. torch.compile 稳定性：QualityGatedFusion 把 enabled 状态放在普通模块属性上以兼容 compile，但如果 site 在编译图捕获后被动态修改，可能造成图不一致；all-or-nothing 的挂载逻辑依赖 batch 内 quality 均匀，改动动态 batch 签名时需要同步审计。
  4. 验证范围：H200 上测试充分，但错误路径、非连续视图、极端 shape 组合只靠单元测试覆盖，跨模型重构下仍有回归风险。

影响面覆盖 diffusion 侧全部主要 DiT 模型家族（FLUX、FLUX.2、GLM-Image、ERNIE-Image、LTX-2、Z-Image）的公共快路径：residual-gate、modulate、BF16 RMSNorm、质量门控挂载。对用户无 API 变化，quality=lossless 保持 bit-for-bit，quality=high 的融合行为由同一套协议统一控制。对团队而言，内核启动约定与 JIT kernel 惯例对齐、模块改为模型无关命名，后续新增扩散模型可直接复用；测试体系从模型本地用例收敛为共享内核用例，降低了维护成本。

跨模型共享路径重构 位精确性依赖硬件验证 inline PTX asm 平台兼容性 LTX-2 旧路径无 half 守卫 测试聚焦 H200

关联 Issue

未识别关联 Issue

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

完整报告

Markdown PDF

参与讨论