attention.py

"""Which tokens see which on a backbone's full-attention layers, in training and serving alike.

| mode | a state token sees | a question token sees |
|---|---|---|
| `causal` | the tokens before it | the tokens before it |
| `state` | the whole state | the whole state, and the question before it |
| `prompt` | the whole prompt: state, question, options | the whole prompt |
| `set` | the state and every question | the state and its own question |

`state` reads the state both ways and never shows it the question, so one state's keys still serve every question
about it. `prompt` reads everything before the answer both ways, the state with its question in view: each question is
then a prompt of its own, with no prefill shared between the questions of one state. `set` keeps `prompt`'s reading
and packs a state's questions after it, `[state][q1]..[qN]`, each ending at its own answer slot: one pass, the state
read once. On these layers a question never attends to another; it hears of them only through the state, which reads
them all. A lone question reads exactly as under `prompt`.

Only full-attention layers change; recurrent, convolution and sliding-window layers stay causal. Under `set` they
run over the packed sequence, so a question also hears the ones before it through them; training packs the same way
(questions of a state in random order), so the served function is the trained one. A
checkpoint trains and serves with one mode, recorded in its config (`d1_attention`, with `d1_set_size`, the most
questions `set` packs after one state).
"""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass, field

import torch
import torch.nn.functional as F
from transformers import AttentionInterface, AttentionMaskInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import sdpa_mask

Tensor = torch.Tensor
SET_SIZE = 4  # the most questions `set` packs after one state when a checkpoint records none (`d1_set_size`)
TRAINED_SET_SIZE = 128  # the largest cap training takes (`--set-size`), so the most a checkpoint records
CHUNK = 4096  # queries a `Sees` mask is built for at once


def _fused(q: torch.Tensor) -> bool:
    """Whether the flash kernels take these queries: a GPU in half precision, ROCm or CUDA from Ampere on."""
    if not q.is_cuda or q.dtype not in (torch.float16, torch.bfloat16):
        return False
    return bool(torch.version.hip) or torch.cuda.get_device_capability(q.device)[0] >= 8


def _windowed(q: Tensor, k: Tensor, v: Tensor, window: int, scale: float) -> Tensor:
    """Causal attention within `window`, as local blocks: each block of `window`
    queries sees its own block and the one before it."""
    batch, heads, length, dim = q.shape
    blocks = -(-length // window)

    def split(t: Tensor) -> Tensor:  # (B, H, T, D) -> (B, blocks, H, window, D)
        t = F.pad(t, (0, 0, 0, blocks * window - length))
        return t.reshape(batch, t.shape[1], blocks, window, dim).transpose(1, 2)

    def with_previous(t: Tensor) -> Tensor:
        return torch.cat((F.pad(t, (0, 0, 0, 0, 0, 0, 1, 0))[:, :-1], t), dim=3)

    qb, kb, vb = split(q), with_previous(split(k)), with_previous(split(v))
    qpos = torch.arange(window, device=q.device)[:, None] + window
    kpos = torch.arange(2 * window, device=q.device)[None, :]
    block = torch.arange(blocks, device=q.device)[:, None, None]
    allowed = (kpos <= qpos) & (kpos > qpos - window) & ((block > 0) | (kpos >= window))  # block 0 has none before it
    mask = allowed[:, None].repeat(batch, 1, 1, 1)
    out = F.scaled_dot_product_attention(qb.flatten(0, 1), kb.flatten(0, 1), vb.flatten(0, 1),
                                         attn_mask=mask, scale=scale, enable_gqa=True)
    out = out.reshape(batch, blocks, heads, window, dim).transpose(1, 2)
    return out.reshape(batch, heads, blocks * window, dim)[:, :, :length]


def causal_attention(q: Tensor, k: Tensor, v: Tensor, window: int | None, scale: float) -> Tensor:
    """Right-padded chains: causal, within `window` if set: on the flash kernels their own window at any length (no
    branch on it, so torch.compile keeps one graph across the window), else blocks."""
    if window is None or not _fused(q) and q.shape[2] <= window:
        return F.scaled_dot_product_attention(q, k, v, is_causal=True, scale=scale, enable_gqa=True)
    if not _fused(q):
        return _windowed(q, k, v, window, scale)
    out = torch.ops.aten._flash_attention_forward(*(t.transpose(1, 2) for t in (q, k, v)), None, None, q.shape[2],
                                                  k.shape[2], 0.0, True, False, scale=scale,
                                                  window_size_left=window - 1, window_size_right=0)[0]
    return out.transpose(1, 2)


def both_ways(mode: str, tokenizer, ids: list[int], suffix: str) -> int | None:
    """How many of a prompt's first tokens read each other both ways: none (`causal`), its state, or all of it
    (`prompt`, and `set` for a lone question)."""
    if mode == "causal":
        return None
    return len(ids) if mode in ("prompt", "set") else state_length(tokenizer, ids, suffix)


def state_length(tokenizer, ids: list[int], suffix: str) -> int:
    """A prompt's tokens before its question: the prompt is the state's text then `suffix`, which opens on a token
    boundary (the state ends on a blank line), so its own tokens are the prompt's last ones."""
    return len(ids) - len(tokenizer.encode(suffix, add_special_tokens=False))


def pack_set(tokenizer, prefix: str, suffixes: list[str],
             header: str = "", slot_offsets: list[list[int]] | None = None,
             first_ids: list[int] | None = None,
             ) -> tuple[list[int], int, list[int], list[list[int]]]:
    """`set`'s sequence for one state, in any prompt layout: `prefix` (the state, ending with the question
    `header`) and the first question tokenized as its lone prompt is, so one question is exactly its `prompt`
    row; each further question its own group, `header` then its suffix. Returns the ids, the state's length, each
    token's question (1..N, -1 on the state) and each question's answer positions. `slot_offsets` names slots
    relative to its suffix's last token (0); without it, each question has that one slot."""
    if not suffixes:
        raise ValueError("a question set cannot be empty")
    offsets = [[0] for _ in suffixes] if slot_offsets is None else slot_offsets
    if len(offsets) != len(suffixes):
        raise ValueError("one slot layout is required per question")

    def positions(k: int, length: int, end: int) -> list[int]:
        own = offsets[k]
        if (not own or own != sorted(set(own)) or own[-1] != 0
                or any(isinstance(o, bool) or not isinstance(o, int) or not -length < o <= 0 for o in own)):
            raise ValueError("slots must be increasing positions inside their question, ending at its last token")
        return [end + o for o in own]

    ids = tokenizer.encode(prefix + suffixes[0], add_special_tokens=False) if first_ids is None else list(first_ids)
    state = state_length(tokenizer, ids, suffixes[0])
    question = [-1] * state + [1] * (len(ids) - state)
    ends = [positions(0, len(ids) - state, len(ids) - 1)]
    for k, suffix in enumerate(suffixes[1:], 2):
        own = tokenizer.encode(header + suffix, add_special_tokens=False)
        ids, question = ids + own, question + [k] * len(own)
        ends.append(positions(k - 1, len(own), len(ids) - 1))
    return ids, state, question, ends


def text_backbone(model):
    """The module that returns hidden states, under any wrapper (DDP, a training wrapper)."""
    m = getattr(model, "module", model)
    if hasattr(m, "get_base_model"):
        m = m.get_base_model()
    return getattr(m, "model", m)


def language_model(model):
    """The language model under any wrapper and a vision-language model."""
    m = text_backbone(model)
    return getattr(m, "language_model", m)


def _runs(*parts: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, int]:
    """Segments, segment i the runs `(start[i], length[i])` of every part back to back: their flat indices, their
    cumulative lengths (int32, as the varlen kernels take them) and the longest."""
    lengths = sum(n for _, n in parts)
    cu = F.pad(lengths.cumsum(0), (1, 0))
    seg = torch.repeat_interleave(torch.arange(len(lengths), device=lengths.device), lengths)
    off = torch.arange(len(seg), device=seg.device) - cu[seg]
    idx = torch.zeros_like(off)
    for start, n in parts:
        idx = torch.where((off >= 0) & (off < n[seg]), start[seg] + off, idx)
        off = off - n[seg]
    return idx, cu.int(), int(lengths.max()) if len(lengths) else 0


def _varlen(q: Tensor, k: Tensor, v: Tensor, cu_q: Tensor, cu_k: Tensor, longest: Tensor, causal: bool,
            scale: float) -> tuple[Tensor, Tensor, Tensor, Tensor]:
    """Flash attention over packed segments: the output [N, H, D], then what its backward reads (the log-sum-exp
    [H, N], the RNG state and its companion). Causal aligns each segment's last query to its last key. `longest`: the
    longest query and key segment, on the CPU. Compiled, the opaque op `d1::varlen`, so the lengths neither break the
    graph nor recompile it; its backward is flash's own (`_varlen_backward`), as eager's is."""
    max_q, max_k = longest.tolist()
    out, lse, rng, unused, _ = torch.ops.aten._flash_attention_forward(q, k, v, cu_q, cu_k, max_q, max_k, 0.0, causal,
                                                                       False, scale=scale)
    return out, lse, rng, unused


def _varlen_backward(ctx, grad, *_):
    q, k, v, out, lse, rng, unused, cu_q, cu_k = ctx.saved_tensors
    max_q, max_k = ctx.longest
    dq, dk, dv = torch.ops.aten._flash_attention_backward(grad.contiguous(), q, k, v, out, lse, cu_q, cu_k, max_q,
                                                          max_k, 0.0, ctx.causal, rng, unused, scale=ctx.scale)
    return dq, dk, dv, None, None, None, None, None


def _varlen_context(ctx, inputs, output):
    q, k, v, cu_q, cu_k, longest, causal, scale = inputs
    ctx.save_for_backward(q, k, v, *output, cu_q, cu_k)
    ctx.longest, ctx.causal, ctx.scale = longest.tolist(), causal, scale


if not hasattr(torch.ops.d1, "varlen"):  # once a process: two repos' copies of this module share the op
    _op = torch.library.custom_op("d1::varlen", _varlen, mutates_args=())
    _op.register_fake(lambda q, k, v, cu_q, cu_k, longest, causal, scale: (
        torch.empty_like(q), q.new_empty(q.shape[1], q.shape[0], dtype=torch.float32),
        q.new_empty((), dtype=torch.int64), q.new_empty((), dtype=torch.int64)))
    _op.register_autograd(_varlen_backward, setup_context=_varlen_context)


@dataclass
class Sees:
    """The full-attention mask of `state_attention`, never built whole (a [B, 1, T, T] tensor takes 33 GB on a
    180k-token row): `attend` takes the attention under it, `rows` a block of queries' part of it."""

    state: torch.Tensor  # [B] each row's state length
    question: torch.Tensor | None  # [B, T] `set`: each token's question, 1..N, and -1 elsewhere
    real: torch.Tensor  # [B, T] the tokens that are not padding
    _calls: dict = field(default_factory=dict, init=False, repr=False, compare=False)  # once a device, every layer's

    def __post_init__(self):
        """On a GPU, the plan on the mask's device, made before any layer: a compiled forward reads it with no graph
        break. Off a GPU, `attend` masks and never plans."""
        if self.real.is_cuda:
            self._calls[self.real.device] = self._plan(self.real.device)

    def rows(self, start: int, width: int, device) -> torch.Tensor:
        """Which keys queries `start` to `start + width` see, [B, 1, width, T]."""
        q = torch.arange(start, start + width, device=device)[None, :, None]
        kv = torch.arange(self.real.shape[1], device=device)[None, None, :]
        s = self.state.to(device)[:, None, None]
        if self.question is None:  # causal, and the state both ways
            m = (kv <= q) | ((q < s) & (kv < s))
        else:  # the state sees every token and every token the state; a question token its own question
            g = self.question.to(device)
            m = (q < s) | (kv < s) | (g[:, start:start + width, None] == g[:, None, :])
        return (m & self.real.to(device)[:, None, :])[:, None]

    def attend(self, q: Tensor, k: Tensor, v: Tensor, scale: float, window: int | None = None) -> Tensor:
        """Attention under this mask; q [B, H, T, D], k and v [B, Hkv, T, D]. On the flash kernels, no mask: the
        states over their keys in one varlen call, the questions over their state and themselves in another. Off
        them (CPU, fp32), `CHUNK` masked queries at a time. A sliding layer's `window`: causal within it, as every
        mode keeps it (right-padded rows: a pad never reaches a real token)."""
        if window is not None:
            return causal_attention(q, k, v, window, scale)
        b, h, t, d = q.shape
        if not _fused(q):
            starts = range(0, t, CHUNK) if t > CHUNK else (0,)  # a range over t would compile a graph per length
            return torch.cat([F.scaled_dot_product_attention(q[:, :, s:s + CHUNK], k, v, scale=scale, enable_gqa=True,
                                                             attn_mask=self.rows(s, min(CHUNK, t - s), q.device))
                              for s in starts], dim=2)
        if q.device not in self._calls:  # a model over GPUs: once on each
            self._calls[q.device] = self._plan(q.device)
        if self._calls[q.device] is None:  # every row reads all of itself, unpadded: plain attention
            return F.scaled_dot_product_attention(q, k, v, scale=scale, enable_gqa=True)
        flat = [x.transpose(1, 2).reshape(b * t, x.shape[1], d) for x in (q, k, v)]
        out = flat[0].new_zeros(b * t, h, d)
        varlen = torch.ops.d1.varlen if torch.compiler.is_compiling() else _varlen
        for qi, cu_q, ki, cu_k, longest, causal in self._calls[q.device]:
            attended = varlen(flat[0][qi], flat[1][ki], flat[2][ki], cu_q, cu_k, longest, causal, scale)[0]
            out = out.index_copy(0, qi, attended)
        return out.view(b, t, h, d).transpose(1, 2)

    def _plan(self, device) -> list | None:
        """The varlen calls, as (queries, their offsets, keys, theirs, the longest of each, causal): each state over
        its row (`set`) or itself (`state`, `prompt`); each question over its state then itself, causally under
        `state`, where the question part is one run a row. Under `set` a row of one question is one segment: all of
        it sees all of it. None when every row is one such segment with no padding."""
        real = self.real.to(device)
        b, t = real.shape
        n = real.sum(1)
        s = torch.minimum(self.state.to(device), n)
        base = torch.arange(b, device=device) * t
        tail = (real & (torch.arange(t, device=device) >= s[:, None])).flatten().nonzero().squeeze(1)
        g = (torch.zeros(b, t, dtype=torch.long, device=device) if self.question is None
             else self.question.to(device)).flatten()[tail]
        new = F.pad((g[1:] != g[:-1]) | (tail[1:] != tail[:-1] + 1), (1, 0), value=True)[:len(tail)]
        first = new.nonzero().squeeze(1)
        start, length = tail[first], torch.diff(F.pad(first, (0, 1), value=len(tail)))
        if self.question is not None:
            lone = torch.bincount(start // t, minlength=b) == 1
            s = torch.where(lone, n, s)
            start, length = start[~lone[start // t]], length[~lone[start // t]]
        if not len(start) and bool((s == t).all()):
            return None
        calls = [(_runs((base, s)), _runs((base, n if self.question is not None else s)), False)]
        if len(start):
            row = start // t
            calls.append((_runs((start, length)), _runs((base[row], s[row]), (start, length)), self.question is None))
        return [(qi, cu_q, ki, cu_k, torch.tensor([max_q, max_k]), causal)
                for (qi, cu_q, max_q), (ki, cu_k, max_k), causal in calls]


def d1_attention(module, q, k, v, attention_mask, scaling=None, sliding_window=None, **kwargs):
    """transformers' attention function for `attn_implementation="d1"` (registered on import): a `Sees` mask attends
    through `Sees.attend` (a sliding layer within its window), anything else is SDPA's."""
    if isinstance(attention_mask, Sees):
        return attention_mask.attend(q, k, v, scaling, sliding_window).transpose(1, 2), None
    return sdpa_attention_forward(module, q, k, v, attention_mask, scaling=scaling, sliding_window=sliding_window,
                                  **kwargs)


AttentionInterface.register("d1", d1_attention)
# the masks transformers builds itself under "d1" (no `state_attention`; a tower, an audio encoder) are SDPA's: an
# implementation without one gets None for every mask, which drops a sliding window and the padding
AttentionMaskInterface.register("d1", sdpa_mask)
ATTENTION_KINDS = ("full_attention", "sliding_attention")  # the layer kinds a `Sees` masks; the others take padding


def mask_mapping(layer_types, sees: Sees, pad: Tensor) -> dict:
    """The per-kind masks of a language model under `Sees`, keyed as its `layer_types`: the `Sees` on attention
    layers, the 2D padding (None without any) on convolution and linear-attention layers."""
    padding = None if pad.all() else pad
    return {kind: sees if kind in ATTENTION_KINDS else padding for kind in set(layer_types)}


@contextmanager
def state_attention(model, state_len: torch.Tensor | None, question: torch.Tensor | None = None):
    """Within the block, the model's language model reads each row's first `state_len` tokens both ways; with
    `question` (`set`: each token's question, 1..N, and -1 elsewhere), as `Sees` says. The mask is swapped in at the
    language model, after the multimodal wrapper has placed its media: as the mask of a language model that reads it
    (`reads_sees`), else as transformers' per-kind masks (`mask_mapping`), which `d1_attention` reads."""
    if state_len is None:
        yield
        return

    def hook(module, args, kwargs):
        pad, x = kwargs.get("attention_mask"), kwargs.get("input_ids")
        if isinstance(pad, dict):
            raise ValueError("state attention replaces a 2D padding mask, not a prepared mapping")
        if pad is None:
            x = kwargs["inputs_embeds"] if x is None else x
            pad = torch.ones(x.shape[:2], dtype=torch.long, device=x.device)
        sees = Sees(state_len.to(pad.device), question, pad.bool())
        if getattr(module, "reads_sees", False):
            kwargs["attention_mask"] = sees
        elif module.config._attn_implementation == "d1":
            kinds = getattr(module.config, "layer_types", None) or ("full_attention",)  # a config naming none: all full
            kwargs["attention_mask"] = mask_mapping(kinds, sees, pad)
        else:
            raise ValueError(f"{type(module).__name__} reads `state_attention` with attn_implementation='d1'")
        return args, kwargs

    handle = language_model(model).register_forward_pre_hook(hook, with_kwargs=True)
    try:
        yield
    finally:
        handle.remove()