tt-model push bevfusion-p150 (container)

7cc5047 verified 19 days ago

20.5 kB

# SPDX-License-Identifier: Apache-2.0
"""fp32 PyTorch re-implementation of the deployed BEVFusion-L network, module by module, with taps.
Built from :mod:`.weights` (the ONNX initializers addressed by node; every BatchNorm folded in float64 where the TT
port folds it), it reproduces ``bevfusion_lidar.onnx`` (S:bevfusion:205-264) from the voxel features on (the VFE mean
is host pre-processing in this port, computed in float32 exactly as the graph does by ``host.preprocess.vfe_mean``):
- ``sparse`` (the 21 ``autoware::ImplicitGemm`` layers): every layer in the **gather form** the device uses (C13 /
  C28): ``y[o] = sum_k W'[:, k, :] @ x[nmap[o, k]] + b'`` over the neighbour maps of ``host.preprocess.SparsePlan``
  (``-1``: no input), ReLU, the basic blocks' residual before their second ReLU; BN folded (``W' = W * s``,
  ``b' = beta - mean * s``). The plan's maps are spconv-faithful (13-query + symmetric write); the independent
  oracle is the research reference's scatter form (``tests/test_reference_cpu.py``);
- ``to-dense``: ``dense[0, x, y, z, :] = conv_out[o]`` on a zero ``[1, 180, 180, 2, 128]`` map, ``Transpose(0, 4, 3,
  1, 2)``, ``Reshape [1, 256, 180, 180]``: channel ``c * 2 + z``, rows x, columns y;
- ``backbone``: SECOND, 2 blocks of 6 3x3 convs + ReLU (first stride 2 in block 1); ``neck``: SECONDFPN deblock0 1x1
  + ReLU, deblock1 ConvTranspose 2x2 s2 + ReLU, concat (1, 512, 180, 180);
- ``head`` dense part: shared conv 3x3 512 -> 128 (``lidar_feat``, no activation), heatmap head (3x3 + ReLU, 3x3 -> 5
  logits), ``heat = sigmoid``, local max ``heat * (heat == local_max)`` with ``local_max`` = MaxPool 3x3 s1 p0 of
  classes 0-3 placed in a zero map (their border cells never survive) and class 4 = the heat itself (passes);
- selection: the top ``K = 500`` of the flattened (5 x 32400) local-max heatmap as ONNX ``TopK`` defines it (largest
  first, **ties to the lower index**: a stable descending sort); ``topk_indices=`` replaces it by given indices
  (teacher forcing, PLAN.md 0.3 item 1: the device top-k breaks ties in no defined order, probe P8);
- query init: ``class = idx // 32400``, ``pos = idx % 32400``; ``query = lidar_feat[:, pos] + class_table[class]``
  (the Conv1d of a one-hot is the table row ``W[:, c] + b``); ``query_pos = (pos // 180 + 0.5, pos % 180 + 0.5)``
  computed arithmetically in float32 (bit-equal to the graph's ``bev_pos`` gather, checked at load);
  ``query_heat = heat_nms[:, :, pos]`` (all 5 classes);
- decoder (one layer, post-norm, 8 heads x 16): self-attention on ``x + pe(query_pos)`` (Q = K = V), cross-attention
  of ``x + pe(query_pos)`` over ``lidar_feat^T + pe(bev_pos)`` (32,400 keys), FFN 128 -> 256 -> 128, LayerNorm eps
  1e-5 after each; ``Q`` divided by ``sqrt(16) = 4`` before ``Q K^T`` (the graph's ``Div``);
- prediction heads (Conv1d 128 -> 64 + BN folded + ReLU, Conv1d 64 -> C: center 2, height 1, dim 3, rot 2, vel 2,
  heatmap 5);
- outputs: ``bbox_pred = [center + query_pos, height, dim, rot, vel]`` (10, K), ``score = sigmoid(heatmap)[label] *
  query_heat[label]`` (K,) (the graph's one-hot product + ReduceMax, equal because both factors are >= 0),
  ``label_pred = class`` (K,) int64.
Taps (``ttaw.golden.TapRegistry``; stable names, the keys of the device gates; dense taps keep the ONNX tensors'
batch axis and layout; sparse taps are ``(N_level, C)`` rows in the plan's level order): :data:`MODULE_TAPS` plus the
opt-in per-layer / large taps :data:`EXTRA_TAPS`. Values are float32 numpy (``head.topk_indices`` and ``label_pred``
int64). torch is imported lazily; the module runs where ttnn is absent.
"""
from __future__ import annotations

from typing import Any, Dict, Optional

import numpy as np

from ..ttaw.golden import NULL_TAPS, TapRegistry
from .config import BEV_SHAPE, SPARSE_LEVELS
from .weights import ENCODER_PLAN, BEVFusionWeights

__all__ = ["BEVFusionTorch", "MODULE_TAPS", "EXTRA_TAPS", "SELECTION_TAPS", "SPARSE_TAPS", "DENSE_TAPS",
           "OUTPUT_NAMES", "stable_topk"]

OUTPUT_NAMES = ("bbox_pred", "score", "label_pred")
SPARSE_TAPS = ("vfe",) + tuple(name for name, *_ in ENCODER_PLAN if not name.endswith(".conv1")) + ("sparse.bev",)
DENSE_TAPS = ("backbone.block0", "backbone.block1", "neck.deblock0", "neck.deblock1", "head.shared", "head.heatmap",
              "head.heat", "head.heat_nms")
# taps downstream of the data-dependent top-k (compared set-based or teacher-forced, PLAN.md 0.3 item 1)
SELECTION_TAPS = ("head.topk_values", "head.topk_indices", "query.feat", "query.pos", "query.heat",
                  "decoder.query_pos_embed", "decoder.self_attn", "decoder.cross_attn", "decoder.out",
                  "heads.center", "heads.height", "heads.dim", "heads.rot", "heads.vel", "heads.heatmap") + OUTPUT_NAMES
MODULE_TAPS = SPARSE_TAPS + DENSE_TAPS + SELECTION_TAPS
# opt-in taps: per layer, the 66 MB neck concat and the input-independent key position embedding
EXTRA_TAPS = (("neck", "head.heatmap.hidden", "decoder.key_pos_embed")
              + tuple(name for name, *_ in ENCODER_PLAN if name.endswith(".conv1"))
              + tuple(f"backbone.block{b}.conv{j}" for b in range(2) for j in range(6)))


def _torch():
    import torch

    return torch


def _tap(taps: TapRegistry, name: str, value):
    """Record a torch tensor as a numpy copy (float32, or int64 for integer tensors) and return it unchanged; never
    through ``ttaw.tensors.to_numpy`` (which would import ttnn)."""
    if taps.enabled and taps.wants(name):
        a = value.detach().cpu().numpy()
        taps.tap(name, a.astype(np.int64 if a.dtype.kind in "iu" else np.float32, copy=True))
    return value


def stable_topk(scores, k: int):
    """ONNX ``TopK(largest=1, sorted=1)`` on the last axis: values descending, equal values in ascending index order
    (a stable descending sort). Returns ``(values, indices)``."""
    torch = _torch()
    order = torch.sort(scores, dim=-1, descending=True, stable=True)
    return order.values[..., :k], order.indices[..., :k]


class BEVFusionTorch:
    """The fp32 reference network. ``threads`` caps torch's intra-op threads (shared host: <= 4)."""

    def __init__(self, weights: BEVFusionWeights, *, threads: Optional[int] = None, device: str = "cpu"):
        torch = _torch()
        if threads is not None:
            torch.set_num_threads(int(threads))
        self.weights = weights
        self.device = device
        t = self._t
        self._sparse = [(layer, t(layer.weight.reshape(layer.out_channels, -1, layer.in_channels)), t(layer.bias))
                        for layer in weights.encoder]
        self._layers: Dict[str, tuple] = {}
        for layer in weights.dense_layers():
            self._layers[layer.name] = (layer, t(layer.weight), t(layer.bias))
        h = weights.head
        d = h.decoder
        for layer in (*d.self_posembed, *d.cross_posembed) + tuple(x for _, a, b in h.prediction_heads for x in (a, b)):
            self._layers[layer.name] = (layer, t(layer.weight), t(layer.bias))
        self.class_table = t(h.class_table)                     # (5, 128)
        self.bev_pos = t(h.bev_pos)                             # (32400, 2)
        self.K = int(h.num_proposals)
        self.bev_h, self.bev_w = BEV_SHAPE
        self.pooled = tuple(h.pooled_classes)
        self._attn = {"self_attn": self._attn_tensors(d.self_attn), "cross_attn": self._attn_tensors(d.cross_attn)}
        self._norms = [(t(g), t(b), float(eps)) for g, b, eps in d.norms]
        self._ffn = (t(d.ffn_w1), t(d.ffn_b1), t(d.ffn_w2), t(d.ffn_b2))
        self._kpe = None                                        # cross_posembed(bev_pos): input-independent, cached

    def _t(self, a: Any):
        torch = _torch()
        return torch.from_numpy(np.array(a, dtype=np.float32, copy=True)).to(self.device)

    def _attn_tensors(self, a) -> dict:
        t = self._t
        return {"wq": t(a.wq), "bq": t(a.bq), "wk": t(a.wk), "bk": t(a.bk), "wv": t(a.wv), "bv": t(a.bv),
                "wo": t(a.wo), "bo": t(a.bo), "heads": a.num_heads, "hd": a.head_dim, "div": float(a.q_divisor)}

    # ---------------------------------------------------------------------------------------- sparse encoder
    def sparse_conv(self, x, nmap: np.ndarray, w, b):
        """Gather-form sparse conv: ``x`` (N_in, Cin) torch, ``nmap`` (N_out, K) int (-1 = none), ``w`` (Cout, K,
        Cin) -> (N_out, Cout) (bias added, no activation). Taps accumulate in tap order, fp32."""
        torch = _torch()
        n_in, cin = x.shape
        xp = torch.cat([x, torch.zeros((1, cin), dtype=x.dtype, device=x.device)], 0)    # explicit zero row
        idx = torch.from_numpy(np.where(nmap >= 0, nmap, n_in).astype(np.int64)).to(x.device)
        out = torch.zeros((idx.shape[0], w.shape[0]), dtype=x.dtype, device=x.device)
        for k in range(idx.shape[1]):
            col = idx[:, k]
            if bool((col < n_in).any()):
                out += xp[col] @ w[:, k, :].T
        return out + b

    def sparse_encoder(self, vfe: np.ndarray, plan: Any, taps: TapRegistry = NULL_TAPS):
        """(V, 4) voxel features + ``host.preprocess.SparsePlan`` -> the BEV (1, 256, 180, 180)."""
        torch = _torch()
        x = _tap(taps, "vfe", torch.as_tensor(np.asarray(vfe, dtype=np.float32), device=self.device))
        if x.shape[0] != plan.levels["l1"].num_active:
            raise ValueError(f"{x.shape[0]} voxel features for {plan.levels['l1'].num_active} active voxels")
        outs: Dict[str, Any] = {}
        for layer, w, b in self._sparse:
            y = self.sparse_conv(x, plan.maps[layer.map], w, b)
            if layer.residual is not None:
                y = y + outs[layer.residual]
            x = _tap(taps, layer.name, torch.relu(y))
            outs[layer.name] = x
        lvl = plan.levels["out"]
        X, Y, Z = lvl.spatial_shape
        if (X, Y) != BEV_SHAPE or lvl.spatial_shape != SPARSE_LEVELS["out"]:
            raise ValueError(f"conv_out grid {lvl.spatial_shape}")
        dense = torch.zeros((1, X, Y, Z, x.shape[1]), dtype=x.dtype, device=x.device)
        oc = torch.from_numpy(lvl.coords).to(x.device)
        dense[0, oc[:, 0], oc[:, 1], oc[:, 2]] = x
        bev = dense.permute(0, 4, 3, 1, 2).reshape(1, x.shape[1] * Z, X, Y)
        return _tap(taps, "sparse.bev", bev)

    # ----------------------------------------------------------------------------------------- dense layers
    def layer(self, name: str, x):
        torch = _torch()
        F = torch.nn.functional
        layer, w, b = self._layers[name]
        if layer.kind == "conv":
            y = F.conv2d(x, w, b, stride=layer.stride, padding=layer.padding)
        elif layer.kind == "conv_transpose":
            y = F.conv_transpose2d(x, w, b, stride=layer.stride, padding=layer.padding)
        else:
            y = F.conv1d(x, w, b, stride=layer.stride, padding=layer.padding)
        return torch.relu(y) if layer.relu else y

    def backbone(self, bev, taps: TapRegistry = NULL_TAPS):
        outs, x = [], bev
        for i, block in enumerate(self.weights.blocks):
            for layer in block:
                x = _tap(taps, layer.name, self.layer(layer.name, x))
            outs.append(_tap(taps, f"backbone.block{i}", x))
        return outs

    def neck(self, blocks, taps: TapRegistry = NULL_TAPS):
        torch = _torch()
        ups = [_tap(taps, layer.name, self.layer(layer.name, x)) for layer, x in zip(self.weights.deblocks, blocks)]
        return _tap(taps, "neck", torch.cat(ups, dim=1))

    def dense_head(self, neck, taps: TapRegistry = NULL_TAPS):
        """-> (lidar_feat (1, 128, 180, 180), heat_nms (1, 5, 180, 180))."""
        torch = _torch()
        F = torch.nn.functional
        lidar = _tap(taps, "head.shared", self.layer("head.shared", neck))
        hid = _tap(taps, "head.heatmap.hidden", self.layer("head.heatmap.hidden", lidar))
        logits = _tap(taps, "head.heatmap", self.layer("head.heatmap", hid))
        heat = _tap(taps, "head.heat", torch.sigmoid(logits))
        k = int(self.weights.head.nms_kernel)
        pad = k // 2
        pooled = list(self.pooled)
        local_max = torch.zeros_like(heat)
        local_max[:, pooled, pad:-pad, pad:-pad] = F.max_pool2d(heat[:, pooled], k, 1, 0)
        rest = [c for c in range(heat.shape[1]) if c not in self.pooled]
        local_max[:, rest] = heat[:, rest]
        heat_nms = _tap(taps, "head.heat_nms", heat * (heat == local_max).to(heat.dtype))
        return lidar, heat_nms

    def select(self, heat_nms, topk_indices: Optional[Any] = None, taps: TapRegistry = NULL_TAPS):
        """Top-K of the flattened local-max heatmap -> (values (1, K), indices (1, K) int64); given ``topk_indices``
        (teacher forcing) are used as they are and their heat values gathered."""
        torch = _torch()
        flat = heat_nms.reshape(1, -1)
        if topk_indices is None:
            values, idx = stable_topk(flat, self.K)
        else:
            idx = torch.as_tensor(np.asarray(topk_indices, dtype=np.int64).reshape(1, -1), device=self.device)
            if idx.shape[1] != self.K or int(idx.min()) < 0 or int(idx.max()) >= flat.shape[1]:
                raise ValueError(f"topk_indices must be {self.K} indices in [0, {flat.shape[1]})")
            values = flat.gather(1, idx)
        return _tap(taps, "head.topk_values", values), _tap(taps, "head.topk_indices", idx)

    def query_init(self, lidar, heat_nms, idx, taps: TapRegistry = NULL_TAPS):
        """-> (query (1, 128, K), query_pos (1, K, 2), query_heat (1, 5, K), class (K,))."""
        torch = _torch()
        n_cells = lidar.shape[2] * lidar.shape[3]
        cls_id = torch.div(idx, n_cells, rounding_mode="floor")
        pos = idx - cls_id * n_cells
        flat = lidar.reshape(1, lidar.shape[1], -1)
        q = flat.gather(2, pos[:, None, :].expand(-1, flat.shape[1], -1)) + self.class_table[cls_id[0]].T[None]
        w = self.bev_w
        qpos = torch.stack([torch.div(pos, w, rounding_mode="floor").to(torch.float32) + 0.5,
                            (pos % w).to(torch.float32) + 0.5], dim=-1)
        heat_flat = heat_nms.reshape(1, heat_nms.shape[1], -1)
        qheat = heat_flat.gather(2, pos[:, None, :].expand(-1, heat_flat.shape[1], -1))
        return (_tap(taps, "query.feat", q), _tap(taps, "query.pos", qpos), _tap(taps, "query.heat", qheat),
                cls_id[0])

    def _posembed(self, which: str, pos):
        """(P, 2) -> (P, 128): Conv1d 2 -> 128 + ReLU (BN folded), Conv1d 128 -> 128 (1x1 convs over P positions)."""
        x = pos.T[None]                                                        # (1, 2, P)
        x = self.layer(f"decoder.{which}.0", x)
        return self.layer(f"decoder.{which}.1", x)[0].T                         # (P, 128)

    def key_pos_embed(self, taps: TapRegistry = NULL_TAPS):
        """``cross_posembed(bev_pos)`` (32400, 128): input-independent, computed once."""
        torch = _torch()
        if self._kpe is None:
            with torch.no_grad():
                self._kpe = self._posembed("cross_posembed", self.bev_pos)
        return _tap(taps, "decoder.key_pos_embed", self._kpe[None])[0]

    def _mha(self, which: str, q_in, kv_in, v_in=None):
        """Multi-head attention as exported: projections, per-head ``softmax((q / div) @ k^T) @ v``, out proj."""
        torch = _torch()
        a = self._attn[which]
        v_in = kv_in if v_in is None else v_in
        q = q_in @ a["wq"] + a["bq"]
        k = kv_in @ a["wk"] + a["bk"]
        v = v_in @ a["wv"] + a["bv"]
        H, D = a["heads"], a["hd"]
        lq, lk = q.shape[0], k.shape[0]
        q = q.reshape(lq, H, D).transpose(0, 1) / a["div"]
        k = k.reshape(lk, H, D).permute(1, 2, 0)
        v = v.reshape(lk, H, D).transpose(0, 1)
        o = torch.softmax(q @ k, dim=-1) @ v                                   # (H, Lq, D)
        o = o.transpose(0, 1).reshape(lq, H * D)
        return o @ a["wo"] + a["bo"]

    def _ln(self, i: int, x):
        torch = _torch()
        g, b, eps = self._norms[i]
        return torch.nn.functional.layer_norm(x, (x.shape[-1],), g, b, eps)

    def decoder(self, query, qpos, lidar, taps: TapRegistry = NULL_TAPS):
        """One post-norm decoder layer -> (K, 128)."""
        torch = _torch()
        qpe = _tap(taps, "decoder.query_pos_embed", self._posembed("self_posembed", qpos[0])[None])[0]
        kpe = self.key_pos_embed(taps)
        x = query[0].T                                                         # (K, 128)
        key = lidar.reshape(lidar.shape[1], -1).T + kpe                        # (32400, 128)
        s = x + qpe
        x = _tap(taps, "decoder.self_attn", self._ln(0, x + self._mha("self_attn", s, s))[None])[0]
        x = _tap(taps, "decoder.cross_attn", self._ln(1, x + self._mha("cross_attn", x + qpe, key))[None])[0]
        w1, b1, w2, b2 = self._ffn
        ff = torch.relu(x @ w1 + b1) @ w2 + b2
        return _tap(taps, "decoder.out", self._ln(2, x + ff)[None])[0]

    def prediction_heads(self, x, qpos, qheat, cls_id, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]:
        torch = _torch()
        y = x.T[None]                                                          # (1, 128, K)
        out = {}
        for name, hid, fin in self.weights.head.prediction_heads:
            out[name] = _tap(taps, fin.name, self.layer(fin.name, self.layer(hid.name, y)))
        center = out["center"] + qpos.transpose(1, 2)                          # float32
        box = torch.cat([center, out["height"], out["dim"], out["rot"], out["vel"]], dim=1)[0]
        cls = (torch.sigmoid(out["heatmap"]) * qheat)[0]                       # (5, K)
        score = cls.gather(0, cls_id[None])[0]
        return {"bbox_pred": _tap(taps, "bbox_pred", box), "score": _tap(taps, "score", score),
                "label_pred": _tap(taps, "label_pred", cls_id.to(torch.int64))}

    # --------------------------------------------------------------------------------------------- whole
    def dense_forward(self, bev: Any, *, topk_indices: Optional[Any] = None,
                      taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]:
        """The plugin-free part (the dense sub-graph of the ONNX): BEV (1, 256, 180, 180) -> outputs."""
        torch = _torch()
        with torch.no_grad():
            bev = torch.as_tensor(np.asarray(bev, dtype=np.float32), device=self.device) \
                if not isinstance(bev, torch.Tensor) else bev
            neck = self.neck(self.backbone(bev, taps), taps)
            lidar, heat_nms = self.dense_head(neck, taps)
            _, idx = self.select(heat_nms, topk_indices, taps)
            query, qpos, qheat, cls_id = self.query_init(lidar, heat_nms, idx, taps)
            x = self.decoder(query, qpos, lidar, taps)
            out = self.prediction_heads(x, qpos, qheat, cls_id, taps)
        return {"bbox_pred": out["bbox_pred"].cpu().numpy().astype(np.float32),
                "score": out["score"].cpu().numpy().astype(np.float32),
                "label_pred": out["label_pred"].cpu().numpy().astype(np.int64)}

    def forward(self, vfe: np.ndarray, plan: Any, *, topk_indices: Optional[Any] = None,
                taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]:
        """Voxel features + sparse plan -> the three graph outputs: ``bbox_pred`` (10, K) float32, ``score`` (K,)
        float32, ``label_pred`` (K,) int64. ``topk_indices`` (K,) or (1, K): teacher-forced query selection."""
        torch = _torch()
        with torch.no_grad():
            bev = self.sparse_encoder(vfe, plan, taps)
        return self.dense_forward(bev, topk_indices=topk_indices, taps=taps)

    def describe(self) -> Dict[str, Any]:
        w = self.weights
        d = w.head.decoder
        return {"onnx_sha256": w.sha256, "num_proposals": self.K, "bev": list(BEV_SHAPE),
                "encoder": [(lay.name, lay.map, lay.in_channels, lay.out_channels, lay.kernel_volume)
                            for lay in w.encoder],
                "blocks": [[(lay.in_channels, lay.out_channels, lay.stride) for lay in b] for b in w.blocks],
                "deblocks": [(lay.kind, lay.in_channels, lay.out_channels, lay.kernel, lay.stride) for lay in w.deblocks],
                "decoder": {"heads": d.self_attn.num_heads, "head_dim": d.self_attn.head_dim,
                            "ffn": int(d.ffn_w1.shape[1])},
                "prediction_heads": {n: fin.out_channels for n, _, fin in w.head.prediction_heads},
                "params_after_folds": w.num_params()}

	# SPDX-License-Identifier: Apache-2.0
	"""fp32 PyTorch re-implementation of the deployed BEVFusion-L network, module by module, with taps.

	Built from :mod:`.weights` (the ONNX initializers addressed by node; every BatchNorm folded in float64 where the TT
	port folds it), it reproduces ``bevfusion_lidar.onnx`` (S:bevfusion:205-264) from the voxel features on (the VFE mean
	is host pre-processing in this port, computed in float32 exactly as the graph does by ``host.preprocess.vfe_mean``):

	- ``sparse`` (the 21 ``autoware::ImplicitGemm`` layers): every layer in the gather form the device uses (C13 /
	C28): ``y[o] = sum_k W'[:, k, :] @ x[nmap[o, k]] + b'`` over the neighbour maps of ``host.preprocess.SparsePlan``
	(``-1``: no input), ReLU, the basic blocks' residual before their second ReLU; BN folded (``W' = W * s``,
	``b' = beta - mean * s``). The plan's maps are spconv-faithful (13-query + symmetric write); the independent
	oracle is the research reference's scatter form (``tests/test_reference_cpu.py``);
	- ``to-dense``: ``dense[0, x, y, z, :] = conv_out[o]`` on a zero ``[1, 180, 180, 2, 128]`` map, ``Transpose(0, 4, 3,
	1, 2)``, ``Reshape [1, 256, 180, 180]``: channel ``c * 2 + z``, rows x, columns y;
	- ``backbone``: SECOND, 2 blocks of 6 3x3 convs + ReLU (first stride 2 in block 1); ``neck``: SECONDFPN deblock0 1x1
	+ ReLU, deblock1 ConvTranspose 2x2 s2 + ReLU, concat (1, 512, 180, 180);
	- ``head`` dense part: shared conv 3x3 512 -> 128 (``lidar_feat``, no activation), heatmap head (3x3 + ReLU, 3x3 -> 5
	logits), ``heat = sigmoid``, local max ``heat * (heat == local_max)`` with ``local_max`` = MaxPool 3x3 s1 p0 of
	classes 0-3 placed in a zero map (their border cells never survive) and class 4 = the heat itself (passes);
	- selection: the top ``K = 500`` of the flattened (5 x 32400) local-max heatmap as ONNX ``TopK`` defines it (largest
	first, ties to the lower index: a stable descending sort); ``topk_indices=`` replaces it by given indices
	(teacher forcing, PLAN.md 0.3 item 1: the device top-k breaks ties in no defined order, probe P8);
	- query init: ``class = idx // 32400``, ``pos = idx % 32400``; ``query = lidar_feat[:, pos] + class_table[class]``
	(the Conv1d of a one-hot is the table row ``W[:, c] + b``); ``query_pos = (pos // 180 + 0.5, pos % 180 + 0.5)``
	computed arithmetically in float32 (bit-equal to the graph's ``bev_pos`` gather, checked at load);
	``query_heat = heat_nms[:, :, pos]`` (all 5 classes);
	- decoder (one layer, post-norm, 8 heads x 16): self-attention on ``x + pe(query_pos)`` (Q = K = V), cross-attention
	of ``x + pe(query_pos)`` over ``lidar_feat^T + pe(bev_pos)`` (32,400 keys), FFN 128 -> 256 -> 128, LayerNorm eps
	1e-5 after each; ``Q`` divided by ``sqrt(16) = 4`` before ``Q K^T`` (the graph's ``Div``);
	- prediction heads (Conv1d 128 -> 64 + BN folded + ReLU, Conv1d 64 -> C: center 2, height 1, dim 3, rot 2, vel 2,
	heatmap 5);
	- outputs: ``bbox_pred = [center + query_pos, height, dim, rot, vel]`` (10, K), ``score = sigmoid(heatmap)[label] *
	query_heat[label]`` (K,) (the graph's one-hot product + ReduceMax, equal because both factors are >= 0),
	``label_pred = class`` (K,) int64.

	Taps (``ttaw.golden.TapRegistry``; stable names, the keys of the device gates; dense taps keep the ONNX tensors'
	batch axis and layout; sparse taps are ``(N_level, C)`` rows in the plan's level order): :data:`MODULE_TAPS` plus the
	opt-in per-layer / large taps :data:`EXTRA_TAPS`. Values are float32 numpy (``head.topk_indices`` and ``label_pred``
	int64). torch is imported lazily; the module runs where ttnn is absent.
	"""
	from __future__ import annotations

	from typing import Any, Dict, Optional

	import numpy as np

	from ..ttaw.golden import NULL_TAPS, TapRegistry
	from .config import BEV_SHAPE, SPARSE_LEVELS
	from .weights import ENCODER_PLAN, BEVFusionWeights

	__all__ = ["BEVFusionTorch", "MODULE_TAPS", "EXTRA_TAPS", "SELECTION_TAPS", "SPARSE_TAPS", "DENSE_TAPS",
	"OUTPUT_NAMES", "stable_topk"]

	OUTPUT_NAMES = ("bbox_pred", "score", "label_pred")
	SPARSE_TAPS = ("vfe",) + tuple(name for name, *_ in ENCODER_PLAN if not name.endswith(".conv1")) + ("sparse.bev",)
	DENSE_TAPS = ("backbone.block0", "backbone.block1", "neck.deblock0", "neck.deblock1", "head.shared", "head.heatmap",
	"head.heat", "head.heat_nms")
	# taps downstream of the data-dependent top-k (compared set-based or teacher-forced, PLAN.md 0.3 item 1)
	SELECTION_TAPS = ("head.topk_values", "head.topk_indices", "query.feat", "query.pos", "query.heat",
	"decoder.query_pos_embed", "decoder.self_attn", "decoder.cross_attn", "decoder.out",
	"heads.center", "heads.height", "heads.dim", "heads.rot", "heads.vel", "heads.heatmap") + OUTPUT_NAMES
	MODULE_TAPS = SPARSE_TAPS + DENSE_TAPS + SELECTION_TAPS
	# opt-in taps: per layer, the 66 MB neck concat and the input-independent key position embedding
	EXTRA_TAPS = (("neck", "head.heatmap.hidden", "decoder.key_pos_embed")
	+ tuple(name for name, *_ in ENCODER_PLAN if name.endswith(".conv1"))
	+ tuple(f"backbone.block{b}.conv{j}" for b in range(2) for j in range(6)))


	def _torch():
	import torch

	return torch


	def _tap(taps: TapRegistry, name: str, value):
	"""Record a torch tensor as a numpy copy (float32, or int64 for integer tensors) and return it unchanged; never
	through ``ttaw.tensors.to_numpy`` (which would import ttnn)."""
	if taps.enabled and taps.wants(name):
	a = value.detach().cpu().numpy()
	taps.tap(name, a.astype(np.int64 if a.dtype.kind in "iu" else np.float32, copy=True))
	return value


	def stable_topk(scores, k: int):
	"""ONNX ``TopK(largest=1, sorted=1)`` on the last axis: values descending, equal values in ascending index order
	(a stable descending sort). Returns ``(values, indices)``."""
	torch = _torch()
	order = torch.sort(scores, dim=-1, descending=True, stable=True)
	return order.values[..., :k], order.indices[..., :k]


	class BEVFusionTorch:
	"""The fp32 reference network. ``threads`` caps torch's intra-op threads (shared host: <= 4)."""

	def __init__(self, weights: BEVFusionWeights, *, threads: Optional[int] = None, device: str = "cpu"):
	torch = _torch()
	if threads is not None:
	torch.set_num_threads(int(threads))
	self.weights = weights
	self.device = device
	t = self._t
	self._sparse = [(layer, t(layer.weight.reshape(layer.out_channels, -1, layer.in_channels)), t(layer.bias))
	for layer in weights.encoder]
	self._layers: Dict[str, tuple] = {}
	for layer in weights.dense_layers():
	self._layers[layer.name] = (layer, t(layer.weight), t(layer.bias))
	h = weights.head
	d = h.decoder
	for layer in (d.self_posembed, d.cross_posembed) + tuple(x for _, a, b in h.prediction_heads for x in (a, b)):
	self._layers[layer.name] = (layer, t(layer.weight), t(layer.bias))
	self.class_table = t(h.class_table) # (5, 128)
	self.bev_pos = t(h.bev_pos) # (32400, 2)
	self.K = int(h.num_proposals)
	self.bev_h, self.bev_w = BEV_SHAPE
	self.pooled = tuple(h.pooled_classes)
	self._attn = {"self_attn": self._attn_tensors(d.self_attn), "cross_attn": self._attn_tensors(d.cross_attn)}
	self._norms = [(t(g), t(b), float(eps)) for g, b, eps in d.norms]
	self._ffn = (t(d.ffn_w1), t(d.ffn_b1), t(d.ffn_w2), t(d.ffn_b2))
	self._kpe = None # cross_posembed(bev_pos): input-independent, cached

	def _t(self, a: Any):
	torch = _torch()
	return torch.from_numpy(np.array(a, dtype=np.float32, copy=True)).to(self.device)

	def _attn_tensors(self, a) -> dict:
	t = self._t
	return {"wq": t(a.wq), "bq": t(a.bq), "wk": t(a.wk), "bk": t(a.bk), "wv": t(a.wv), "bv": t(a.bv),
	"wo": t(a.wo), "bo": t(a.bo), "heads": a.num_heads, "hd": a.head_dim, "div": float(a.q_divisor)}

	# ---------------------------------------------------------------------------------------- sparse encoder
	def sparse_conv(self, x, nmap: np.ndarray, w, b):
	"""Gather-form sparse conv: ``x`` (N_in, Cin) torch, ``nmap`` (N_out, K) int (-1 = none), ``w`` (Cout, K,
	Cin) -> (N_out, Cout) (bias added, no activation). Taps accumulate in tap order, fp32."""
	torch = _torch()
	n_in, cin = x.shape
	xp = torch.cat([x, torch.zeros((1, cin), dtype=x.dtype, device=x.device)], 0) # explicit zero row
	idx = torch.from_numpy(np.where(nmap >= 0, nmap, n_in).astype(np.int64)).to(x.device)
	out = torch.zeros((idx.shape[0], w.shape[0]), dtype=x.dtype, device=x.device)
	for k in range(idx.shape[1]):
	col = idx[:, k]
	if bool((col < n_in).any()):
	out += xp[col] @ w[:, k, :].T
	return out + b

	def sparse_encoder(self, vfe: np.ndarray, plan: Any, taps: TapRegistry = NULL_TAPS):
	"""(V, 4) voxel features + ``host.preprocess.SparsePlan`` -> the BEV (1, 256, 180, 180)."""
	torch = _torch()
	x = _tap(taps, "vfe", torch.as_tensor(np.asarray(vfe, dtype=np.float32), device=self.device))
	if x.shape[0] != plan.levels["l1"].num_active:
	raise ValueError(f"{x.shape[0]} voxel features for {plan.levels['l1'].num_active} active voxels")
	outs: Dict[str, Any] = {}
	for layer, w, b in self._sparse:
	y = self.sparse_conv(x, plan.maps[layer.map], w, b)
	if layer.residual is not None:
	y = y + outs[layer.residual]
	x = _tap(taps, layer.name, torch.relu(y))
	outs[layer.name] = x
	lvl = plan.levels["out"]
	X, Y, Z = lvl.spatial_shape
	if (X, Y) != BEV_SHAPE or lvl.spatial_shape != SPARSE_LEVELS["out"]:
	raise ValueError(f"conv_out grid {lvl.spatial_shape}")
	dense = torch.zeros((1, X, Y, Z, x.shape[1]), dtype=x.dtype, device=x.device)
	oc = torch.from_numpy(lvl.coords).to(x.device)
	dense[0, oc[:, 0], oc[:, 1], oc[:, 2]] = x
	bev = dense.permute(0, 4, 3, 1, 2).reshape(1, x.shape[1] * Z, X, Y)
	return _tap(taps, "sparse.bev", bev)

	# ----------------------------------------------------------------------------------------- dense layers
	def layer(self, name: str, x):
	torch = _torch()
	F = torch.nn.functional
	layer, w, b = self._layers[name]
	if layer.kind == "conv":
	y = F.conv2d(x, w, b, stride=layer.stride, padding=layer.padding)
	elif layer.kind == "conv_transpose":
	y = F.conv_transpose2d(x, w, b, stride=layer.stride, padding=layer.padding)
	else:
	y = F.conv1d(x, w, b, stride=layer.stride, padding=layer.padding)
	return torch.relu(y) if layer.relu else y

	def backbone(self, bev, taps: TapRegistry = NULL_TAPS):
	outs, x = [], bev
	for i, block in enumerate(self.weights.blocks):
	for layer in block:
	x = _tap(taps, layer.name, self.layer(layer.name, x))
	outs.append(_tap(taps, f"backbone.block{i}", x))
	return outs

	def neck(self, blocks, taps: TapRegistry = NULL_TAPS):
	torch = _torch()
	ups = [_tap(taps, layer.name, self.layer(layer.name, x)) for layer, x in zip(self.weights.deblocks, blocks)]
	return _tap(taps, "neck", torch.cat(ups, dim=1))

	def dense_head(self, neck, taps: TapRegistry = NULL_TAPS):
	"""-> (lidar_feat (1, 128, 180, 180), heat_nms (1, 5, 180, 180))."""
	torch = _torch()
	F = torch.nn.functional
	lidar = _tap(taps, "head.shared", self.layer("head.shared", neck))
	hid = _tap(taps, "head.heatmap.hidden", self.layer("head.heatmap.hidden", lidar))
	logits = _tap(taps, "head.heatmap", self.layer("head.heatmap", hid))
	heat = _tap(taps, "head.heat", torch.sigmoid(logits))
	k = int(self.weights.head.nms_kernel)
	pad = k // 2
	pooled = list(self.pooled)
	local_max = torch.zeros_like(heat)
	local_max[:, pooled, pad:-pad, pad:-pad] = F.max_pool2d(heat[:, pooled], k, 1, 0)
	rest = [c for c in range(heat.shape[1]) if c not in self.pooled]
	local_max[:, rest] = heat[:, rest]
	heat_nms = _tap(taps, "head.heat_nms", heat * (heat == local_max).to(heat.dtype))
	return lidar, heat_nms

	def select(self, heat_nms, topk_indices: Optional[Any] = None, taps: TapRegistry = NULL_TAPS):
	"""Top-K of the flattened local-max heatmap -> (values (1, K), indices (1, K) int64); given ``topk_indices``
	(teacher forcing) are used as they are and their heat values gathered."""
	torch = _torch()
	flat = heat_nms.reshape(1, -1)
	if topk_indices is None:
	values, idx = stable_topk(flat, self.K)
	else:
	idx = torch.as_tensor(np.asarray(topk_indices, dtype=np.int64).reshape(1, -1), device=self.device)
	if idx.shape[1] != self.K or int(idx.min()) < 0 or int(idx.max()) >= flat.shape[1]:
	raise ValueError(f"topk_indices must be {self.K} indices in [0, {flat.shape[1]})")
	values = flat.gather(1, idx)
	return _tap(taps, "head.topk_values", values), _tap(taps, "head.topk_indices", idx)

	def query_init(self, lidar, heat_nms, idx, taps: TapRegistry = NULL_TAPS):
	"""-> (query (1, 128, K), query_pos (1, K, 2), query_heat (1, 5, K), class (K,))."""
	torch = _torch()
	n_cells = lidar.shape[2] * lidar.shape[3]
	cls_id = torch.div(idx, n_cells, rounding_mode="floor")
	pos = idx - cls_id * n_cells
	flat = lidar.reshape(1, lidar.shape[1], -1)
	q = flat.gather(2, pos[:, None, :].expand(-1, flat.shape[1], -1)) + self.class_table[cls_id[0]].T[None]
	w = self.bev_w
	qpos = torch.stack([torch.div(pos, w, rounding_mode="floor").to(torch.float32) + 0.5,
	(pos % w).to(torch.float32) + 0.5], dim=-1)
	heat_flat = heat_nms.reshape(1, heat_nms.shape[1], -1)
	qheat = heat_flat.gather(2, pos[:, None, :].expand(-1, heat_flat.shape[1], -1))
	return (_tap(taps, "query.feat", q), _tap(taps, "query.pos", qpos), _tap(taps, "query.heat", qheat),
	cls_id[0])

	def _posembed(self, which: str, pos):
	"""(P, 2) -> (P, 128): Conv1d 2 -> 128 + ReLU (BN folded), Conv1d 128 -> 128 (1x1 convs over P positions)."""
	x = pos.T[None] # (1, 2, P)
	x = self.layer(f"decoder.{which}.0", x)
	return self.layer(f"decoder.{which}.1", x)[0].T # (P, 128)

	def key_pos_embed(self, taps: TapRegistry = NULL_TAPS):
	"""``cross_posembed(bev_pos)`` (32400, 128): input-independent, computed once."""
	torch = _torch()
	if self._kpe is None:
	with torch.no_grad():
	self._kpe = self._posembed("cross_posembed", self.bev_pos)
	return _tap(taps, "decoder.key_pos_embed", self._kpe[None])[0]

	def _mha(self, which: str, q_in, kv_in, v_in=None):
	"""Multi-head attention as exported: projections, per-head ``softmax((q / div) @ k^T) @ v``, out proj."""
	torch = _torch()
	a = self._attn[which]
	v_in = kv_in if v_in is None else v_in
	q = q_in @ a["wq"] + a["bq"]
	k = kv_in @ a["wk"] + a["bk"]
	v = v_in @ a["wv"] + a["bv"]
	H, D = a["heads"], a["hd"]
	lq, lk = q.shape[0], k.shape[0]
	q = q.reshape(lq, H, D).transpose(0, 1) / a["div"]
	k = k.reshape(lk, H, D).permute(1, 2, 0)
	v = v.reshape(lk, H, D).transpose(0, 1)
	o = torch.softmax(q @ k, dim=-1) @ v # (H, Lq, D)
	o = o.transpose(0, 1).reshape(lq, H * D)
	return o @ a["wo"] + a["bo"]

	def _ln(self, i: int, x):
	torch = _torch()
	g, b, eps = self._norms[i]
	return torch.nn.functional.layer_norm(x, (x.shape[-1],), g, b, eps)

	def decoder(self, query, qpos, lidar, taps: TapRegistry = NULL_TAPS):
	"""One post-norm decoder layer -> (K, 128)."""
	torch = _torch()
	qpe = _tap(taps, "decoder.query_pos_embed", self._posembed("self_posembed", qpos[0])[None])[0]
	kpe = self.key_pos_embed(taps)
	x = query[0].T # (K, 128)
	key = lidar.reshape(lidar.shape[1], -1).T + kpe # (32400, 128)
	s = x + qpe
	x = _tap(taps, "decoder.self_attn", self._ln(0, x + self._mha("self_attn", s, s))[None])[0]
	x = _tap(taps, "decoder.cross_attn", self._ln(1, x + self._mha("cross_attn", x + qpe, key))[None])[0]
	w1, b1, w2, b2 = self._ffn
	ff = torch.relu(x @ w1 + b1) @ w2 + b2
	return _tap(taps, "decoder.out", self._ln(2, x + ff)[None])[0]

	def prediction_heads(self, x, qpos, qheat, cls_id, taps: TapRegistry = NULL_TAPS) -> Dict[str, Any]:
	torch = _torch()
	y = x.T[None] # (1, 128, K)
	out = {}
	for name, hid, fin in self.weights.head.prediction_heads:
	out[name] = _tap(taps, fin.name, self.layer(fin.name, self.layer(hid.name, y)))
	center = out["center"] + qpos.transpose(1, 2) # float32
	box = torch.cat([center, out["height"], out["dim"], out["rot"], out["vel"]], dim=1)[0]
	cls = (torch.sigmoid(out["heatmap"]) * qheat)[0] # (5, K)
	score = cls.gather(0, cls_id[None])[0]
	return {"bbox_pred": _tap(taps, "bbox_pred", box), "score": _tap(taps, "score", score),
	"label_pred": _tap(taps, "label_pred", cls_id.to(torch.int64))}

	# --------------------------------------------------------------------------------------------- whole
	def dense_forward(self, bev: Any, *, topk_indices: Optional[Any] = None,
	taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]:
	"""The plugin-free part (the dense sub-graph of the ONNX): BEV (1, 256, 180, 180) -> outputs."""
	torch = _torch()
	with torch.no_grad():
	bev = torch.as_tensor(np.asarray(bev, dtype=np.float32), device=self.device) \
	if not isinstance(bev, torch.Tensor) else bev
	neck = self.neck(self.backbone(bev, taps), taps)
	lidar, heat_nms = self.dense_head(neck, taps)
	_, idx = self.select(heat_nms, topk_indices, taps)
	query, qpos, qheat, cls_id = self.query_init(lidar, heat_nms, idx, taps)
	x = self.decoder(query, qpos, lidar, taps)
	out = self.prediction_heads(x, qpos, qheat, cls_id, taps)
	return {"bbox_pred": out["bbox_pred"].cpu().numpy().astype(np.float32),
	"score": out["score"].cpu().numpy().astype(np.float32),
	"label_pred": out["label_pred"].cpu().numpy().astype(np.int64)}

	def forward(self, vfe: np.ndarray, plan: Any, *, topk_indices: Optional[Any] = None,
	taps: TapRegistry = NULL_TAPS) -> Dict[str, np.ndarray]:
	"""Voxel features + sparse plan -> the three graph outputs: ``bbox_pred`` (10, K) float32, ``score`` (K,)
	float32, ``label_pred`` (K,) int64. ``topk_indices`` (K,) or (1, K): teacher-forced query selection."""
	torch = _torch()
	with torch.no_grad():
	bev = self.sparse_encoder(vfe, plan, taps)
	return self.dense_forward(bev, topk_indices=topk_indices, taps=taps)

	def describe(self) -> Dict[str, Any]:
	w = self.weights
	d = w.head.decoder
	return {"onnx_sha256": w.sha256, "num_proposals": self.K, "bev": list(BEV_SHAPE),
	"encoder": [(lay.name, lay.map, lay.in_channels, lay.out_channels, lay.kernel_volume)
	for lay in w.encoder],
	"blocks": [[(lay.in_channels, lay.out_channels, lay.stride) for lay in b] for b in w.blocks],
	"deblocks": [(lay.kind, lay.in_channels, lay.out_channels, lay.kernel, lay.stride) for lay in w.deblocks],
	"decoder": {"heads": d.self_attn.num_heads, "head_dim": d.self_attn.head_dim,
	"ffn": int(d.ffn_w1.shape[1])},
	"prediction_heads": {n: fin.out_channels for n, _, fin in w.head.prediction_heads},
	"params_after_folds": w.num_params()}