LoRA & Parameter-Efficient Fine-Tuning

"""LoRA, transcribed from the objective.

No ML library. The point of writing it this way is that three claims
about this technique become things you can read rather than things you
are told:

    h = W0 @ x + (alpha / r) * B @ (A @ x)

First, W0 never appears on the left of an assignment after setup. It is
frozen, and freezing it is not a flag on an optimizer -- it is the
absence of any code that updates it.

Second, dL/dW0 is never formed. Look at backward() below: it computes
the gradient flowing BACK through the layer (needed, because layers
below this one still have adapters), and it computes dL/dA and dL/dB.
The d-by-k weight gradient that full fine-tuning would materialize is
simply not there. That single missing array is the entire memory
argument.

Third, B starts at zero, so the adapted model starts EXACTLY at the
pretrained model. Not approximately. There is nothing to warm up.
"""

import math
import random

Matrix = list[list[float]]
Vector = list[float]


def zeros(rows: int, cols: int) -> Matrix:
    return [[0.0 for _ in range(cols)] for _ in range(rows)]


def matvec(matrix: Matrix, vector: Vector) -> Vector:
    out = []
    for row in matrix:
        total = 0.0
        for index in range(len(vector)):
            total += row[index] * vector[index]
        out.append(total)
    return out


def matvec_transposed(matrix: Matrix, vector: Vector) -> Vector:
    """matrix.T @ vector, without building the transpose."""
    cols = len(matrix[0])
    out = [0.0] * cols
    for row_index, row in enumerate(matrix):
        scale = vector[row_index]
        if scale == 0.0:
            continue
        for col_index in range(cols):
            out[col_index] += row[col_index] * scale
    return out


def add_outer(target: Matrix, left: Vector, right: Vector, scale: float) -> None:
    """target += scale * outer(left, right), in place."""
    for row_index, left_value in enumerate(left):
        if left_value == 0.0:
            continue
        row = target[row_index]
        for col_index, right_value in enumerate(right):
            row[col_index] += scale * left_value * right_value


class LoRALinear:
    """One adapted projection: a frozen W0 plus a learned rank-r update.

    The shapes are the whole story. W0 is d-by-k and holds d*k numbers
    that will never change. A is r-by-k and B is d-by-r, and together
    they hold r*(d+k), which for d = k = 4096 and r = 8 is about 65
    thousand against 16.8 million.
    """

    def __init__(self, d_out, d_in, rank, alpha, rng):
        if rank <= 0 or rank > min(d_out, d_in):
            raise ValueError("rank must be in 1..min(d_out, d_in)")

        # Frozen. Stands in for a pretrained weight; in reality this is
        # loaded, and it is the only large array in this class.
        self.w0 = [
            [rng.gauss(0.0, 1.0 / math.sqrt(d_in)) for _ in range(d_in)]
            for _ in range(d_out)
        ]

        # A is initialized randomly and B at zero, so B @ A == 0 at step
        # zero and the layer is exactly W0. Reversing which one is zero
        # would work too; what must not happen is both being random,
        # which perturbs a pretrained model before training starts.
        self.a = [
            [rng.gauss(0.0, 1.0 / math.sqrt(d_in)) for _ in range(d_in)]
            for _ in range(rank)
        ]
        self.b = zeros(d_out, rank)

        self.rank = rank
        self.alpha = alpha
        # alpha / r, not alpha. Dividing by the rank is what lets you
        # raise r without also raising the magnitude of the update --
        # the difference between sweeping capacity and sweeping two
        # things at once.
        self.scale = alpha / rank

        self.grad_a = zeros(rank, d_in)
        self.grad_b = zeros(d_out, rank)
        self.cached_x = None
        self.cached_ax = None

    def forward(self, x: Vector) -> Vector:
        base = matvec(self.w0, x)
        projected = matvec(self.a, x)              # r numbers
        update = matvec(self.b, projected)         # d numbers
        self.cached_x = x
        self.cached_ax = projected
        return [base[i] + self.scale * update[i] for i in range(len(base))]

    def backward(self, grad_h: Vector) -> Vector:
        """Accumulate dL/dA and dL/dB; return dL/dx.

        dL/dx still needs W0, because layers below this one have
        adapters of their own and their gradients have to arrive. This
        is why LoRA saves memory and not much time: the backward pass
        walks the full depth exactly as it would otherwise.
        """
        # dL/dB = scale * grad_h (A x)^T
        add_outer(self.grad_b, grad_h, self.cached_ax, self.scale)

        # dL/dA = scale * (B^T grad_h) x^T
        back_through_b = matvec_transposed(self.b, grad_h)
        add_outer(self.grad_a, back_through_b, self.cached_x, self.scale)

        # dL/dx = W0^T grad_h + scale * A^T B^T grad_h.
        # Note what is absent: no outer product against W0 is ever
        # formed. That array is dL/dW0, and it is the one full
        # fine-tuning cannot avoid.
        through_base = matvec_transposed(self.w0, grad_h)
        through_adapter = matvec_transposed(self.a, back_through_b)
        return [
            through_base[i] + self.scale * through_adapter[i]
            for i in range(len(through_base))
        ]

    def zero_grad(self) -> None:
        for row in self.grad_a:
            for index in range(len(row)):
                row[index] = 0.0
        for row in self.grad_b:
            for index in range(len(row)):
                row[index] = 0.0

    def trainable_parameters(self) -> int:
        return self.rank * (len(self.a[0]) + len(self.b))

    def frozen_parameters(self) -> int:
        return len(self.w0) * len(self.w0[0])

    def merge(self) -> Matrix:
        """W0 + (alpha / r) B A, materialized.

        Zero added latency afterwards, because the result is an ordinary
        weight matrix. And the adapter is gone: one merged model serves
        one adapter, which is exactly the property the multi-tenant case
        is built on not having.
        """
        merged = [row[:] for row in self.w0]
        for out_index in range(len(self.b)):
            b_row = self.b[out_index]
            target = merged[out_index]
            for rank_index in range(self.rank):
                weight = self.scale * b_row[rank_index]
                if weight == 0.0:
                    continue
                a_row = self.a[rank_index]
                for in_index in range(len(a_row)):
                    target[in_index] += weight * a_row[in_index]
        return merged


class Adam:
    """Two moment buffers per trainable parameter, allocated here.

    Written out because this is the cost LoRA removes. Every buffer
    below is sized to the adapter, not to W0, and that ratio -- not the
    forward arithmetic -- is why the technique exists.
    """

    def __init__(self, shapes, lr, beta1=0.9, beta2=0.999, eps=1e-8):
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.step_count = 0
        self.m = [zeros(rows, cols) for rows, cols in shapes]
        self.v = [zeros(rows, cols) for rows, cols in shapes]

    def step(self, params, grads):
        self.step_count += 1
        bias1 = 1.0 - self.beta1 ** self.step_count
        bias2 = 1.0 - self.beta2 ** self.step_count

        for slot, (param, grad) in enumerate(zip(params, grads)):
            m_buffer = self.m[slot]
            v_buffer = self.v[slot]
            for row_index, row in enumerate(param):
                grad_row = grad[row_index]
                m_row = m_buffer[row_index]
                v_row = v_buffer[row_index]
                for col_index in range(len(row)):
                    gradient = grad_row[col_index]
                    m_row[col_index] = (
                        self.beta1 * m_row[col_index] + (1.0 - self.beta1) * gradient
                    )
                    v_row[col_index] = (
                        self.beta2 * v_row[col_index]
                        + (1.0 - self.beta2) * gradient * gradient
                    )
                    m_hat = m_row[col_index] / bias1
                    v_hat = v_row[col_index] / bias2
                    row[col_index] -= self.lr * m_hat / (math.sqrt(v_hat) + self.eps)


def squared_error(prediction: Vector, target: Vector) -> float:
    return sum((p - t) ** 2 for p, t in zip(prediction, target)) / len(target)


def squared_error_grad(prediction: Vector, target: Vector) -> Vector:
    scale = 2.0 / len(target)
    return [scale * (p - t) for p, t in zip(prediction, target)]


def train(layer: LoRALinear, samples, epochs: int, lr: float):
    """The ordinary loop. The loss is the ordinary loss.

    Nothing in here is LoRA-specific, which is the point: the objective
    did not change, only the set of parameters the optimizer was handed.
    """
    optimizer = Adam(
        shapes=[(len(layer.a), len(layer.a[0])), (len(layer.b), len(layer.b[0]))],
        lr=lr,
    )

    history = []
    for _ in range(epochs):
        total = 0.0
        layer.zero_grad()
        for x, target in samples:
            prediction = layer.forward(x)
            total += squared_error(prediction, target)
            layer.backward(squared_error_grad(prediction, target))
        optimizer.step([layer.a, layer.b], [layer.grad_a, layer.grad_b])
        history.append(total / len(samples))
    return history


def rank_ceiling_demo(d: int, target_rank: int, adapter_rank: int, seed: int = 0):
    """The experiment that distinguishes a capacity ceiling from a data
    problem, run small enough to read.

    A target update of rank target_rank is planted, then fitted with an
    adapter of rank adapter_rank. When adapter_rank < target_rank the
    loss flattens at a positive floor and stays there -- not because the
    data is bad, not because the learning rate is wrong, but because the
    answer is outside the feasible set. In production this presents as a
    week of unhelpful sweeps.
    """
    rng = random.Random(seed)
    layer = LoRALinear(d, d, adapter_rank, alpha=adapter_rank, rng=rng)

    # Plant a rank-target_rank delta and generate data through W0 + it.
    left = [[rng.gauss(0.0, 0.4) for _ in range(target_rank)] for _ in range(d)]
    right = [[rng.gauss(0.0, 0.4) for _ in range(d)] for _ in range(target_rank)]

    samples = []
    for _ in range(64):
        x = [rng.gauss(0.0, 1.0) for _ in range(d)]
        base = matvec(layer.w0, x)
        delta = matvec(left, matvec(right, x))
        samples.append((x, [base[i] + delta[i] for i in range(d)]))

    history = train(layer, samples, epochs=400, lr=0.05)
    return {
        "final_loss": history[-1],
        "trainable": layer.trainable_parameters(),
        "frozen": layer.frozen_parameters(),
        "ratio": layer.frozen_parameters() / layer.trainable_parameters(),
    }

"""LoRA with the three silent failures made loud.

The literal version accepted all of these without complaint:

  * an adapter loaded against a base it was not trained on, which
    degrades quality and raises nothing;
  * an alpha that differs from the one the adapter was trained with,
    which applies a different update with no indication;
  * a merge left in place after an evaluation raised, so the next
    caller gets a model that is not the one they asked for.

Each becomes an exception or a context manager here. The fourth thing
that is added is diagnostic rather than defensive: effective_rank() and
retained_capability() are the two measurements that distinguish this
technique working from this technique appearing to work.
"""

from __future__ import annotations

import contextlib
import json
from collections.abc import Iterator, Sequence
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
import numpy.typing as npt

FloatArray = npt.NDArray[np.float32]


class AdapterError(Exception):
    """Base for every recoverable failure in adaptation."""


class InvalidAdapterConfig(AdapterError):
    """Rank, alpha or target modules cannot describe a usable adapter."""


class BaseModelMismatch(AdapterError):
    """The adapter was trained against a different base than this one.

    Includes quantization: an adapter trained against a four-bit base
    learned to compensate for that base's quantization error, so the
    same checkpoint in full precision is a different model as far as
    this adapter is concerned.
    """


class MergeNotSupported(AdapterError):
    """Merging would silently change the model that was evaluated."""


@dataclass(frozen=True)
class BaseFingerprint:
    """What an adapter is pinned to. All three parts, not just the name."""

    checkpoint_id: str
    quantization: str
    tokenizer_id: str

    def describe(self) -> str:
        return f"{self.checkpoint_id}[{self.quantization}]/{self.tokenizer_id}"


@dataclass(frozen=True)
class AdapterConfig:
    """Validated once, at construction, rather than trusted per call.

    Frozen because these values are part of the artefact's meaning:
    changing alpha after training does not adjust the adapter, it
    applies a different update with the same weights.
    """

    rank: int
    alpha: float
    target_modules: tuple[str, ...]
    dropout: float = 0.0

    def __post_init__(self) -> None:
        if self.rank <= 0:
            raise InvalidAdapterConfig(f"rank must be positive, got {self.rank}")
        if self.alpha <= 0.0:
            raise InvalidAdapterConfig(f"alpha must be positive, got {self.alpha}")
        if not self.target_modules:
            raise InvalidAdapterConfig(
                "no target modules: an adapter attached to nothing trains nothing"
            )
        if not 0.0 <= self.dropout < 1.0:
            raise InvalidAdapterConfig(f"dropout must be in [0, 1), got {self.dropout}")

    @property
    def scale(self) -> float:
        """alpha / r. Read from the config, never recomputed by a caller.

        Sweeping rank with alpha fixed sweeps this too, which is why a
        rank study that holds alpha constant measures two things and
        resolves neither.
        """
        return self.alpha / self.rank


@dataclass
class LoRAAdapter:
    """A frozen projection plus its rank-r update.

    The weight is a view onto the base model's array, never a copy. That
    is not only memory: a copy would let the adapter and the base drift
    apart, and an adapter whose frozen half has drifted is the hardest
    version of this bug to find.
    """

    name: str
    base_weight: FloatArray
    config: AdapterConfig
    fingerprint: BaseFingerprint
    down: FloatArray = field(init=False)
    up: FloatArray = field(init=False)
    _merged: bool = field(default=False, init=False)

    def __post_init__(self) -> None:
        out_features, in_features = self.base_weight.shape
        if self.config.rank > min(out_features, in_features):
            raise InvalidAdapterConfig(
                f"rank {self.config.rank} exceeds min dimension "
                f"{min(out_features, in_features)} of {self.name}; "
                "a full-rank update is just fine-tuning with extra steps"
            )

        rng = np.random.default_rng(abs(hash(self.name)) % (2**32))
        self.down = rng.normal(
            0.0, in_features**-0.5, size=(self.config.rank, in_features)
        ).astype(np.float32)
        # Zero, so the adapted model starts exactly at the pretrained
        # one. An untrained adapter is a provable no-op, not an
        # approximate one.
        self.up = np.zeros((out_features, self.config.rank), dtype=np.float32)

    @property
    def trainable_parameters(self) -> int:
        return int(self.down.size + self.up.size)

    @property
    def frozen_parameters(self) -> int:
        return int(self.base_weight.size)

    def delta(self) -> FloatArray:
        """The update alone, scaled. Materialized only for diagnostics."""
        return self.config.scale * (self.up @ self.down)

    def forward(self, activations: FloatArray) -> FloatArray:
        if activations.shape[-1] != self.base_weight.shape[1]:
            raise ValueError(
                f"{self.name}: expected {self.base_weight.shape[1]} input features, "
                f"got {activations.shape[-1]}"
            )

        base = activations @ self.base_weight.T
        if self._merged:
            return base
        return base + self.config.scale * (activations @ self.down.T) @ self.up.T

    @contextlib.contextmanager
    def merged(self) -> Iterator[None]:
        """Temporarily fold the update into the weight.

        A context manager rather than a merge()/unmerge() pair because
        the evaluation inside will eventually raise, and an unmerge that
        depends on an exception not happening is an unmerge that will be
        skipped exactly when it matters. On exit the base weight is
        restored bit-for-bit.
        """
        if self.fingerprint.quantization != "none":
            raise MergeNotSupported(
                f"{self.name}: base is {self.fingerprint.quantization}; the adapter "
                "was trained against that quantization error, so merging into a "
                "dequantized weight yields a different model than was evaluated"
            )

        original = self.base_weight.copy()
        self.base_weight += self.delta()
        self._merged = True
        try:
            yield
        finally:
            self.base_weight[...] = original
            self._merged = False


def effective_rank(adapter: LoRAAdapter, energy: float = 0.99) -> int:
    """Directions the update actually uses, from its spectrum.

    A rank-64 adapter whose delta has four meaningful singular values is
    reporting that the rank sweep was never run. Cheap to compute and
    almost never computed, which is why rank is usually a number someone
    copied from a blog post.
    """
    if not 0.0 < energy <= 1.0:
        raise ValueError(f"energy must be in (0, 1], got {energy}")

    singular_values = np.linalg.svd(adapter.delta(), compute_uv=False)
    total = float(singular_values.sum())
    if total == 0.0:
        return 0

    cumulative = np.cumsum(singular_values) / total
    return int(np.searchsorted(cumulative, energy) + 1)


def retained_capability(
    before: Sequence[float], after: Sequence[float], tolerance: float = 0.02
) -> tuple[bool, float]:
    """Did adaptation cost something the in-domain metric cannot see?

    Guard clauses rather than nesting, because the interesting path is
    the comparison and the uninteresting paths are all early exits.
    """
    if len(before) != len(after):
        raise ValueError(f"score counts differ: {len(before)} vs {len(after)}")
    if not before:
        raise ValueError("no held-out capability scores; forgetting is unmeasured")

    drop = float(np.mean(before) - np.mean(after))
    return drop <= tolerance, drop


class AdapterSet:
    """Adapters attached to one base, sharing one fingerprint."""

    def __init__(self, fingerprint: BaseFingerprint, config: AdapterConfig) -> None:
        self.fingerprint = fingerprint
        self.config = config
        self._adapters: dict[str, LoRAAdapter] = {}

    def attach(self, name: str, base_weight: FloatArray) -> LoRAAdapter:
        if name in self._adapters:
            raise InvalidAdapterConfig(f"{name} already has an adapter attached")
        if not any(pattern in name for pattern in self.config.target_modules):
            raise InvalidAdapterConfig(
                f"{name} matches none of {self.config.target_modules}; "
                "the target-module set moves quality more than rank does"
            )

        adapter = LoRAAdapter(
            name=name,
            base_weight=base_weight,
            config=self.config,
            fingerprint=self.fingerprint,
        )
        self._adapters[name] = adapter
        return adapter

    @property
    def trainable_parameters(self) -> int:
        return sum(item.trainable_parameters for item in self._adapters.values())

    @property
    def frozen_parameters(self) -> int:
        return sum(item.frozen_parameters for item in self._adapters.values())

    def save(self, directory: Path) -> None:
        """Weights and the scale that gives them meaning, together.

        Storing alpha and rank beside the tensors is the whole point.
        An adapter file without them is a set of numbers that can be
        applied at any magnitude, and it will be.
        """
        directory.mkdir(parents=True, exist_ok=True)
        manifest = {
            "rank": self.config.rank,
            "alpha": self.config.alpha,
            "target_modules": list(self.config.target_modules),
            "base": {
                "checkpoint_id": self.fingerprint.checkpoint_id,
                "quantization": self.fingerprint.quantization,
                "tokenizer_id": self.fingerprint.tokenizer_id,
            },
            "modules": sorted(self._adapters),
        }
        (directory / "adapter_config.json").write_text(json.dumps(manifest, indent=2))

        tensors = {}
        for name, adapter in self._adapters.items():
            tensors[f"{name}.lora_down"] = adapter.down
            tensors[f"{name}.lora_up"] = adapter.up
        np.savez(directory / "adapter_model.npz", **tensors)

    @classmethod
    def load(
        cls, directory: Path, serving_fingerprint: BaseFingerprint
    ) -> tuple[AdapterConfig, dict[str, FloatArray]]:
        """Refuse a mismatched base rather than degrade against it."""
        manifest = json.loads((directory / "adapter_config.json").read_text())
        trained_against = BaseFingerprint(**manifest["base"])

        if trained_against != serving_fingerprint:
            raise BaseModelMismatch(
                f"adapter trained against {trained_against.describe()}, "
                f"serving {serving_fingerprint.describe()}"
            )

        config = AdapterConfig(
            rank=manifest["rank"],
            alpha=manifest["alpha"],
            target_modules=tuple(manifest["target_modules"]),
        )
        with np.load(directory / "adapter_model.npz") as archive:
            tensors = {key: archive[key].astype(np.float32) for key in archive.files}
        return config, tensors

"""LoRA serving, shaped for the case the technique exists for.

One frozen base, many adapters, mixed in a single batch. The layout
follows from that: adapters stacked contiguously and selected by index,
the scale folded into the weights at load, the update accumulated into
the base output in place, and scratch allocated once.

What this does NOT change is worth stating first, because it is the
thing people expect it to change. The backward pass still walks the
full depth of the frozen model, so training throughput here is close to
full fine-tuning's. The saving was always memory.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import numpy.typing as npt

FloatArray = npt.NDArray[np.float32]
IndexArray = npt.NDArray[np.int32]


@dataclass(frozen=True)
class ServingShape:
    max_batch: int
    rank: int
    in_features: int
    out_features: int


class StackedAdapters:
    """Every tenant's adapter for one module, in two contiguous tensors.

    down is (n_adapters, rank, in_features) and up is
    (n_adapters, out_features, rank), both C-contiguous float32. One
    allocation each, so selecting an adapter is a strided read rather
    than a pointer chase through n_adapters separate arrays -- which is
    the difference between a gather that streams and one that stalls.

    The cost of this layout is that the set is immutable: adding a
    tenant reallocates both tensors. That is the right trade for a
    serving process that reloads on a schedule, and the wrong one for a
    system that wants to register adapters at runtime.
    """

    def __init__(self, shape: ServingShape, n_adapters: int) -> None:
        self.shape = shape
        self.down = np.zeros(
            (n_adapters, shape.rank, shape.in_features), dtype=np.float32
        )
        self.up = np.zeros(
            (n_adapters, shape.out_features, shape.rank), dtype=np.float32
        )

        # Scratch at the largest batch shape, sliced per call. A serving
        # loop that allocates per request spends a measurable share of a
        # cheap layer inside the allocator.
        self._projection = np.empty((shape.max_batch, shape.rank), dtype=np.float32)
        self._gathered_down = np.empty(
            (shape.max_batch, shape.rank, shape.in_features), dtype=np.float32
        )
        self._gathered_up = np.empty(
            (shape.max_batch, shape.out_features, shape.rank), dtype=np.float32
        )

    def load(self, slot: int, down: FloatArray, up: FloatArray, scale: float) -> None:
        """Fold alpha/r into up once, here, and never again.

        The scale is a constant of the adapter, so applying it per
        forward pass is a full out_features-wide multiply repeated for
        the life of the deployment. Folding it in at load costs one
        multiply per weight, total.

        It also means the stored tensors are no longer the trained ones,
        so the manifest's alpha must not be re-applied downstream --
        exactly the kind of bookkeeping that has to be written down.
        """
        np.copyto(self.down[slot], down)
        np.multiply(up, np.float32(scale), out=self.up[slot], casting="same_kind")

    def forward(
        self, activations: FloatArray, base_weight: FloatArray, slots: IndexArray
    ) -> FloatArray:
        """Mixed-adapter batch in two batched matmuls.

        activations is (batch, in_features); slots[i] names the adapter
        for row i. The base matmul is shared across the whole batch --
        one frozen weight, which is the entire operational argument for
        this technique -- and only the thin update varies per row.
        """
        batch = activations.shape[0]
        if batch > self.shape.max_batch:
            raise ValueError(
                f"batch {batch} exceeds scratch shape {self.shape.max_batch}"
            )

        # Shared across every tenant in the batch, at full BLAS
        # throughput. This is the dominant term and it is paid once.
        out = activations @ base_weight.T

        # Gather into preallocated scratch rather than fancy-indexing
        # into a fresh array, which would allocate per call.
        np.take(self.down, slots, axis=0, out=self._gathered_down[:batch])
        np.take(self.up, slots, axis=0, out=self._gathered_up[:batch])

        # (batch, 1, in) @ (batch, in, rank) -> (batch, 1, rank).
        # einsum keeps the shape contract readable; the arithmetic is
        # the same batched GEMM either way.
        projection = np.einsum(
            "bi,bri->br",
            activations,
            self._gathered_down[:batch],
            out=self._projection[:batch],
            optimize=True,
        )

        # Accumulate in place. The out_features-wide intermediate the
        # earlier stages returned is never allocated; at batch scale it
        # is the largest short-lived array in the layer.
        out += np.einsum("br,bor->bo", projection, self._gathered_up[:batch])
        return out


class HomogeneousBatch:
    """The faster path when a batch happens to share one adapter.

    Worth having separately because it is common -- a single-tenant
    deployment, or a scheduler that groups by adapter -- and because the
    gather disappears entirely: one (rank, in) matmul for the whole
    batch instead of a per-row slab. Grouping requests by adapter before
    serving turns the general path into this one, at the cost of holding
    requests until a group forms.
    """

    def __init__(self, down: FloatArray, up: FloatArray, scale: float) -> None:
        self.down = np.ascontiguousarray(down, dtype=np.float32)
        self.up = np.ascontiguousarray(up * np.float32(scale), dtype=np.float32)

    def forward(self, activations: FloatArray, base_weight: FloatArray) -> FloatArray:
        out = activations @ base_weight.T
        # Two matmuls, no gather, no per-row slab. Fused into out so the
        # intermediate never lands.
        out += (activations @ self.down.T) @ self.up.T
        return out


def group_by_adapter(slots: IndexArray) -> list[tuple[int, IndexArray]]:
    """Turn a mixed batch into homogeneous runs.

    A sort plus a boundary scan, vectorized rather than a Python loop
    over requests. Whether this pays depends on the tenant distribution:
    a batch of sixty-four distinct adapters gains nothing and has paid
    for a sort, while skewed traffic -- which is what real multi-tenant
    traffic looks like -- collapses to a handful of groups.
    """
    order = np.argsort(slots, kind="stable")
    ordered = slots[order]
    boundaries = np.flatnonzero(np.diff(ordered)) + 1
    runs = np.split(order, boundaries)
    return [(int(ordered[run[0]]), run.astype(np.int32)) for run in runs]


def memory_breakdown(
    n_parameters: int, trainable: int, bytes_per_parameter: int = 4
) -> dict[str, float]:
    """Where the memory actually went, in gigabytes.

    Included because the number people plan capacity from -- the
    trainable-parameter ratio -- is not the number that decides whether
    a run fits. Optimizer state scales with trainable parameters and
    collapses; frozen weights and activations do not move at all, and
    activations are usually what fails first.
    """
    gigabyte = 1024**3
    frozen = n_parameters * bytes_per_parameter / gigabyte
    gradients = trainable * bytes_per_parameter / gigabyte
    # Adam keeps two fp32 moments per trainable parameter. This is the
    # term the technique removes, and the only one.
    optimizer = trainable * 2 * 4 / gigabyte
    return {
        "frozen_weights_gb": frozen,
        "gradients_gb": gradients,
        "optimizer_state_gb": optimizer,
        "trainable_fraction": trainable / n_parameters,
        "note": "activation memory is unchanged and is excluded here on purpose",
    }