triton_fused_norm ¶

def fused_path_supported(
    residual: torch.Tensor,
    x: torch.Tensor,
    gate: torch.Tensor | int,
    shift: torch.Tensor | None,
    scale: torch.Tensor | None,
    norm: torch.nn.Module,
) -> bool:
    """Cheap eligibility check; any False falls back to the eager path."""
    if not _HAS_TRITON or envs.FASTVIDEO_DISABLE_FUSED_NORM.get():
        return False
    if torch.is_grad_enabled():
        return False
    if torch.compiler.is_compiling():
        # Dynamo must not trace into the Triton launcher (see the custom-op
        # boundaries in fastvideo/models/dits/minimax_h3_fusions/modulation.py).
        # Until this kernel is wrapped in torch.library.custom_op + register_fake
        # the way that module does, compiled forwards take the eager path, so
        # enable_torch_compile keeps its pre-fusion behavior instead of hitting
        # an unregistered launch inside the compiled graph.
        return False
    # LayerNorm family only (RMS keeps its own path); weight/bias fp32 as
    # FP32LayerNorm guarantees after .float().
    if not isinstance(norm, torch.nn.LayerNorm):
        return False
    if residual.dim() != 3 or residual.shape != x.shape:
        return False
    if residual.dtype != x.dtype or residual.dtype not in _SUPPORTED_STREAM_DTYPES:
        return False
    if not (residual.is_cuda and x.is_cuda):
        return False
    if not (residual.is_contiguous() and x.is_contiguous()):
        return False
    batch, seq, hidden = residual.shape
    if hidden > _MAX_HIDDEN or norm.normalized_shape != (hidden, ):
        return False
    if isinstance(gate, torch.Tensor):
        if gate.dtype not in (torch.float32, residual.dtype) or not gate.is_cuda:
            return False
        if _broadcast_strides(gate, batch, seq, hidden) is None:
            return False
    elif gate != 1:
        return False
    if (shift is None) != (scale is None):
        return False
    if shift is not None and scale is not None:
        for t in (shift, scale):
            if t.dtype not in (torch.float32, residual.dtype) or not t.is_cuda:
                return False
            if _broadcast_strides(t, batch, seq, hidden) is None:
                return False
    return True

def fused_residual_norm_mod(
    residual: torch.Tensor,
    x: torch.Tensor,
    gate: torch.Tensor | int,
    shift: torch.Tensor | None,
    scale: torch.Tensor | None,
    norm: torch.nn.Module,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run the fused kernel. Caller must have checked fused_path_supported().

    Returns (modulated_or_normalized, residual_out), both in the stream dtype.
    """
    batch, seq, hidden = residual.shape
    stream_dtype = residual.dtype
    out = torch.empty_like(residual)
    res_out = torch.empty_like(residual)

    gate_is_tensor = isinstance(gate, torch.Tensor)
    has_mod = shift is not None
    # Eager promotion keeps the chain fp32 (no intermediate rounding) exactly
    # when the gate is a fp32 tensor; the scalar-1 gate path materializes
    # stream-dtype intermediates.
    intermediate_round = (not gate_is_tensor or gate.dtype == stream_dtype) and stream_dtype == torch.bfloat16

    if gate_is_tensor:
        g_sb, g_ss = _broadcast_strides(gate, batch, seq, hidden)  # type: ignore[misc]
        gate_arg = gate
    else:
        g_sb = g_ss = 0
        gate_arg = residual  # unused placeholder pointer
    if has_mod:
        sh_sb, sh_ss = _broadcast_strides(shift, batch, seq, hidden)  # type: ignore[arg-type,misc]
        sc_sb, sc_ss = _broadcast_strides(scale, batch, seq, hidden)  # type: ignore[arg-type,misc]
        shift_arg, scale_arg = shift, scale
    else:
        sh_sb = sh_ss = sc_sb = sc_ss = 0
        shift_arg = scale_arg = residual  # unused placeholder pointers

    weight = norm.weight
    bias = norm.bias
    has_affine = weight is not None
    has_bias = bias is not None

    block_h = triton.next_power_of_2(hidden)
    num_warps = 4 if block_h <= 2048 else 8

    _fused_residual_norm_mod_kernel[(batch * seq, )](
        x,
        residual,
        gate_arg,
        weight if has_affine else residual,
        bias if has_bias else residual,
        shift_arg,
        scale_arg,
        out,
        res_out,
        seq,
        hidden,
        x.stride(0),
        x.stride(1),
        residual.stride(0),
        residual.stride(1),
        g_sb,
        g_ss,
        sh_sb,
        sh_ss,
        sc_sb,
        sc_ss,
        out.stride(0),
        out.stride(1),
        res_out.stride(0),
        res_out.stride(1),
        norm.eps,
        HAS_GATE_TENSOR=gate_is_tensor,
        INTERMEDIATE_ROUND=intermediate_round,
        HAS_AFFINE=has_affine,
        HAS_BIAS=has_bias,
        HAS_MOD=has_mod,
        OUT_BF16=stream_dtype == torch.bfloat16,
        BLOCK_H=block_h,
        num_warps=num_warps,
    )
    return out, res_out