modulation ¶

def fused_residual_gate_rmsnorm_modulate(
    residual: torch.Tensor,
    branch: torch.Tensor,
    gate: torch.Tensor,
    weight: torch.Tensor,
    scale: torch.Tensor,
    shift: torch.Tensor,
    index: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fuse residual update, row-indexed gate, RMSNorm, and modulation.

    ``index`` values must lie in ``[0, table_rows)``; see
    :func:`fused_rmsnorm_modulate` for why the wrapper does not check them.
    """
    _require_forward_only(residual, branch, gate, weight, scale, shift)
    if torch.compiler.is_compiling():
        return torch.ops.fastvideo._minimax_h3_residual_gate_rmsnorm_modulate(
            residual,
            branch,
            gate,
            weight,
            scale,
            shift,
            index,
            eps,
        )
    return _fused_residual_gate_rmsnorm_modulate_impl(
        residual,
        branch,
        gate,
        weight,
        scale,
        shift,
        index,
        eps,
    )

def fused_rmsnorm_modulate(
    x: torch.Tensor,
    weight: torch.Tensor,
    scale: torch.Tensor,
    shift: torch.Tensor,
    index: torch.Tensor,
    eps: float,
) -> torch.Tensor:
    """Run RMSNorm and row-indexed modulation in one strict Triton kernel.

    ``index`` values must lie in ``[0, table_rows)``. Unlike eager
    ``index_select``, the kernel does not raise on out-of-range values (a
    device-side bounds check would synchronize); callers are safe by
    construction (``timestep_indices * 3 + token_tags``, SP pads with 0).
    """
    _require_forward_only(x, weight, scale, shift)
    if torch.compiler.is_compiling():
        return torch.ops.fastvideo._minimax_h3_rmsnorm_modulate(x, weight, scale, shift, index, eps)
    return _fused_rmsnorm_modulate_impl(x, weight, scale, shift, index, eps)