| |
| """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") |
|
|
|
|
| |
|
|
| @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 |
| bev_pos: np.ndarray |
| self_posembed: PosEmbedMLP |
| cross_posembed: PosEmbedMLP |
| self_attn: MHAWeights |
| cross_attn: MHAWeights |
| norms: Tuple[Tuple[np.ndarray, np.ndarray, float], ...] |
| ffn: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] |
| 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") |
|
|
| |
| 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 |
|
|
|
|
| |
|
|
| 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 |
|
|
| |
| 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 |
|
|
| |
| @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} |
|
|
| |
| 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()} |
|
|
|
|
| |
|
|
| 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) |
| 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)} |
|
|
|
|
| |
|
|
| 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] |
| 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]} |
|
|
