code/tt_diffusion_planner/ttaw/models/transfusion_head.py

# SPDX-License-Identifier: Apache-2.0
"""C26: the TransFusion query head on the device (PLAN.md section 1.2 row C26; TransFusion builds it, BEVFusion and
PTv3 reuse it; probes P4, P7, P8, P13).

The dense part of a TransFusion head (shared conv, heatmap convs) is C17 / C24 work; this module is what follows the
local-max heatmap (C22): query selection (C21), query initialisation, ONE post-norm decoder layer and the prediction
heads, as one traceable device graph with one packed output. The model-agnostic weights record
:class:`QueryHeadWeights` (host numpy, BN folded) describes it:

- ``class_table`` ``(C, d)``: the class encoding (``Conv1d(one_hot(c))`` = ``W[:, c] + b``, an exact row gather);
- ``bev_pos`` ``(cells, 2)``: the model's constant BEV position of every feature cell (TF / PT ``(x + 0.5, y + 0.5)``
  with ``k = y * W + x``; BF ``(i + 0.5, j + 0.5)`` with ``k = i * 180 + j``); a query's position is ``bev_pos[pos]``;
- ``self_posembed`` / ``cross_posembed``: the position MLPs ``Conv1d 2 -> d + ReLU, Conv1d d -> d`` (BN folded);
- the decoder layer: self-attention on ``s = x + pe(query_pos)`` (Q = K = V = s), add + LayerNorm, cross-attention
  of ``x + pe(query_pos)`` over ``key = lidar_feat + pe_k(bev_pos)`` (K = V = key), add + LN, FFN (ReLU), add + LN;
- the prediction heads ``(name, Conv1d d -> h + ReLU, Conv1d h -> c)``.

Device form (every rewrite exact in real arithmetic; the oracle :class:`QueryHeadOracle` runs the literal forms):

| step | device ops |
|---|---|
| selection | C21 ``TopKSelect``: class-major bf16 row of the local-max heatmap, ``topk_large_indices`` k = ``padded_k(K)``, exact fp32 decode -> uint32 ``pos`` / ``cls`` rows (or given indices: teacher forcing) |
| query init | ``ttnn.embedding`` of the ``lidar_feat`` rows (bf16) by ``pos`` + the class-table row by ``cls``; the query position embedding is a function of ``pos`` only, so it is a constant **QPE table** ``self_posembed(bev_pos)`` (host float64 -> bf16) gathered by ``pos``: ``query_pos`` itself never passes through bf16 (S:transfusion:241) |
| keys | ``key = lidar_feat + KPE`` with the constant ``KPE = cross_posembed(bev_pos)`` table (the 0.6 GMAC position branch leaves the per-frame graph, S:transfusion:279) |
| attention | C20 ``sdpa`` (``is_causal=False``, ``scale`` from the weights, chunk table), the head dim zero-padded to 32 in the projection weights (``pad_head_columns`` / ``pad_head_rows``, exact); self-attention on the logical ``K`` queries with no mask, Q / K / V from one fused projection; cross-attention over every cell |
| LayerNorm | ``ttnn.layer_norm(o, residual_input_tensor=x, epsilon=<the model's>)`` (never ttnn's 1e-12 default) |
| heads | merged: ``relu(x @ [Wh_1 ... Wh_n] + bh)`` then the block-diagonal ``@ Wo + bo`` (padded to 32), fp32 out |
| query heat | ``ttnn.embedding`` of the local-max heatmap rows (classes zero-padded to 32) by ``pos``: every class at the query's cell (the TF / mmdeploy score ``query_heat * sigmoid(heatmap)``; BF / PT one-hot it on the host) |

Output (``__call__``): ``{"heads": fp32 [1, 1, K, 32], "query_heat": bf16 [1, 1, Kp, 32], "indices": fp32 [1, 1, 1,
Kp]}`` (``Kp = padded_k(K)``; the first ``K`` indices are the proposals, fp32 holds them exactly), ready for
``ttaw.trace.pack_outputs``. :func:`assemble_transfusion` turns the read-back arrays into the TF / mmdeploy graph
outputs ``cls_score0``, ``bbox_pred0``, ``dir_cls_pred0`` (``center + query_pos`` in float32 on the host).

Limits: one decoder layer (the TF, BF and PT heads have one; a second layer would re-embed the predicted centres,
which no constant table can do); ``bev_pos`` must be finite. Precision (``ttaw.precision`` policy, module names
``<prefix>.self_attn``, ``.cross_attn``, ``.ffn``, ``.heads``, ``.norm``): fp32 weights, HiFi4 + fp32 accumulation,
bf16 activations (SDPA takes bf16); the tables are bf16 (their rows are model activations, rounded once).

numpy only at import; ttnn / torch are imported inside the functions that need them.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple

import numpy as np

from ..ops.attention import pad_head_columns, pad_head_rows, padded_head_dim
from ..ops.topk import TopKSelect, padded_k
from ..precision import Precision, PrecisionPolicy
from .second import RowLinear, free_maps

__all__ = ["MHAWeights", "PosEmbedMLP", "PredictionHead", "QueryHeadWeights", "TransFusionQueryHead",
           "QueryHeadOracle", "assemble_transfusion", "DEFAULT_PRECISION", "default_policy", "HEAD_PAD"]

HEAD_PAD = 32
DEFAULT_PRECISION = "HiFi4+fp32:w=fp32"


def default_policy() -> PrecisionPolicy:
    """fp32 weights, HiFi4 + fp32 accumulation, bf16 activations for every module of the head."""
    return PrecisionPolicy(default=DEFAULT_PRECISION)


def _f32(a: Any) -> np.ndarray:
    """A writable float32 copy (ONNX-derived arrays may be read-only; torch warns about those)."""
    return np.array(a, dtype=np.float32, copy=True, order="C")


# ------------------------------------------------------------------------------------------------- weights

@dataclass(frozen=True)
class MHAWeights:
    """``nn.MultiheadAttention`` as exported, ``[in, out]`` layout: ``q = x_q @ wq + bq`` (likewise k / v), per head
    ``softmax(scale * q k^T) v``, ``out = concat_heads @ wo + bo``."""

    wq: np.ndarray
    bq: np.ndarray
    wk: np.ndarray
    bk: np.ndarray
    wv: np.ndarray
    bv: np.ndarray
    wo: np.ndarray
    bo: np.ndarray
    num_heads: int
    scale: float

    @property
    def head_dim(self) -> int:
        return int(np.shape(self.wq)[1]) // int(self.num_heads)

    @property
    def dim(self) -> int:
        return int(np.shape(self.wq)[0])


@dataclass(frozen=True)
class PosEmbedMLP:
    """``Conv1d 2 -> d`` (BN folded) + ReLU, ``Conv1d d -> d``; ``[in, out]`` weights."""

    w0: np.ndarray
    b0: np.ndarray
    w1: np.ndarray
    b1: np.ndarray

    def numpy(self, pos: Any) -> np.ndarray:
        """``(P, 2)`` positions -> ``(P, d)`` float64."""
        p = np.asarray(pos, np.float64)
        h = np.maximum(p @ np.asarray(self.w0, np.float64) + np.asarray(self.b0, np.float64), 0.0)
        return h @ np.asarray(self.w1, np.float64) + np.asarray(self.b1, np.float64)


@dataclass(frozen=True)
class PredictionHead:
    """``Conv1d d -> h`` + ReLU, ``Conv1d h -> c`` (``[in, out]`` weights)."""

    name: str
    w_hidden: np.ndarray
    b_hidden: np.ndarray
    w_out: np.ndarray
    b_out: np.ndarray

    @property
    def channels(self) -> int:
        return int(np.shape(self.w_out)[1])


@dataclass(frozen=True)
class QueryHeadWeights:
    """Everything the query head needs, model-agnostic (module docstring). ``height`` / ``width``: the feature grid
    (``cells = height * width`` keys, row-major ``k = y * width + x`` as the dense maps are laid out)."""

    height: int
    width: int
    num_classes: int
    num_proposals: int
    class_table: np.ndarray                       # (C, d)
    bev_pos: np.ndarray                           # (cells, 2)
    self_posembed: PosEmbedMLP
    cross_posembed: PosEmbedMLP
    self_attn: MHAWeights
    cross_attn: MHAWeights
    norms: Tuple[Tuple[np.ndarray, np.ndarray, float], ...]   # (gamma, beta, eps) x 3: after self, cross, FFN
    ffn: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]  # w1 (d, f), b1, w2 (f, d), b2
    heads: Tuple[PredictionHead, ...]

    @property
    def cells(self) -> int:
        return int(self.height) * int(self.width)

    @property
    def dim(self) -> int:
        return int(np.shape(self.class_table)[1])

    def validate(self) -> None:
        d, c, cells = self.dim, int(self.num_classes), self.cells
        if np.shape(self.class_table) != (c, d):
            raise ValueError(f"class_table {np.shape(self.class_table)} != ({c}, {d})")
        if np.shape(self.bev_pos) != (cells, 2) or not np.all(np.isfinite(self.bev_pos)):
            raise ValueError(f"bev_pos must be finite ({cells}, 2), got {np.shape(self.bev_pos)}")
        for a in (self.self_attn, self.cross_attn):
            if a.dim != d or np.shape(a.wo) != (a.num_heads * a.head_dim, d):
                raise ValueError("attention weights do not match the model width")
            if not (np.isfinite(a.scale) and a.scale > 0):
                raise ValueError(f"attention scale {a.scale}")
        if len(self.norms) != 3:
            raise ValueError("a decoder layer has three LayerNorms")
        if np.shape(self.ffn[0])[0] != d or np.shape(self.ffn[2])[1] != d:
            raise ValueError("FFN weights do not match the model width")
        if not self.heads or any(np.shape(h.w_hidden)[0] != d for h in self.heads):
            raise ValueError("prediction heads do not read the model width")
        if sum(h.channels for h in self.heads) > HEAD_PAD:
            raise ValueError(f"the merged head output is padded to {HEAD_PAD} channels")
        if c > HEAD_PAD:
            raise ValueError(f"the query heat rows hold at most {HEAD_PAD} classes")
        padded_k(self.num_proposals)
        if self.num_proposals > cells * c:
            raise ValueError("more proposals than heatmap entries")

    # ---- host tables (float64, rounded once by the device upload) ------------------------------------------
    def query_pos_table(self) -> np.ndarray:
        """``self_posembed(bev_pos)`` ``(cells, d)``: the query position embedding of every cell."""
        return self.self_posembed.numpy(self.bev_pos)

    def key_pos_table(self) -> np.ndarray:
        """``cross_posembed(bev_pos)`` ``(cells, d)``: the input-independent key position embedding."""
        return self.cross_posembed.numpy(self.bev_pos)

    def merged_heads(self, pad_to: int = HEAD_PAD):
        """``(Wh, bh, Wo, bo, slices)``: ``relu(x @ Wh + bh) @ Wo + bo`` with the block-diagonal ``Wo`` padded to
        ``pad_to`` columns; ``slices`` = ``[(name, start, stop)]`` of each head's output channels (exact rewrite)."""
        wh = np.concatenate([_f32(h.w_hidden) for h in self.heads], axis=1)
        bh = np.concatenate([_f32(h.b_hidden) for h in self.heads])
        rows = wh.shape[1]
        cols = sum(h.channels for h in self.heads)
        width = max(pad_to, -(-cols // pad_to) * pad_to)
        wo = np.zeros((rows, width), np.float32)
        bo = np.zeros(width, np.float32)
        slices, r, c = [], 0, 0
        for h in self.heads:
            hr, hc = np.shape(h.w_hidden)[1], h.channels
            wo[r:r + hr, c:c + hc] = _f32(h.w_out)
            bo[c:c + hc] = _f32(h.b_out)
            slices.append((h.name, c, c + hc))
            r, c = r + hr, c + hc
        return wh, bh, wo, bo, slices


# ---------------------------------------------------------------------------------------------- device

class _LayerNorm:
    def __init__(self, gamma: Any, beta: Any, eps: float, precision: Precision, name: str):
        self.gamma, self.beta, self.eps = _f32(gamma).reshape(1, -1).copy(), _f32(beta).reshape(1, -1).copy(), float(eps)
        self.precision = precision
        self.name = name
        self._dev: Optional[Tuple[Any, Any, Any]] = None

    def upload(self, device: Any):
        import torch
        import ttnn

        def up(a):
            return ttnn.from_torch(torch.from_numpy(a), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=device,
                                   memory_config=ttnn.DRAM_MEMORY_CONFIG)

        self._dev = (up(self.gamma), up(self.beta), self.precision.compute_kernel_config())
        return self._dev

    def __call__(self, x: Any, residual: Any):
        """``LayerNorm(x + residual)`` (one ``ttnn.layer_norm`` with a fused residual add)."""
        import ttnn

        g, b, cfg = self._dev or self.upload(x.device())
        return ttnn.layer_norm(x, residual_input_tensor=residual, weight=g, bias=b, epsilon=self.eps,
                               compute_kernel_config=cfg, memory_config=ttnn.DRAM_MEMORY_CONFIG)

    def release(self) -> None:
        if self._dev is not None:
            free_maps(self._dev[0], self._dev[1])
            self._dev = None


class TransFusionQueryHead:
    """The query head of :class:`QueryHeadWeights` on the device (module docstring). Device weights and tables are
    uploaded by :meth:`prepare` (or the first call), outside any capture; the forward is traceable."""

    def __init__(self, weights: QueryHeadWeights, *, policy: Optional[PrecisionPolicy] = None, prefix: str = "head",
                 attn_fp32_acc: bool = False, name: str = "query_head"):
        weights.validate()
        self.weights = weights
        self.prefix = prefix
        self.name = name
        self.policy = policy or default_policy()
        self.attn_fp32_acc = bool(attn_fp32_acc)
        self.dim = weights.dim
        self.cells = weights.cells
        self.classes = int(weights.num_classes)
        self.k = int(weights.num_proposals)
        self.select = TopKSelect(self.cells, self.classes, self.k, name=f"{prefix}.topk")
        self.kp = self.select.kp
        sa, ca = weights.self_attn, weights.cross_attn
        self.heads_self, self.hd_self = int(sa.num_heads), sa.head_dim
        self.heads_cross, self.hd_cross = int(ca.num_heads), ca.head_dim
        self.scale_self, self.scale_cross = float(sa.scale), float(ca.scale)

        def prec(module: str) -> Precision:
            return self.policy.resolve(f"{prefix}.{module}")

        def cols(w, b, a: MHAWeights):
            return pad_head_columns(_f32(w), _f32(b), a.num_heads, a.head_dim)

        p_sa, p_ca, p_ffn, p_heads = prec("self_attn"), prec("cross_attn"), prec("ffn"), prec("heads")
        q, bq = cols(sa.wq, sa.bq, sa)
        k, bk = cols(sa.wk, sa.bk, sa)
        v, bv = cols(sa.wv, sa.bv, sa)
        self.self_qkv = RowLinear(np.concatenate([q, k, v], axis=1), np.concatenate([bq, bk, bv]), precision=p_sa,
                                  name=f"{prefix}.self_attn.qkv")
        self.self_out = RowLinear(pad_head_rows(_f32(sa.wo), sa.num_heads, sa.head_dim), _f32(sa.bo),
                                  precision=p_sa, name=f"{prefix}.self_attn.out")
        self.cross_q = RowLinear(*cols(ca.wq, ca.bq, ca), precision=p_ca, name=f"{prefix}.cross_attn.q")
        self.cross_k = RowLinear(*cols(ca.wk, ca.bk, ca), precision=p_ca, name=f"{prefix}.cross_attn.k")
        self.cross_v = RowLinear(*cols(ca.wv, ca.bv, ca), precision=p_ca, name=f"{prefix}.cross_attn.v")
        self.cross_out = RowLinear(pad_head_rows(_f32(ca.wo), ca.num_heads, ca.head_dim), _f32(ca.bo),
                                   precision=p_ca, name=f"{prefix}.cross_attn.out")
        w1, b1, w2, b2 = weights.ffn
        self.ffn1 = RowLinear(_f32(w1), _f32(b1), activation="relu", precision=p_ffn, name=f"{prefix}.ffn.0")
        self.ffn2 = RowLinear(_f32(w2), _f32(b2), precision=p_ffn, name=f"{prefix}.ffn.1")
        wh, bh, wo, bo, self.head_slices = weights.merged_heads()
        self.heads_hidden = RowLinear(wh, bh, activation="relu", precision=p_heads, name=f"{prefix}.heads.hidden")
        self.heads_out = RowLinear(wo, bo, precision=p_heads, output_dtype="float32", name=f"{prefix}.heads.out")
        p_norm = prec("norm")
        self.norms = [_LayerNorm(g, b, eps, p_norm, f"{prefix}.norm{i}") for i, (g, b, eps) in enumerate(weights.norms)]
        self._tables: Optional[Dict[str, Any]] = None

    # ---- device constants --------------------------------------------------------------------------------
    def prepare(self, device: Any) -> None:
        """Upload the tables (bf16: class table and QPE as ROW_MAJOR gather tables, KPE as a TILE map) and the
        LayerNorm rows. The linears upload their weights on their first call (the warm-up)."""
        import torch
        import ttnn

        if self._tables is not None:
            return
        w = self.weights

        def up(a: np.ndarray, layout):
            t = torch.from_numpy(_f32(a).reshape(1, 1, a.shape[0], a.shape[1]))
            return ttnn.from_torch(t, dtype=ttnn.bfloat16, layout=layout, device=device,
                                   memory_config=ttnn.DRAM_MEMORY_CONFIG)

        self._tables = {"class": up(_f32(w.class_table), ttnn.ROW_MAJOR_LAYOUT),
                        "qpe": up(w.query_pos_table(), ttnn.ROW_MAJOR_LAYOUT),
                        "kpe": up(w.key_pos_table(), ttnn.TILE_LAYOUT)}
        for n in self.norms:
            n.upload(device)

    def _t(self, device: Any) -> Dict[str, Any]:
        if self._tables is None:
            self.prepare(device)
        return self._tables  # type: ignore[return-value]

    # ---- forward pieces ------------------------------------------------------------------------------------
    @staticmethod
    def _rows(t: Any, n: int):
        """``[1, n, D]`` / ``[1, 1, n, D]`` -> ``[1, 1, n, D]`` (a view)."""
        import ttnn

        return ttnn.reshape(t, (1, 1, int(n), int(t.shape[-1])))

    def _gather(self, index: Any, table: Any):
        import ttnn

        return self._rows(ttnn.embedding(index, table, layout=ttnn.TILE_LAYOUT), int(index.shape[-1]))

    def init_queries(self, lidar_feat: Any, pos: Any, cls: Any) -> Tuple[Any, Any]:
        """``lidar_feat`` TILE ``[1, 1, cells, d]``, ``pos`` / ``cls`` uint32 ROW_MAJOR ``[1, Kp]`` -> ``(x, qpe)``
        TILE ``[1, 1, K, d]``: the initial queries ``lidar_feat[pos] + class_table[cls]`` and their position
        embeddings ``QPE[pos]``."""
        import ttnn

        t = self._t(lidar_feat.device())
        table = ttnn.to_layout(lidar_feat, ttnn.ROW_MAJOR_LAYOUT)
        feat = self._gather(pos, table)
        free_maps(table)
        cemb = self._gather(cls, t["class"])
        x_p = ttnn.add(feat, cemb)
        free_maps(feat, cemb)
        qpe_p = self._gather(pos, t["qpe"])
        if self.kp == self.k:
            return x_p, qpe_p
        x = ttnn.slice(x_p, [0, 0, 0, 0], [1, 1, self.k, self.dim])
        qpe = ttnn.slice(qpe_p, [0, 0, 0, 0], [1, 1, self.k, self.dim])
        free_maps(x_p, qpe_p)
        return x, qpe

    def keys(self, lidar_feat: Any):
        """``lidar_feat + KPE`` (TILE ``[1, 1, cells, d]``): the cross-attention key / value input."""
        import ttnn

        return ttnn.add(lidar_feat, self._t(lidar_feat.device())["kpe"])

    def _sdpa(self, q, k, v, scale: float):
        from ..ops.attention import sdpa

        return sdpa(q, k, v, scale=scale, fp32_acc=self.attn_fp32_acc, concat_heads=True)

    def decoder(self, x: Any, qpe: Any, key: Any, taps: Optional[Dict[str, Any]] = None):
        """One post-norm decoder layer on ``x`` ``[1, 1, K, d]`` -> ``[1, 1, K, d]`` (``key``: :meth:`keys`)."""
        import ttnn

        from ..ops.attention import split_heads, split_qkv

        s = ttnn.add(x, qpe)
        qkv = self.self_qkv(s)
        free_maps(s)
        q, k, v = split_qkv(qkv, self.heads_self)
        free_maps(qkv)
        a = self._sdpa(q, k, v, self.scale_self)
        free_maps(q, k, v)
        o = self.self_out(a)
        free_maps(a)
        x1 = self.norms[0](o, x)
        free_maps(o)
        if taps is not None:
            taps["decoder.self_attn"] = x1
        qc_in = ttnn.add(x1, qpe)
        qc = self.cross_q(qc_in)
        free_maps(qc_in)
        kc, vc = self.cross_k(key), self.cross_v(key)
        qh, kh, vh = (split_heads(t, self.heads_cross) for t in (qc, kc, vc))
        free_maps(qc, kc, vc)
        a = self._sdpa(qh, kh, vh, self.scale_cross)
        free_maps(qh, kh, vh)
        o = self.cross_out(a)
        free_maps(a)
        x2 = self.norms[1](o, x1)
        free_maps(o)
        if taps is not None:
            taps["decoder.cross_attn"] = x2
        else:
            free_maps(x1)
        h = self.ffn1(x2)
        f = self.ffn2(h)
        free_maps(h)
        x3 = self.norms[2](f, x2)
        free_maps(f)
        if taps is None:
            free_maps(x2)
        return x3

    def predict(self, x: Any):
        """The merged prediction heads: ``[1, 1, K, d]`` -> fp32 ``[1, 1, K, 32]`` (``head_slices`` channels)."""
        hid = self.heads_hidden(x)
        out = self.heads_out(hid)
        free_maps(hid)
        return out

    def query_heat(self, heat_nms: Any, pos: Any):
        """The local-max heatmap rows (every class, zero-padded to 32) at the proposals' cells: bf16 TILE
        ``[1, 1, Kp, 32]``."""
        import ttnn

        padded = ttnn.pad(heat_nms, [(0, 0), (0, 0), (0, 0), (0, HEAD_PAD - self.classes)], 0.0)
        table = ttnn.to_layout(padded, ttnn.ROW_MAJOR_LAYOUT)
        free_maps(padded)
        out = self._gather(pos, table)
        free_maps(table)
        return out

    def _indices_tile(self, indices: Any):
        """uint32 ROW_MAJOR ``[1, Kp]`` -> fp32 ``[1, 1, 1, Kp]`` ROW_MAJOR-ready tensor (exact below 2**24)."""
        import ttnn

        t = ttnn.to_layout(indices, ttnn.TILE_LAYOUT)
        f = ttnn.typecast(t, ttnn.float32)
        free_maps(t)
        return ttnn.reshape(f, (1, 1, 1, int(indices.shape[-1])))

    def __call__(self, lidar_feat: Any, heat_nms: Any, indices: Optional[Any] = None,
                 taps: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """``lidar_feat`` TILE ``[1, 1, cells, d]`` (bf16), ``heat_nms`` TILE ``[1, 1, cells, C]`` bf16 (C22's
        output), ``indices``: None (device top-k) or a uint32 ROW_MAJOR ``[1, Kp]`` input (teacher forcing). Returns
        ``{"heads", "query_heat", "indices"}`` (module docstring); ``taps`` (a dict) also collects ``query.feat``,
        ``decoder.query_pos_embed``, ``decoder.self_attn``, ``decoder.cross_attn``, ``decoder.out``."""
        idx = self.select.indices(heat_nms) if indices is None else indices
        pos, cls = self.select.decode(idx)
        x, qpe = self.init_queries(lidar_feat, pos, cls)
        free_maps(cls)
        key = self.keys(lidar_feat)
        y = self.decoder(x, qpe, key, taps)
        free_maps(key)
        heads = self.predict(y)
        qheat = self.query_heat(heat_nms, pos)
        free_maps(pos)
        idx_f = self._indices_tile(idx)
        if indices is None:
            free_maps(idx)
        if taps is not None:
            taps.update({"query.feat": x, "decoder.query_pos_embed": qpe, "decoder.out": y})
        else:
            free_maps(x, qpe, y)
        return {"heads": heads, "query_heat": qheat, "indices": idx_f}

    # ---- host side -----------------------------------------------------------------------------------------
    def unpack_heads(self, heads: Any) -> Dict[str, np.ndarray]:
        """The read-back ``heads`` array -> ``{name: (c, K)}`` float32."""
        a = np.asarray(heads, np.float32).reshape(-1, HEAD_PAD)[: self.k]
        return {name: np.ascontiguousarray(a[:, s:e].T) for name, s, e in self.head_slices}

    def release(self) -> None:
        for m in (self.self_qkv, self.self_out, self.cross_q, self.cross_k, self.cross_v, self.cross_out,
                  self.ffn1, self.ffn2, self.heads_hidden, self.heads_out, *self.norms):
            m.release()
        if self._tables is not None:
            free_maps(*self._tables.values())
            self._tables = None

    def describe(self) -> Dict[str, Any]:
        return {"name": self.name, "cells": self.cells, "classes": self.classes, "proposals": self.k,
                "padded_k": self.kp, "dim": self.dim, "heads": [self.heads_self, self.heads_cross],
                "head_dim": [self.hd_self, self.hd_cross],
                "padded_head_dim": [padded_head_dim(self.hd_self), padded_head_dim(self.hd_cross)],
                "scale": [self.scale_self, self.scale_cross], "attn_fp32_acc": self.attn_fp32_acc,
                "head_slices": [list(s) for s in self.head_slices], "precision": self.policy.describe(),
                "topk": self.select.describe()}


# ---------------------------------------------------------------------------------------------- oracle

class QueryHeadOracle:
    """fp32 (or fp64) torch oracle of the LITERAL head (no tables, no padding, no merged heads): the reference of
    the device tests. Inputs NCHW ``lidar_feat (1, d, H, W)``, ``heat_nms (1, C, H, W)`` and the class-major
    indices of the proposals."""

    def __init__(self, weights: QueryHeadWeights, *, dtype: str = "float32"):
        weights.validate()
        self.w = weights
        self.dtype = dtype

    def _t(self, a: Any):
        import torch

        return torch.as_tensor(np.asarray(a, np.float64)).to(getattr(torch, self.dtype))

    def _mha(self, a: MHAWeights, q_in, kv_in, v_in):
        import torch

        t = self._t
        q = q_in @ t(a.wq) + t(a.bq)
        k = kv_in @ t(a.wk) + t(a.bk)
        v = v_in @ t(a.wv) + t(a.bv)
        h, d = int(a.num_heads), a.head_dim
        q = q.reshape(-1, h, d).transpose(0, 1)
        k = k.reshape(-1, h, d).transpose(0, 1)
        v = v.reshape(-1, h, d).transpose(0, 1)
        o = torch.softmax((q @ k.transpose(1, 2)) * float(a.scale), dim=-1) @ v
        return o.transpose(0, 1).reshape(q_in.shape[0], h * d) @ t(a.wo) + t(a.bo)

    def _ln(self, i: int, x):
        import torch

        g, b, eps = self.w.norms[i]
        return torch.nn.functional.layer_norm(x, (x.shape[-1],), self._t(g), self._t(b), float(eps))

    def _posembed(self, mlp: PosEmbedMLP, pos):
        import torch

        t = self._t
        return torch.relu(pos @ t(mlp.w0) + t(mlp.b0)) @ t(mlp.w1) + t(mlp.b1)

    def __call__(self, lidar_feat: Any, heat_nms: Any, indices: Any) -> Dict[str, np.ndarray]:
        """-> ``{"query.feat" (K, d), "decoder.query_pos_embed", "decoder.self_attn", "decoder.cross_attn",
        "decoder.out" (K, d), "heads" {name: (c, K)}, "query_heat" (C, K), "query_pos" (K, 2)}``."""
        import torch

        w = self.w
        idx = np.asarray(indices).astype(np.int64).reshape(-1)[: w.num_proposals]
        pos, cls = idx % w.cells, idx // w.cells
        with torch.no_grad():
            feat = self._t(np.asarray(lidar_feat, np.float64).reshape(w.dim, w.cells).T)      # (cells, d)
            heat = np.asarray(heat_nms, np.float64).reshape(w.num_classes, w.cells)
            x = feat[pos] + self._t(w.class_table)[cls]
            qpos = self._t(np.asarray(w.bev_pos, np.float64)[pos])
            qpe = self._posembed(w.self_posembed, qpos)
            key = feat + self._posembed(w.cross_posembed, self._t(w.bev_pos))
            s = x + qpe
            x1 = self._ln(0, x + self._mha(w.self_attn, s, s, s))
            x2 = self._ln(1, x1 + self._mha(w.cross_attn, x1 + qpe, key, key))
            w1, b1, w2, b2 = (self._t(a) for a in w.ffn)
            x3 = self._ln(2, x2 + torch.relu(x2 @ w1 + b1) @ w2 + b2)
            heads = {}
            for h in w.heads:
                y = torch.relu(x3 @ self._t(h.w_hidden) + self._t(h.b_hidden)) @ self._t(h.w_out) + self._t(h.b_out)
                heads[h.name] = y.T.numpy().astype(np.float32)

        def f(a):
            return a.numpy().astype(np.float32)

        return {"query.feat": f(x), "decoder.query_pos_embed": f(qpe), "decoder.self_attn": f(x1),
                "decoder.cross_attn": f(x2), "decoder.out": f(x3), "heads": heads,
                "query_heat": heat[:, pos].astype(np.float32), "query_pos": f(qpos)}


# -------------------------------------------------------------------------------------------- assembly

def assemble_transfusion(heads: Dict[str, np.ndarray], query_heat: Any, indices: Any, bev_pos: Any, *,
                         num_classes: int, num_proposals: int, cells: int) -> Dict[str, np.ndarray]:
    """TransFusion's (mmdeploy-patched) graph outputs from the head's read-back arrays, in float32 on the host:
    ``cls_score0 = query_heat * sigmoid(heads.heatmap)`` ``(1, C, K)``, ``bbox_pred0 = [center + query_pos,
    height, dim, vel]`` ``(1, 8, K)``, ``dir_cls_pred0 = rot`` ``(1, 2, K)``; ``query_pos = bev_pos[pos]``
    (float32: never through bf16, S:transfusion:241). ``heads``: :meth:`TransFusionQueryHead.unpack_heads`;
    ``query_heat``: the ``[.., Kp, 32]`` rows; ``indices``: the class-major proposal indices."""
    k = int(num_proposals)
    idx = np.rint(np.asarray(indices, np.float64).reshape(-1)[:k]).astype(np.int64)
    pos = idx % int(cells)
    qpos = np.asarray(bev_pos, np.float32)[pos]                                       # (K, 2)
    qheat = np.asarray(query_heat, np.float32).reshape(-1, HEAD_PAD)[:k, :int(num_classes)].T
    logits = np.asarray(heads["heatmap"], np.float32)
    sig = (1.0 / (1.0 + np.exp(-logits.astype(np.float64)))).astype(np.float32)
    cls = (qheat * sig).astype(np.float32)[None]
    center = (np.asarray(heads["center"], np.float32) + qpos.T).astype(np.float32)
    bbox = np.concatenate([center, heads["height"], heads["dim"], heads["vel"]], axis=0).astype(np.float32)[None]
    return {"cls_score0": cls, "bbox_pred0": bbox, "dir_cls_pred0": np.asarray(heads["rot"], np.float32)[None]}