# 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=)`` (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 ``.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]}