code/models/tt_dit/layers/lora.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""
Reusable LoRA mixin for the Linear family.

LoRA-aware Linears keep a *host bank* of registered adapters and apply
the active adapter via one of two execution paths, chosen at construction
with ``lora_mode``:

    fuse     — merge the delta into the base weight at bind time:
                   W += scale * A.T @ B.T   on bind
                   W -= scale * A.T @ B.T   on unbind
               Forward runs the un-modified base path with every fused
               kernel intact (matmul+gelu, matmul+addcmul, etc.) and
               forward-time LoRA overhead is zero. Swap cost is bounded
               to ~one small matmul + an in-place add per Linear.

    runtime  — keep A, B on device and add their effect to every forward:
                   y = base(x) + scale * (x @ A) @ B
               Swap cost is microseconds (a host-side index assignment +
               re-uploading A/B). Per-forward cost is two extra small
               matmuls. Useful when the per-forward overhead is cheaper
               than the bind-time matmul (very tall stacks of cheap adapters,
               or quick A/B experiments).

Both A and B are uploaded with sharding inferred from W's own ``mesh_axes``
so the matmul lands the right shard:
  - A is [in, rank] → shard A's dim-0 the same way W's dim-0 is sharded.
  - B is [rank, out] → shard B's dim-1 the same way W's dim-1 is sharded.

Runtime mode supports replicated outputs only. Chunked-output Linears
(``chunks`` set on ColParallel) must use fuse mode — splitting the LoRA
delta across the same chunk boundaries is not implemented and would
defeat the runtime path's main advantage (zero CCL).

Runtime mode is also incompatible with the fused-addcmul paths
(``RowParallelLinear.forward_fused_addcmul``,
``WanAttention._to_out_fused_addcmul``): those kernels combine matmul,
reduce-scatter, and addcmul in a single op and read ``weight.data``
directly, so the runtime A/B tensors have nowhere to participate. Fuse
mode works on these paths automatically because the delta lives in
``weight.data``.

TODO: ``lora_enabled`` is currently threaded through 5 constructors
(``WanPipeline`` → ``WanTransformer3DModel`` → ``WanTransformerBlock`` →
``WanAttention`` / ``ParallelFeedForward``) to pick between the LoRA
class and its base. The LoRA classes could instead be drop-in
replacements with a zero-overhead empty-bank fast path in
``LoRAMixin.forward``, removing the flag entirely. Deferred to a
follow-up — needs perf validation that the empty-bank forward is
truly identical to the base.

Note: bind/unbind in fuse mode adds and subtracts the delta in bf16,
which is not bit-exactly inverse. W drifts by ~O(N · eps_bf16) over N
swap cycles. In production this is bounded automatically by
``dynamic_load=True`` (``cache.load_model`` restores the cached base
on every page-in). For static pipelines (``dynamic_load=False``) that
need bit-exact restore across many swaps, snapshot the state dict once
at construction and call ``load_torch_state_dict`` to re-seed W when
drift becomes material. (With the A/B cache enabled — ``cache_capacity>0``
— unbind subtracts the *same* device delta that bind added for a cached
adapter, so those swap cycles are exact and do not drift.)

Fuse-mode A/B cache (``cache_capacity``)
----------------------------------------
Each bind/unbind/reload otherwise re-uploads A and B from host and
re-lays them out for W's sharding. With ``cache_capacity>0`` the uploaded
device factors are kept in a per-module LRU (keyed by bank index), so a
re-bind of an already-seen adapter — and, under ``dynamic_load``, the
re-merge after every page-in — skips the host upload. The cache holds
only the rank-sized A/B (never the full delta), survives page-out (the
factors are plain tensors, not tracked ``Parameters``, so
``deallocate_weights`` leaves them alone), and is bounded to
``cache_capacity`` adapters per Linear. ``cache_capacity=0`` disables it
and restores the upload-then-free behavior.
"""
from __future__ import annotations

from collections import OrderedDict
from dataclasses import dataclass

import torch

import ttnn

from ..utils.tensor import from_torch as _from_torch


@dataclass
class LoRAAdapter:
    """Host-resident LoRA adapter record. ``A`` and ``B`` are torch tensors
    in PyTorch LoRA layout (A: [rank, in], B: [out, rank])."""

    name: str
    A: torch.Tensor
    B: torch.Tensor
    rank: int
    scale: float


class LoRAMixin:
    """Adds a LoRA bank + active-adapter state to a Linear-family Module.

    Subclasses call ``_init_lora_state(mode=...)`` at the end of their
    ``__init__``. The mixin handles bind/unbind and (for ``runtime`` mode)
    overrides ``forward`` to add the delta.
    """

    # Layout-neutral (adds no __dict__/__weakref__ of its own) so a plain
    # Linear instance can be promoted in place via ``__class__`` assignment
    # (see experimental/lora/promote.py); state still lives in Module's __dict__.
    __slots__ = ()

    LORA_MODES = ("fuse", "runtime")

    # ---- one-time state init ----
    def _init_lora_state(self, mode: str = "fuse", cache_capacity: int = 0) -> None:
        if mode not in self.LORA_MODES:
            raise ValueError(f"lora_mode must be one of {self.LORA_MODES}; got {mode!r}")
        self.lora_mode = mode
        self.lora_bank: list[LoRAAdapter | None] = []
        self.active_idx: int | None = None
        self.active_scale: float = 1.0
        # Fuse-mode LRU of device-resident A/B factors, keyed by bank index:
        # {idx: (A_dev, B_dev, scale)}, MRU last. Bounded to ``lora_cache_capacity``
        # adapters; 0 disables caching (upload-then-free per bind).
        self.lora_cache_capacity: int = int(cache_capacity)
        self._ab_cache: OrderedDict[int, tuple[ttnn.Tensor, ttnn.Tensor, float]] = OrderedDict()
        # ``active_idx`` records intent (which adapter should be used).
        # ``_delta_applied`` (fuse mode) records whether the merge is
        # currently present on the on-device weight. These can diverge
        # under dynamic_load: when the transformer is paged out the
        # device weight is gone; when it pages back in the cached base
        # weight is restored and the delta must be re-applied.
        self._delta_applied: bool = False
        # Device-resident A/B for runtime mode (None when no adapter
        # is bound or when the layer is paged out).
        self._runtime_A: ttnn.Tensor | None = None
        self._runtime_B: ttnn.Tensor | None = None

    # ---- introspection ----
    @property
    def is_lora_active(self) -> bool:
        """True if any LoRA is currently bound (whether or not the on-device
        side state is currently realized — e.g. paged out under dynamic_load)."""
        return self.active_idx is not None

    # ---- bank management ----
    def register_lora(
        self,
        A_torch: torch.Tensor,
        B_torch: torch.Tensor,
        scale: float = 1.0,
        name: str = "",
    ) -> int:
        # swiglu doubles out_features AND head-permutes W in _prepare_torch_state;
        # the LoRA delta would not undergo the same permutation, so the merged
        # weight's gate/value split would be wrong. Refuse rather than silently
        # corrupt — and avoid the confusing 2*out validation error.
        if getattr(self, "fused_activation_fn", None) is None and getattr(self, "activation_fn", None) == "swiglu":
            raise ValueError(
                "LoRA adapters are not supported on swiglu layers; the merged "
                "delta would not match the swiglu head-permuted weight layout"
            )

        rank = A_torch.shape[0]
        if A_torch.shape != (rank, self.in_features):
            raise ValueError(
                f"A must be [rank, in_features]; got {tuple(A_torch.shape)} " f"expected (*, {self.in_features})"
            )
        if B_torch.shape != (self.out_features, rank):
            raise ValueError(
                f"B must be [out_features, rank]; got {tuple(B_torch.shape)} " f"expected ({self.out_features}, {rank})"
            )

        adapter = LoRAAdapter(
            name=name,
            A=A_torch.detach().clone(),
            B=B_torch.detach().clone(),
            rank=rank,
            scale=float(scale),
        )
        self.lora_bank.append(adapter)
        return len(self.lora_bank) - 1

    def unregister_lora(self, idx: int) -> None:
        if not (0 <= idx < len(self.lora_bank)):
            return
        if self.active_idx == idx:
            self.unbind_active()
        self._drop_cached_ab(idx)
        self.lora_bank[idx] = None

    def set_lora_cache_capacity(self, n: int) -> None:
        """Set how many adapters keep their device A/B factors resident (LRU).
        Shrinking evicts the least-recently-used down to ``n``; ``n<=0`` frees
        the whole cache and disables further caching."""
        self.lora_cache_capacity = int(n)
        if self.lora_cache_capacity <= 0:
            self._free_ab_cache()
        else:
            self._evict_ab_cache()

    # ---- bind / unbind ----
    def bind_active(self, idx: int, scale: float | None = None) -> None:
        """Make bank[idx] the active LoRA.

        If a different adapter is currently bound, it is removed first.
        When the layer's weights are not currently on device (dynamic_load
        page-out), the merge is deferred — ``active_idx`` and ``active_scale``
        record the intent and ``reapply_after_load`` will pick it up on
        the next load.
        """
        # Validate the target slot BEFORE touching the current active adapter,
        # so a bad idx doesn't leave W with the prior delta half-undone.
        if not (0 <= idx < len(self.lora_bank)) or self.lora_bank[idx] is None:
            raise IndexError(f"invalid lora slot {idx}")

        adapter = self.lora_bank[idx]
        use_scale = adapter.scale if scale is None else float(scale)

        # Idempotent: don't subtract-and-re-add the same delta — that round-trip
        # accumulates bf16 quantization drift in W.
        if self.active_idx == idx and self.active_scale == use_scale:
            return

        if self.active_idx is not None:
            self._undo_active()

        self.active_idx = idx
        self.active_scale = use_scale
        if self.is_loaded():
            self._realize_active()

    def unbind_active(self) -> None:
        if self.active_idx is None:
            return
        self._undo_active()

    def _realize_active(self) -> None:
        """Materialize the on-device effect of the active adapter.

        fuse → add the delta into ``W``. runtime → upload ``A`` and ``B``
        for the forward path to consume.
        """
        if self.lora_mode == "fuse":
            self._apply_delta(self.active_idx, self.active_scale, sign=+1)
            self._delta_applied = True
        else:  # runtime
            adapter = self.lora_bank[self.active_idx]
            self._upload_runtime_ab(adapter.A, adapter.B)

    def _undo_active(self) -> None:
        if self.is_loaded():
            if self.lora_mode == "fuse" and self._delta_applied:
                if self.lora_bank[self.active_idx] is not None:
                    self._apply_delta(self.active_idx, self.active_scale, sign=-1)
            elif self.lora_mode == "runtime":
                self._free_runtime_ab()
        self._delta_applied = False
        self.active_idx = None
        self.active_scale = 1.0

    def reapply_after_load(self) -> None:
        """Re-materialize the active adapter after a dynamic_load reload.

        Called by the pipeline after a reload restores the cached base
        weights and (for runtime mode) wipes the device-resident A/B.
        No-op when no adapter is active or the effect is already in place.
        """
        if self.active_idx is None or not self.is_loaded():
            return
        if self.lora_mode == "fuse" and self._delta_applied:
            return
        if self.lora_mode == "runtime" and self._runtime_A is not None:
            return
        self._realize_active()

    def deallocate_weights(self) -> None:  # type: ignore[override]
        # Weights about to be freed → any merged delta goes with them
        # and runtime A/B are tied to the layer being resident.
        self._delta_applied = False
        if self._runtime_A is not None or self._runtime_B is not None:
            self._free_runtime_ab()
        super().deallocate_weights()

    def deallocate_lora(self) -> None:
        if self.active_idx is not None:
            self.unbind_active()
        self._free_ab_cache()
        self.lora_bank = []

    # ---- forward override (runtime mode) ----
    def forward(self, x: ttnn.Tensor, *args, **kwargs):  # type: ignore[override]
        if self.lora_mode != "runtime" or self._runtime_A is None:
            return super().forward(x, *args, **kwargs)

        # Reference math: y = activation(linear(x, W, b) + delta). The base
        # Linear may fuse the activation into the matmul; if so, disable it
        # for this call so we can add the delta BEFORE the activation runs.
        saved_fused = getattr(self, "fused_activation_fn", None)
        saved_activation = getattr(self, "activation_fn", None)
        try:
            if saved_fused is not None:
                self.fused_activation_fn = None
            if saved_activation is not None:
                self.activation_fn = None
            base = super().forward(x, *args, **kwargs)
        finally:
            self.fused_activation_fn = saved_fused
            self.activation_fn = saved_activation

        if isinstance(base, list):
            raise RuntimeError(
                "runtime LoRA delta does not support chunked-output Linears; "
                "construct with lora_mode='fuse' for chunks>1"
            )
        # Honor a caller-passed compute_kernel_config override (the base
        # Linear path does), so runtime LoRA and the base matmul use the
        # same precision settings within one forward.
        delta_compute_config = kwargs.get("compute_kernel_config", self.compute_config)
        ax = ttnn.matmul(x, self._runtime_A, compute_kernel_config=delta_compute_config)
        delta = ttnn.matmul(ax, self._runtime_B, compute_kernel_config=delta_compute_config)
        ttnn.deallocate(ax)
        out = ttnn.add(base, delta)
        ttnn.deallocate(base)
        ttnn.deallocate(delta)

        if saved_fused is not None:
            # Only gelu is currently hoisted into fused_activation_fn.
            if saved_fused[0] == ttnn.UnaryOpType.GELU:
                out = ttnn.gelu(out)
            else:
                raise NotImplementedError(f"runtime LoRA cannot replay fused activation {saved_fused}")
        if saved_activation is not None:
            from .linear import _apply_activation_fn

            out = _apply_activation_fn(out, saved_activation)
        return out

    def forward_fused_addcmul(self, *args, **kwargs):  # type: ignore[override]
        # The fused matmul + reduce-scatter + addcmul kernel reads weight.data
        # directly and has nowhere to add a separate LoRA delta — runtime mode
        # would silently no-op. Fuse mode is fine because the delta is already
        # merged into weight.data.
        if self.lora_mode == "runtime" and self.is_lora_active:
            raise RuntimeError(
                "runtime LoRA mode is incompatible with forward_fused_addcmul; "
                "construct this layer with lora_mode='fuse'"
            )
        return super().forward_fused_addcmul(*args, **kwargs)

    # ---- internals ----
    def _apply_delta(self, idx: int, scale: float, sign: int) -> None:
        """Add (``sign>0``) or subtract (``sign<0``) bank[idx]'s ``+scale`` delta
        into ``self.weight.data`` in place (fuse mode). The device A/B factors
        come from the LRU cache when enabled, so a re-bind or a post-reload
        re-merge skips the host upload; the same cached ``+scale`` delta serves
        both bind and unbind, making the pair an exact negation. The full-size
        delta itself is never cached (that would cost a whole weight per adapter)."""
        A_dev, B_dev, owned = self._acquire_delta_ab(idx, float(scale))
        delta = ttnn.matmul(A_dev, B_dev, compute_kernel_config=self.compute_config)
        if owned:
            ttnn.deallocate(A_dev)
            ttnn.deallocate(B_dev)
        if sign > 0:
            ttnn.add(self.weight.data, delta, output_tensor=self.weight.data)
        else:
            ttnn.subtract(self.weight.data, delta, output_tensor=self.weight.data)
        ttnn.deallocate(delta)

    def _acquire_delta_ab(self, idx: int, scale: float) -> tuple[ttnn.Tensor, ttnn.Tensor, bool]:
        """Return ``(A_dev, B_dev, owned)`` for bank[idx]'s ``+scale`` delta
        factors, uploading + laying them out for W's sharding on a miss.
        ``owned=True`` means the factors are uncached and the caller must
        deallocate them; cached factors are owned by the LRU (``owned=False``)."""
        adapter = self.lora_bank[idx]
        if self.lora_cache_capacity <= 0:
            A_dev, B_dev = self._upload_ab_for_w_sharding(adapter.A, adapter.B, scale=scale)
            return A_dev, B_dev, True
        cached = self._ab_cache.get(idx)
        if cached is not None and cached[2] == scale:
            self._ab_cache.move_to_end(idx)  # mark most-recently-used
            return cached[0], cached[1], False
        if cached is not None:  # same slot bound at a new scale → rebuild
            self._drop_cached_ab(idx)
        A_dev, B_dev = self._upload_ab_for_w_sharding(adapter.A, adapter.B, scale=scale)
        self._ab_cache[idx] = (A_dev, B_dev, scale)  # inserted at the MRU end
        self._evict_ab_cache()
        return A_dev, B_dev, False

    def _evict_ab_cache(self) -> None:
        # Trim from the LRU end. The active adapter is always the MRU (just
        # inserted or touched), so for capacity >= 1 it is never evicted.
        while len(self._ab_cache) > self.lora_cache_capacity:
            _, (A_old, B_old, _) = self._ab_cache.popitem(last=False)
            ttnn.deallocate(A_old)
            ttnn.deallocate(B_old)

    def _drop_cached_ab(self, idx: int) -> None:
        entry = self._ab_cache.pop(idx, None)
        if entry is not None:
            ttnn.deallocate(entry[0])
            ttnn.deallocate(entry[1])

    def _free_ab_cache(self) -> None:
        for A_dev, B_dev, _ in self._ab_cache.values():
            ttnn.deallocate(A_dev)
            ttnn.deallocate(B_dev)
        self._ab_cache.clear()

    def _upload_runtime_ab(self, A_torch: torch.Tensor, B_torch: torch.Tensor) -> None:
        """Upload A and B for the runtime forward path. The active scale
        is baked into B so forward stays as two pure matmuls + one add."""
        self._runtime_A, self._runtime_B = self._upload_ab_for_w_sharding(A_torch, B_torch, scale=self.active_scale)

    def _free_runtime_ab(self) -> None:
        if self._runtime_A is not None:
            ttnn.deallocate(self._runtime_A)
            self._runtime_A = None
        if self._runtime_B is not None:
            ttnn.deallocate(self._runtime_B)
            self._runtime_B = None

    def _upload_ab_for_w_sharding(
        self,
        A_torch: torch.Tensor,
        B_torch: torch.Tensor,
        *,
        scale: float,
    ) -> tuple[ttnn.Tensor, ttnn.Tensor]:
        """Upload ``A`` ([rank,in] → [in,rank]) and ``scale*B`` ([out,rank] → [rank,out])
        with the sharding of ``W``:
          - A's dim 0 sharded like W's dim 0
          - B's dim 1 sharded like W's dim 1
        So that matmul(A, B) yields a tensor with W's combined sharding."""
        w_axes = self.weight.mesh_axes
        # Always take the multiply path so bind (+scale) and unbind (-scale)
        # exercise the same code with the same fp32→bf16 rounding boundaries;
        # the on-device deltas then cleanly negate each other.
        B_eff = B_torch * scale
        A_dev = _from_torch(
            A_torch.transpose(0, 1).contiguous().to(torch.bfloat16),
            device=self.mesh_device,
            layout=ttnn.Layout.TILE,
            dtype=ttnn.bfloat16,
            mesh_axes=(w_axes[0], None),
        )
        B_dev = _from_torch(
            B_eff.transpose(0, 1).contiguous().to(torch.bfloat16),
            device=self.mesh_device,
            layout=ttnn.Layout.TILE,
            dtype=ttnn.bfloat16,
            mesh_axes=(None, w_axes[1]),
        )
        return A_dev, B_dev