def fused_dora_compose(
    lora: torch.Tensor,
    base: torch.Tensor,
    mag_norm_scale: torch.Tensor,
    scale: float,
    inplace: bool = True,
) -> torch.Tensor:
    """
    Fused DoRA composition: out = (mag_norm_scale - 1) * base + mag_norm_scale * (scale * lora).

    Algebraically equivalent to ``mag * (scale * lora + base) - base`` but uses
    the numerically stable form that avoids catastrophic cancellation in bf16/fp16.

    Args:
        lora: LoRA result tensor [..., out_features]
        base: Base layer result tensor, same shape as lora
        mag_norm_scale: Magnitude/norm scale, broadcastable to lora shape.
                        Expected shape [1, out_features] or [out_features].
        scale: Scalar LoRA scaling factor.
        inplace: If True, write result into lora tensor.

    Returns:
        The composed result tensor (may be lora if inplace=True).
    """
    if (
        not _is_dynamo_compiling()
        and _TRITON_AVAILABLE
        and lora.is_cuda
        and lora.is_contiguous()
        and base.is_contiguous()
        and mag_norm_scale.is_contiguous()
        and _mag_broadcasts_last_dim(mag_norm_scale, lora)
        and lora.dtype == base.dtype == mag_norm_scale.dtype
    ):
        return _fused_dora_compose_triton(lora, base, mag_norm_scale, scale, inplace)
    return _fused_dora_compose_torch(lora, base, mag_norm_scale, scale, inplace)

def fused_dora_forward_and_inner(
    lora: torch.Tensor,
    base: torch.Tensor,
    mag_norm_scale: torch.Tensor,
    scale: float,
) -> tuple:
    """
    Compute both the DoRA composition output and the ``inner`` tensor in one pass.

    Returns:
        (out, inner) where:
            out   = (mag_norm_scale - 1) * base + mag_norm_scale * (scale * lora)
            inner = scale * lora + base

    When Triton is available and tensors are on CUDA, a single fused kernel
    computes both outputs simultaneously — the intermediate ``scaled_lora``
    never leaves SRAM registers.  This eliminates the VRAM spike caused by
    sequential PyTorch ops in ``FusedDoRAComposeFunction.forward``.
    """
    if (
        not _is_dynamo_compiling()
        and _TRITON_AVAILABLE
        and lora.is_cuda
        and lora.is_contiguous()
        and base.is_contiguous()
        and mag_norm_scale.is_contiguous()
        and _mag_broadcasts_last_dim(mag_norm_scale, lora)
        and lora.dtype == base.dtype == mag_norm_scale.dtype
    ):
        return _fused_dora_forward_and_inner_triton(lora, base, mag_norm_scale, scale)
    return _fused_dora_forward_and_inner_torch(lora, base, mag_norm_scale, scale)

def fused_norm_assembly(
    w_norm_sq: torch.Tensor,
    cross_term: torch.Tensor,
    ba_norm_sq: torch.Tensor,
    scale: float,
) -> tuple:
    """
    Fused norm assembly: compute weight_norm from components.

    Args:
        w_norm_sq: ||W||^2 per row, shape [out_features]
        cross_term: <W, BA> per row, shape [out_features]
        ba_norm_sq: ||BA||^2 per row, shape [out_features]
        scale: LoRA scaling factor

    Returns:
        (weight_norm,) — magnitude division is always done in PyTorch by the caller.
    """
    if (
        not _is_dynamo_compiling()
        and _TRITON_AVAILABLE
        and w_norm_sq.is_cuda
        and w_norm_sq.is_contiguous()
        and cross_term.is_contiguous()
        and ba_norm_sq.is_contiguous()
        # The Triton norm kernel uses inline PTX (sqrt.rn.f32) which is
        # NVIDIA-specific.  On ROCm/HIP (torch.version.hip is set), fall
        # back to PyTorch to avoid a crash from NVIDIA-only assembly.
        and torch.version.hip is None
    ):
        return _fused_norm_assembly_triton(
            w_norm_sq,
            cross_term,
            ba_norm_sq,
            scale,
        )
    return _fused_norm_assembly_torch(
        w_norm_sq,
        cross_term,
        ba_norm_sq,
        scale,
    )

def fused_dora_compose_autograd(
    lora: torch.Tensor,
    base: torch.Tensor,
    mag_norm_scale: torch.Tensor,
    scale: float,
) -> torch.Tensor:
    """
    DoRA composition with custom autograd for fused backward.

    This is the main entry point for training-time composition that benefits
    from the fused backward pass.

    When ``torch.compile`` is active and a ``torch.library.custom_op`` is
    registered (PyTorch 2.4+), this routes through ``peft::fused_dora_compose``
    so that Dynamo sees a single opaque graph node.  In eager mode, this
    always uses ``FusedDoRAComposeFunction.apply`` with Triton kernels.

    Args:
        lora: LoRA output tensor [..., out_features]
        base: Base result tensor, same shape as lora
        mag_norm_scale: Magnitude/norm scale [1, out_features]
        scale: Scalar LoRA scaling factor

    Returns:
        out = (mag_norm_scale - 1) * base + mag_norm_scale * (scale * lora)
    """
    if _is_dynamo_compiling():
        if _HAS_CUSTOM_OP:
            return _fused_dora_compose_custom_op(lora, base, mag_norm_scale, scale)
        # Fallback for older PyTorch without custom_op support.
        return (mag_norm_scale - 1) * base + mag_norm_scale * (scale * lora)
    return FusedDoRAComposeFunction.apply(lora, base, mag_norm_scale, scale)

def _fused_dora_compose_torch(
    lora: torch.Tensor,
    base: torch.Tensor,
    mag_norm_scale: torch.Tensor,
    scale: float,
    inplace: bool = True,
) -> torch.Tensor:
    """Pure PyTorch fallback for fused DoRA composition.

    Uses the numerically stable form ``(mag - 1) * base + mag * (scale * lora)``
    to avoid catastrophic cancellation in bf16/fp16 when mag ≈ 1.

    Note:
        The caller (``dora.py``) is responsible for ensuring ``lora`` is a
        freshly-computed tensor (e.g. ``lora_B(lora_A(x))``) and does not
        share storage with other live tensors when ``inplace=True``.
    """
    if inplace:
        if torch.promote_types(torch.promote_types(lora.dtype, base.dtype), mag_norm_scale.dtype) != lora.dtype:
            result = (mag_norm_scale - 1) * base + mag_norm_scale * (scale * lora)
            lora.copy_(result)
            return lora

        # Numerically stable in-place: lora = (mag-1)*base + mag*(scale*lora)
        # Two mul_ calls preserve canonical evaluation order mag*(scale*lora),
        # matching the out-of-place path and the Triton kernel for same-dtype
        # bitwise parity.
        lora.mul_(scale)
        lora.mul_(mag_norm_scale)
        lora.add_(base * (mag_norm_scale - 1))
        return lora
    else:
        result = (mag_norm_scale - 1) * base + mag_norm_scale * (scale * lora)
        # Cast to lora.dtype for consistent output contract (matches inplace branch).
        if result.dtype != lora.dtype:
            result = result.to(lora.dtype)
        return result

def _fused_dora_forward_and_inner_torch(
    lora: torch.Tensor,
    base: torch.Tensor,
    mag_norm_scale: torch.Tensor,
    scale: float,
) -> tuple:
    """Pure PyTorch fallback for fused forward-and-inner computation.

    Computes both outputs without Triton.  ``scaled_lora`` is a temporary
    that is freed after ``inner`` and ``out`` are computed.
    """
    result_dtype = lora.dtype
    scaled_lora = scale * lora
    inner = scaled_lora + base
    out = (mag_norm_scale - 1) * base + mag_norm_scale * scaled_lora
    if out.dtype != result_dtype:
        out = out.to(result_dtype)
    # Cast inner to match the Triton path's output contract (inner in
    # lora.dtype).  Without this, mixed-dtype inputs produce an fp32 inner
    # that doubles the activation VRAM saved for backward (d_mag reduction).
    if inner.dtype != result_dtype:
        inner = inner.to(result_dtype)
    return out, inner

def _fused_norm_assembly_torch(
    w_norm_sq: torch.Tensor,
    cross_term: torch.Tensor,
    ba_norm_sq: torch.Tensor,
    scale: float,
) -> tuple:
    """Pure PyTorch fallback for fused norm assembly."""
    # Use Python float directly instead of torch.as_tensor to avoid allocating
    # a new scalar tensor on every call.  PyTorch handles scalar-tensor
    # promotion natively in the arithmetic below.
    s = float(scale)
    norm_sq = w_norm_sq + (2.0 * s) * cross_term + (s * s) * ba_norm_sq
    # Keep this out-of-place to avoid in-place autograd hazards if this helper
    # is called with grad tracking enabled.
    norm_sq = norm_sq.clamp_min(0)
    weight_norm = torch.sqrt(norm_sq)

    return (weight_norm,)

def _fused_backward_torch(d_out, inner, mag_norm_scale, scale, needs_lora_grad, needs_base_grad, needs_mag_grad):
    """Pure PyTorch fallback for fused backward.

    ``inner`` is ``scale * lora + base`` (precomputed in forward).
    """
    d_lora = d_base = d_mag = None

    if needs_lora_grad:
        d_lora = mag_norm_scale * scale * d_out

    if needs_base_grad:
        d_base = (mag_norm_scale - 1) * d_out

    if needs_mag_grad:
        assert inner is not None, "inner must be saved when mag requires grad"
        # inner was computed under no_grad in forward and may have a different
        # dtype than d_out when AMP autocast is active (e.g. inner in fp32
        # while d_out in fp16).  Align dtypes to avoid silent precision loss.
        if inner.dtype != d_out.dtype:
            inner = inner.to(d_out.dtype)
        d_mag_full = inner * d_out
        # Reduce over dimensions that were broadcast from mag_norm_scale.
        # For Linear [B, F] with mag [1, F] → reduce dim 0.
        # For Conv2d [N, C, H, W] with mag [1, C, 1, 1] → reduce dims 0,2,3.
        sum_dims = _broadcast_reduce_dims(d_out.shape, mag_norm_scale.shape)
        if sum_dims:
            d_mag = d_mag_full.sum(dim=sum_dims, keepdim=True)
        else:
            d_mag = d_mag_full
        # Reshape to match mag_norm_scale shape
        d_mag = d_mag.reshape(mag_norm_scale.shape)

    return d_lora, d_base, d_mag