| """Attention Residuals (AttnRes) — arXiv:2603.15031 (Moonshot / Kimi). |
| |
| Full AttnRes: replace uniform residual sum with depth-wise softmax over prior |
| branch outputs. Each destination step has one learnable pseudo-query w ∈ R^d: |
| |
| α_i = softmax_i( w^T RMSNorm(v_i) ) |
| h = Σ_i α_i v_i |
| |
| We treat each Attention and MLP as a separate depth step (paper §2). |
| """ |
| from __future__ import annotations |
|
|
| from typing import List |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import Tensor, nn |
|
|
|
|
| def _as_local(t: Tensor) -> Tensor: |
| """Localize DTensor params/acts for depth-mix arithmetic.""" |
| to_local = getattr(t, "to_local", None) |
| if callable(to_local): |
| return to_local() |
| return t |
|
|
|
|
| def _match_ref(out: Tensor, ref: Tensor) -> Tensor: |
| """Preserve DTensor wrapper of `ref` when depth-mix ran on local tensors.""" |
| if type(out) is type(ref): |
| return out |
| from_local = getattr(type(ref), "from_local", None) |
| if from_local is None or not hasattr(ref, "device_mesh"): |
| return out |
| try: |
| return from_local(out, device_mesh=ref.device_mesh, placements=ref.placements) |
| except Exception: |
| return out |
|
|
|
|
| class AttnResBank(nn.Module): |
| """Reusable value bank + depth-wise queries for Full AttnRes.""" |
|
|
| def __init__(self, hidden_size: int, max_slots: int = 256, eps: float = 1e-6) -> None: |
| super().__init__() |
| self.hidden_size = int(hidden_size) |
| self.max_slots = int(max_slots) |
| self.eps = float(eps) |
| |
| self.queries = nn.Parameter(torch.empty(self.max_slots, self.hidden_size)) |
| nn.init.normal_(self.queries, mean=0.0, std=0.02) |
| self._values: List[Tensor] = [] |
|
|
| def clear(self) -> None: |
| self._values = [] |
|
|
| def snapshot(self) -> List[Tensor]: |
| """Detached copies so a later segment can mix over prior loops as constants.""" |
| return [v.detach() for v in self._values] |
|
|
| def restore(self, values: List[Tensor]) -> None: |
| self._values = list(values) |
|
|
| def detach_inplace(self) -> None: |
| self._values = [v.detach() for v in self._values] |
|
|
| @property |
| def size(self) -> int: |
| return len(self._values) |
|
|
| def append(self, v: Tensor) -> None: |
| if len(self._values) >= self.max_slots: |
| raise RuntimeError( |
| f"AttnResBank overflow: {len(self._values)} >= max_slots={self.max_slots}" |
| ) |
| self._values.append(v) |
|
|
| def mix(self, *, next_layer: bool = False) -> Tensor: |
| """Attend over current values. |
| |
| next_layer=False: input to the next branch (|V|=n → query slot n-1). |
| next_layer=True: readout after last append (|V|=n → query slot n). |
| """ |
| n = len(self._values) |
| if n == 0: |
| raise RuntimeError("AttnResBank.mix called with empty values") |
| if n == 1 and not next_layer: |
| return self._values[0] |
| slot = n if next_layer else (n - 1) |
| slot = min(max(int(slot), 0), self.max_slots - 1) |
| ref = self._values[-1] |
| locals_v = [_as_local(v) for v in self._values] |
| V = torch.stack(locals_v, dim=0) |
| K = F.rms_norm(V, (V.shape[-1],), eps=self.eps) |
| w = _as_local(self.queries[slot]) |
| |
| logits = torch.einsum("d,...d->...", w, K) |
| alpha = torch.softmax(logits, dim=0) |
| out = (alpha.unsqueeze(-1) * V).sum(dim=0) |
| return _match_ref(out, ref) |
|
|
