Blame for code/tt_bevformer/tt/encoder.py

# SPDX-License-Identifier: Apache-2.0
"""The BEV encoder on the device: BEV queries, the rotated history, 3 x (TSA -> LN -> rebatched SCA -> LN -> FFN ->
LN). Stock ttnn ops only; every tensor of the BEV stream is a TILE ``[1, 1, 22500, 256]`` bf16 map in DRAM.

Per frame (RT-dev values written by the host before the replay, ``tt.model``):

- ``can_bus`` ``[1, 1, 1, 18]`` fp32 -> ``can_bus_mlp`` (Linear-ReLU-Linear-ReLU-LayerNorm, fp32) -> broadcast-added
  to ``bev_embedding``: the BEV queries;
- ``rot_index`` ``[1, Q]`` uint32: the prev_bev rotation as a row gather (``RotateTRT`` nearest, round half away,
  zero fill: ``host.geometry.rotation_index_map``) from the ``prev_bev`` state ``[1, 1, Q + 1, 256]`` ROW_MAJOR bf16
  (the previous ``bev_embed`` plus a zero row: the gather table; PADDED ``ttnn.embedding``, C19 rules);
- ``use`` / ``omu`` (``use_prev_bev`` and ``1 - use_prev_bev``, exactly 0 / 1): the TSA queue values
  ``use [rot(prev_bev); bev_queries] + (1 - use) [query; query]``;
- ``tsa_shift`` ``[1, 16, 1, 32]``: ``2 shift use`` added to the queue-0 grids (``tt.params.tsa_shift_rows``);
- per calibration (re-uploaded only when ``lidar2img`` changes): ``sca_index`` ``[1, 6 * CAP]`` (the BEV query of every
  rebatched row, ``Q`` = padding), ``sca_base`` ``[1, 1, 6 * CAP, 32]`` (the rows' ``2 ref_cam - 1``), ``sca_g0..2``
  ``[1, Q]`` (the gather-form scatter-add: rows of the stacked ``[6 * CAP + 1]`` SCA output, ``6 * CAP`` = zero row)
  and ``sca_scale`` ``[1, 1, Q, 1]`` (1 / cameras seeing the query);
- the SCA value maps of the 3 layers (the image segment's hand-off state).

TSA (``tt.params.tsa``): ``qc = [v0 | query + bev_pos]`` -> one folded sampling projection (16 grid blocks, ``n =
queue * 8 + head``) + ``tsa_base`` + ``tsa_shift`` -> grid; one attention projection -> softmax weights; value_proj of
``[v0; v1]`` -> 16 value maps 150 x 150; fused MSDA; the 16 heads merged into ``[Q, 512]`` and the queue mean folded
into ``output_proj``. SCA: gather the query rows of every camera (``SCA_CAPACITY`` per camera), the folded
projections on the 36,864 rows, fused MSDA over 48 value maps (``n = head * 6 + camera``), heads merged, scatter-add
as 3 gathers of the stacked output + its zero row, times the 1 / count scale, ``output_proj``. Residual adds are
fused into ``ttnn.layer_norm`` (eps 1e-5 passed explicitly: the ttnn default is 1e-12).
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional

import numpy as np

from ..reference import config as C
from ..ttaw.models.second import RowLinear
from ..ttaw.ops.gather import gather_rows, with_zero_rows
from . import kconf
from .image import _Const
from .msda import GridSampleMSDA, attend, free, site_modes, value_maps
from .params import PX_SCALE, BEVFormerParams, LayerNormP, compact_blocks, tsa_base, tsa_base_px

__all__ = ["DEFAULT_PRECISION", "TtLayerNorm", "TtEncoder", "SCA_ROWS"]

DEFAULT_PRECISION = "HiFi4+fp32+l1acc:w=fp32"
Q = C.NUM_QUERY_BEV
E = C.EMBED
SCA_ROWS = C.NUM_CAMS * C.SCA_CAPACITY


class TtLayerNorm:
    """``LayerNorm(x [+ residual])`` with the model's epsilon and fp32 affine rows (uploaded on first use)."""

    def __init__(self, p: LayerNormP, *, precision: Any = DEFAULT_PRECISION):
        from ..ttaw.precision import Precision

        self.p = p
        self.cfg = Precision.parse(precision).compute_kernel_config()
        self.g = _Const(np.asarray(p.gamma, np.float32).reshape(1, -1), "float32")
        self.b = _Const(np.asarray(p.beta, np.float32).reshape(1, -1), "float32")

    def __call__(self, x: Any, residual: Optional[Any] = None, l1: bool = False):
        import ttnn

        dev = x.device()
        kw = {} if residual is None else {"residual_input_tensor": residual}
        mc = ttnn.L1_MEMORY_CONFIG if l1 else ttnn.DRAM_MEMORY_CONFIG
        return ttnn.layer_norm(x, weight=self.g(dev), bias=self.b(dev), epsilon=float(self.p.eps),
                               compute_kernel_config=self.cfg, memory_config=mc, **kw)

    def release(self) -> None:
        self.g.release()
        self.b.release()


_TUNED = {"on": False}       # BEVFORMER_BEV_MM=tuned (set by TtEncoder before its linears are built)


_VP_L1 = {"on": False}           # BEVFORMER_TSA_VP_L1 (set by TtEncoder)


def _lin(p, name: str, precision: Any, *, activation: Optional[str] = None, out: str = "bfloat16") -> RowLinear:
    cls = _TunedRowLinear if _TUNED["on"] else RowLinear
    return cls(p.w, p.b, activation=activation, precision=kconf.prec(name, precision), output_dtype=out, name=name)


class _TunedRowLinear(RowLinear):
    """``BEVFORMER_BEV_MM=tuned``: ttnn.linear with the :data:`tt.linconv.TUNED_BEV` config of the shape (same K
    blocking as the default: bit-identical), else the default."""

    def __call__(self, x: Any) -> Any:
        import ttnn

        from ..ttaw.tensors import ttnn_dtype
        from .linconv import bev_program_config

        weight, bias, grid, cfg = self._device or self._upload(x.device())
        m = int(x.shape[0]) * -(-int(x.shape[-2]) // 32) * 32            # rows padded to tiles per batch
        pc = bev_program_config(m, self.in_channels, self.out_channels, grid, self.activation)
        if pc is None:
            return RowLinear.__call__(self, x)
        return ttnn.linear(x, weight, bias=bias, program_config=pc, compute_kernel_config=cfg,
                           memory_config=ttnn.DRAM_MEMORY_CONFIG, dtype=ttnn_dtype(self.output_dtype))


def _linear_to_l1(lin: Any, x: Any) -> Any:
    """``lin(x)`` (a ``RowLinear`` / ``_TunedRowLinear``) with its output in L1 interleaved: the same ttnn.linear
    call (program config, compute config) but for the output memory config, so the same values."""
    import ttnn

    from ..ttaw.tensors import ttnn_dtype

    weight, bias, grid, cfg = lin._device or lin._upload(x.device())
    if isinstance(lin, _TunedRowLinear):
        from .linconv import bev_program_config

        m = int(x.shape[0]) * -(-int(x.shape[-2]) // 32) * 32
        pc = bev_program_config(m, lin.in_channels, lin.out_channels, grid, lin.activation)
        if pc is not None:
            return ttnn.linear(x, weight, bias=bias, program_config=pc, compute_kernel_config=cfg,
                               memory_config=ttnn.L1_MEMORY_CONFIG, dtype=ttnn_dtype(lin.output_dtype))
    return ttnn.linear(x, weight, bias=bias, activation=lin.activation, compute_kernel_config=cfg, core_grid=grid,
                       memory_config=ttnn.L1_MEMORY_CONFIG, dtype=ttnn_dtype(lin.output_dtype))


class _PinnedLinear(RowLinear):
    """The compact sampling projection (``BEVFORMER_MSDA_LOC=compact``): ``x @ W + b`` in fp32 with the 1D config
    ttnn.linear picks for the head-block projection except N (K advanced one tile per step, ``in0_block_w`` 1, as
    measured in the profile), so every kept column equals the head-block projection's column bit for bit."""

    def __call__(self, x: Any) -> Any:
        import ttnn

        weight, bias, grid, cfg = self._device or self._upload(x.device())
        if _TUNED["on"]:
            from .linconv import bev_program_config

            pc = bev_program_config(-(-int(x.shape[-2]) // 32) * 32, self.in_channels, self.out_channels, grid)
            if pc is not None:
                return ttnn.linear(x, weight, bias=bias, program_config=pc, compute_kernel_config=cfg,
                                   memory_config=ttnn.DRAM_MEMORY_CONFIG, dtype=ttnn.float32)
        m = -(-int(x.shape[-2]) // 32)
        nt = self.out_channels // 32
        pcm = -(-m // (int(grid.x) * int(grid.y)))
        sh = 2 if pcm % 2 == 0 else 1
        sw = 2 if nt % 2 == 0 else 1
        pc = ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(
            compute_with_storage_grid_size=(int(grid.x), int(grid.y)), in0_block_w=1, out_subblock_h=sh,
            out_subblock_w=sw, out_block_h=pcm, out_block_w=nt, per_core_M=pcm, per_core_N=nt, fuse_batch=True,
            fused_activation=None, mcast_in0=False)
        return ttnn.linear(x, weight, bias=bias, program_config=pc, compute_kernel_config=cfg,
                           memory_config=ttnn.DRAM_MEMORY_CONFIG, dtype=ttnn.float32)


def _sampling(p, name: str, precision: Any, blocks: int, cols: int, compact: bool):
    if compact:
        c = compact_blocks(p, blocks, cols)
        return _PinnedLinear(c.w, c.b, precision=kconf.prec(name, precision), output_dtype="float32",
                             name=name + ".compact")
    return _lin(p, name, precision, out="float32")


class _TSA:
    def __init__(self, params: BEVFormerParams, i: int, precision: Any, msda_mode: str, grid_msda: GridSampleMSDA,
                 px: bool = False, tile: bool = False, skip: bool = False, compact: bool = False,
                 sm: bool = False, pairs: bool = False, mk: bool = False):
        self.msda_mode, self.grid_msda = msda_mode, grid_msda
        self.px = bool(px) and msda_mode == "kernel"
        self.tile = bool(tile) and msda_mode == "kernel"
        self.skip = bool(skip)
        self.loc_hs = 2 * C.TSA_POINTS if compact and self.px else 0
        p = params.tsa(i)
        n = f"encoder.layers.{i}.tsa"
        self.value_proj = _lin(p["value_proj"], f"{n}.value_proj", precision)
        self.sampling = _sampling(p["sampling_px" if self.px else "sampling"], f"{n}.sampling", precision,
                                  C.TSA_QUEUE * C.NUM_HEADS, self.loc_hs, bool(self.loc_hs))
        self.sm = bool(sm) and bool(self.loc_hs)
        self.attention = _sampling(p["attention"], f"{n}.attention", precision, C.TSA_QUEUE * C.NUM_HEADS,
                                   C.TSA_POINTS, self.sm)
        # MSDA_PAIRS=all: the TSA value as pixel-pair maps [16, 150, 151, 64] in L1 (tt.pairmap_k)
        self.pairs = bool(pairs) and self.sm and not self.skip
        # MSDA_MK (tsa / all): concat + both projections + loc glue + msda_k as one tt.msda_mk program
        self.mk = None
        if mk and self.pairs:
            from .msda_mk import MkWeights

            self.mk = MkWeights(self.sampling, self.attention)
        self.output_proj = _lin(p["output_proj"], f"{n}.output_proj", precision)
        from ..ttaw.precision import Precision

        self.cfg = Precision.parse(precision).compute_kernel_config()

    def __call__(self, query, v0, v1, bev_pos, base, shift, vin=None):
        """-> ``output_proj(mean_queue(MSDA))`` ``[1, 1, Q, 256]`` (no residual). ``vin``: the ``[v0; v1]`` value
        input when the caller holds it (``bev_prev``: the same for every layer)."""
        import ttnn

        heads = C.TSA_QUEUE * C.NUM_HEADS
        qpos = ttnn.add(query, bev_pos, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        if self.mk is not None:
            cols = logits = None                  # computed inside msda_mk from [v0 | qpos]
        else:
            qc = ttnn.concat([v0, qpos], dim=3, memory_config=ttnn.DRAM_MEMORY_CONFIG)
            free(qpos)
            cols, logits = self.sampling(qc), self.attention(qc)
            free(qc)
        own = vin is None
        if own:
            vin = ttnn.concat([v0, v1], dim=0, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        vp_l1 = self.pairs and _VP_L1["on"]
        vp = _linear_to_l1(self.value_proj, vin) if vp_l1 else self.value_proj(vin)
        if own:
            free(vin)
        if self.pairs:
            from .pairmap_k import pair_maps

            # TSA_VP_L1: the projection was written to L1 directly (no DRAM -> L1 copy of the 23 MB maps' input)
            vp1 = vp if vp_l1 else ttnn.to_memory_config(vp, ttnn.L1_MEMORY_CONFIG)
            value = pair_maps(vp1, heads=C.NUM_HEADS, sub=1, hw=(C.BEV_H, C.BEV_W))
            if not vp_l1:
                free(vp1)
        else:
            value = value_maps(vp, num_heads=C.NUM_HEADS, hw=(C.BEV_H, C.BEV_W), l1=self.msda_mode == "kernel",
                               tile=self.tile)
        free(vp)
        if self.mk is not None:
            from .msda_mk import msda_block

            # BEVFORMER_ENC_L1 also keeps the TSA megakernel's [22528, 512] output in L1 for the output projection
            merged = msda_block(value, v0, qpos, self.mk, base, shift, heads=heads, points=C.TSA_POINTS, sub=1,
                                hw=(C.BEV_H, C.BEV_W), nq=Q,
                                memory_config=ttnn.L1_MEMORY_CONFIG if kconf.FFN_L1 else None)
            free(value, qpos)
        else:
            merged = attend(value, cols, logits, num_heads=heads, points=C.TSA_POINTS, base=base, shift=shift,
                            mode=self.msda_mode, grid_msda=self.grid_msda, compute_kernel_config=self.cfg,
                            px=self.px, hw=(C.BEV_H, C.BEV_W), skip=self.skip, loc_hs=self.loc_hs, softmax=self.sm)
            free(value, cols, logits)
        # BEVFORMER_ENC_L1=all (D20 round 2): the output projection writes L1 for the LayerNorm
        y = _linear_to_l1(self.output_proj, merged) if kconf.ENC_L1_ALL else self.output_proj(merged)
        free(merged)
        return y

    def linears(self) -> List[RowLinear]:
        return [self.value_proj, self.sampling, self.attention, self.output_proj]


class _SCA:
    def __init__(self, params: BEVFormerParams, i: int, precision: Any, msda_mode: str, grid_msda: GridSampleMSDA,
                 px: bool = False, skip: bool = False, compact: bool = False, sm: bool = False,
                 tile_gather: bool = False, value_dram: bool = False, mk: bool = False):
        self.msda_mode, self.grid_msda = msda_mode, grid_msda
        self.value_dram = bool(value_dram)   # SCA_VALUE=dram: msda_k reads the DRAM state (no L1 copy)
        self.tile_gather = bool(tile_gather)
        self.px = bool(px) and msda_mode == "kernel"
        self.skip = bool(skip)
        self.loc_hs = 2 * C.SCA_POINTS if compact and self.px else 0
        p = params.sca(i)
        n = f"encoder.layers.{i}.sca"
        self.sampling = _sampling(p["sampling_px" if self.px else "sampling"], f"{n}.sampling", precision,
                                  C.NUM_HEADS, self.loc_hs, bool(self.loc_hs))
        self.sm = bool(sm) and bool(self.loc_hs)
        self.attention = _sampling(p["attention"], f"{n}.attention", precision, C.NUM_HEADS, C.SCA_POINTS,
                                   self.sm)
        self.output_proj = _lin(p["output_proj"], f"{n}.output_proj", precision)
        # MSDA_MK (sca / all): both projections + loc glue + msda_k as one tt.msda_mk program (pair maps only)
        self.mk = None
        if mk and self.sm and self.tile_gather:
            from .msda_mk import MkWeights

            self.mk = MkWeights(self.sampling, self.attention)
        from ..ttaw.precision import Precision

        self.cfg = Precision.parse(precision).compute_kernel_config()

    def __call__(self, query, value, ctx: Dict[str, Any]):
        """``query`` ``[1, 1, Q, 256]``, ``value`` ``(48, 23, 40, 32)``; ``ctx`` holds the per-calibration tables ->
        ``output_proj(sum over cameras / count)`` ``[1, 1, Q, 256]`` (no residual)."""
        import ttnn

        cap, nh, ncam = C.SCA_CAPACITY, C.NUM_HEADS, C.NUM_CAMS
        if self.tile_gather:
            from .rowgather_k import tile_gather

            q_l1 = query.memory_config().buffer_type == ttnn.BufferType.L1
            q1 = query if q_l1 else ttnn.to_memory_config(query, ttnn.L1_MEMORY_CONFIG)
            rows = tile_gather(ctx["sca_index"], q1)                         # [1, 1, 6 CAP, 256] TILE
            if not q_l1:
                free(q1)
        else:
            table = with_zero_rows(query)                                    # [1, 1, Q + 1, 256] RM
            rows = gather_rows(ctx["sca_index"], table, sentinel=Q)          # [1, 1, 6 CAP, 256] TILE
            free(table)
        aff = ctx.get("sca_base_aff")
        if aff is not None and self.mk is not None and int(value.shape[-1]) == 64:
            base = ctx["sca_base"]                                          # raw: msda_mk applies S, O
        else:
            base = ctx["sca_base_c"] if self.loc_hs else (ctx["sca_base_px"] if self.px else ctx["sca_base"])
        if self.mk is not None and int(value.shape[-1]) == 64:
            from .msda_mk import msda_block

            from . import kconf as kconf_

            own = value.memory_config().buffer_type != ttnn.BufferType.L1 and not self.value_dram
            v = value
            if own:
                if kconf_.SCA_COPY_K:      # BEVFORMER_SCA_COPY=kernel (D20 round 3): repage_k (bit-identical)
                    from .repage_k import repage

                    v = repage(value, ttnn.L1_MEMORY_CONFIG)
                else:
                    v = ttnn.to_memory_config(value, ttnn.L1_MEMORY_CONFIG)
            from . import kconf

            act = ctx.get("sca_nblk") if kconf.SCA_ACTIVE else None
            merged = msda_block(v, rows, None, self.mk, base, None, heads=nh, points=C.SCA_POINTS, sub=ncam,
                                hw=C.FEAT_HW, nq=int(rows.shape[-2]), memory_config=ttnn.L1_MEMORY_CONFIG,
                                active=act, base_aff=aff)
            free(rows)
            if own:
                free(v)
        else:
            cols, logits = self.sampling(rows), self.attention(rows)
            free(rows)
            merged = attend(value, cols, logits, num_heads=nh, points=C.SCA_POINTS, base=base, sub=ncam,
                            mode=self.msda_mode, grid_msda=self.grid_msda, compute_kernel_config=self.cfg,
                            px=self.px, skip=self.skip, loc_hs=self.loc_hs, softmax=self.sm,
                            out_l1=self.tile_gather and self.sm, value_in_place=self.value_dram)
            free(cols, logits)                                               # merged: [1, 1, 6 CAP, 256]
        if self.tile_gather:
            from .rowgather_k import tile_gather

            if merged.memory_config().buffer_type == ttnn.BufferType.L1:
                stacked = merged                                             # msda_k wrote it in L1
            else:
                stacked = ttnn.to_memory_config(merged, ttnn.L1_MEMORY_CONFIG)   # index 6 CAP -> zero row
                free(merged)
        else:
            stacked = with_zero_rows(merged)                                 # [1, 1, 6 CAP + 1, 256] RM
            free(merged)
        from . import kconf

        if self.tile_gather and kconf.SCA_SCATTER_K and C.SCA_MAX_CAMS_PER_QUERY == 3:
            from .scatter_k import scatter_mean

            # BEVFORMER_SCA_SCATTER=kernel: the 3 gathers, 2 adds and the 1 / cameras multiply in one program
            slots = scatter_mean(stacked, [ctx[f"sca_g{k}"] for k in range(3)], ctx["sca_scale"])
            free(stacked)
            y = _linear_to_l1(self.output_proj, slots) if kconf.ENC_L1_ALL else self.output_proj(slots)
            free(slots)
            return y
        acc = None
        for k in range(C.SCA_MAX_CAMS_PER_QUERY):
            if self.tile_gather:
                g = tile_gather(ctx[f"sca_g{k}"], stacked)                   # [1, 1, Q, 256] TILE
            else:
                g = gather_rows(ctx[f"sca_g{k}"], stacked, sentinel=SCA_ROWS)    # [1, 1, Q, 256] TILE
            if acc is None:
                acc = g
            else:
                s = ttnn.add(acc, g, memory_config=ttnn.DRAM_MEMORY_CONFIG)
                free(acc, g)
                acc = s
        free(stacked)
        slots = ttnn.multiply(acc, ctx["sca_scale"], memory_config=ttnn.DRAM_MEMORY_CONFIG)
        free(acc)
        y = _linear_to_l1(self.output_proj, slots) if kconf.ENC_L1_ALL else self.output_proj(slots)
        free(slots)
        return y

    def linears(self) -> List[RowLinear]:
        return [self.sampling, self.attention, self.output_proj]


class TtEncoder:
    """BEV queries + rotated history + 3 encoder layers (module docstring). ``__call__(ctx, values, taps=None)``
    returns ``bev_embed`` ``[1, 1, Q, 256]`` TILE bf16; ``ctx`` maps the RT-dev names to device tensors."""

    def __init__(self, params: BEVFormerParams, *, precision: Any = DEFAULT_PRECISION, msda: str = "grid",
                 msda_px: bool = False, msda_tile: bool = False, msda_skip: bool = False, msda_loc: str = "heads",
                 msda_softmax: str = "op", sca_gather: str = "rm", bev_mm: str = "default",
                 tsa_pairs: bool = False, sca_value_dram: bool = False, vin_shared: bool = False,
                 vp_l1: bool = False, msda_mk: str = "off"):
        if msda_mk not in ("off", "sca", "tsa", "all"):
            raise ValueError(f"msda_mk must be off, sca, tsa or all, got {msda_mk!r}")
        self.vin_shared = bool(vin_shared)
        _VP_L1["on"] = bool(vp_l1)
        if bev_mm not in ("default", "tuned", "fast"):
            raise ValueError(f"bev_mm must be 'default', 'tuned' or 'fast', got {bev_mm!r}")
        _TUNED["on"] = bev_mm in ("tuned", "fast")
        from .linconv import BEV_FAST

        BEV_FAST["on"] = bev_mm == "fast"
        if sca_gather not in ("rm", "tile"):
            raise ValueError(f"sca_gather must be 'rm' or 'tile', got {sca_gather!r}")
        if msda_loc not in ("heads", "compact"):
            raise ValueError(f"msda_loc must be 'heads' or 'compact', got {msda_loc!r}")
        if msda_softmax not in ("op", "kernel"):
            raise ValueError(f"msda_softmax must be 'op' or 'kernel', got {msda_softmax!r}")
        sm = msda_softmax == "kernel" and not msda_skip
        self.precision = precision
        self.msda_mode = msda
        modes = site_modes(msda)
        self.tsa_px = bool(msda_px) and modes["tsa"] == "kernel"
        self.sca_px = bool(msda_px) and modes["sca"] == "kernel"
        self.grid_msda = GridSampleMSDA()
        l0, l2, ln = params.can_bus_mlp()
        self.cb0 = _lin(l0, "can_bus_mlp.0", precision, activation="relu", out="float32")
        self.cb2 = _lin(l2, "can_bus_mlp.2", precision, activation="relu", out="float32")
        self.cb_norm = TtLayerNorm(ln, precision=precision)
        self.bev_embedding = _Const(params.bev_embedding().reshape(1, 1, Q, E))
        self.bev_pos = _Const(params.bev_pos().reshape(1, 1, Q, E))
        self.tsa_base = _Const(tsa_base_px() if self.tsa_px else tsa_base(), "float32")
        # BEVFORMER_MSDA_LOC=compact (pixel units only): the sampling projections write only the 2P real columns of
        # each head (tsa 16 heads x 8, sca 8 heads x 16) and msda_k reads that merged map: no head split, 4x / 2x
        # fewer columns through the base / shift adds and the typecast; the same values element for element
        compact = msda_loc == "compact"
        self.tsa_hs = 2 * C.TSA_POINTS if compact and self.tsa_px else 0
        self.sca_hs = 2 * C.SCA_POINTS if compact and self.sca_px else 0
        nb = C.TSA_QUEUE * C.NUM_HEADS
        self.tsa_base_c = _Const(np.ascontiguousarray(np.tile(tsa_base_px()[..., :self.tsa_hs], (1, 1, 1, nb))),
                                 "float32") if self.tsa_hs else None
        # pixel units: tsa_shift (2 shift use, grid units) x size / 2 x 2^16; sca_base (2 ref - 1) -> (ref size - 0.5) 2^16
        sx = np.zeros((1, 1, 1, 32), np.float32)
        sx[..., 0:2 * C.TSA_POINTS:2], sx[..., 1:2 * C.TSA_POINTS:2] = C.BEV_W / 2.0 * PX_SCALE, C.BEV_H / 2.0 * PX_SCALE
        self.shift_scale = _Const(sx, "float32")
        from .msda_mk import base_affine

        self.sca_aff = _Const(base_affine(C.FEAT_HW, C.SCA_POINTS), "float32")
        self.layers = []
        for i in range(C.ENC_LAYERS):
            lp = params.encoder_layer(i)
            pre = f"encoder.layers.{i}"
            self.layers.append({
                "tsa": _TSA(params, i, precision, modes["tsa"], self.grid_msda, px=msda_px, tile=msda_tile,
                            skip=msda_skip, compact=bool(self.tsa_hs), sm=sm, pairs=tsa_pairs,
                            mk=msda_mk in ("tsa", "all")),
                "sca": _SCA(params, i, precision, modes["sca"], self.grid_msda, px=msda_px, skip=msda_skip,
                            compact=bool(self.sca_hs), sm=sm, tile_gather=sca_gather == "tile",
                            value_dram=sca_value_dram, mk=msda_mk in ("sca", "all")),
                "ffn1": _lin(lp["ffn1"], f"{pre}.ffn.0", precision, activation="relu"),
                "ffn2": _lin(lp["ffn2"], f"{pre}.ffn.1", precision),
                "norms": [TtLayerNorm(n, precision=precision) for n in lp["norms"]],
            })

    def bev_queries(self, can_bus: Any):
        """``can_bus`` ``[1, 1, 1, 18]`` fp32 -> ``bev_embedding + can_bus_mlp(can_bus)`` ``[1, 1, Q, 256]`` bf16."""
        import ttnn

        h = self.cb0(can_bus)
        h2 = self.cb2(h)
        free(h)
        cb = self.cb_norm(h2)
        free(h2)
        cbb = ttnn.typecast(cb, ttnn.bfloat16)
        free(cb)
        q = ttnn.add(self.bev_embedding(can_bus.device()), cbb, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        free(cbb)
        return q

    @staticmethod
    def _blend(a, b, use, omu):
        """``use * a + (1 - use) * b`` (use exactly 0 or 1)."""
        import ttnn

        x = ttnn.multiply(a, use, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        y = ttnn.multiply(b, omu, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        out = ttnn.add(x, y, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        free(x, y)
        return out

    def __call__(self, ctx: Dict[str, Any], values: List[Any], taps: Optional[Dict[str, Any]] = None,
                 use_const: Optional[int] = None):
        """``use_const`` 1 (the ``bev_prev`` variant): ``use_prev_bev`` is known to be 1, so the TSA queue values
        are ``v0 = rot(prev_bev)``, ``v1 = bev_queries`` without the blend ops (``x * 1 + y * 0`` returns ``x``)."""
        import ttnn

        dev = ctx["can_bus"].device()
        use, omu = ctx["use"], ctx["omu"]
        bev_q = self.bev_queries(ctx["can_bus"])
        rot = gather_rows(ctx["rot_index"], ctx["prev_bev"], sentinel=Q)          # [1, 1, Q, 256] TILE
        bev_pos, base = self.bev_pos(dev), self.tsa_base(dev)
        shift = ctx["tsa_shift"]
        if self.tsa_px:
            shift = ttnn.multiply(ctx["tsa_shift"], self.shift_scale(dev), memory_config=ttnn.DRAM_MEMORY_CONFIG)
        if self.tsa_hs:
            from .msda import compact_cols

            sc = compact_cols(shift, self.tsa_hs, blocks=C.TSA_QUEUE * C.NUM_HEADS)
            free(shift)
            shift = sc
            base = self.tsa_base_c(dev)
        from . import kconf

        sca_mk = all(L["sca"].mk is not None for L in self.layers) and all(int(v.shape[-1]) == 64 for v in values)
        if self.sca_px and self.sca_hs == 16 and sca_mk and kconf.SCA_BASE_K:
            ctx = dict(ctx, sca_base_aff=self.sca_aff(dev))     # msda_mk computes base * S + O (BEVFORMER_SCA_BASE)
        elif self.sca_px:
            from .msda_k import loc_consts

            s_row, o_row = loc_consts(dev, C.FEAT_HW, C.SCA_POINTS)
            sb = ttnn.multiply(ctx["sca_base"], s_row, memory_config=ttnn.DRAM_MEMORY_CONFIG)
            ctx = dict(ctx, sca_base_px=ttnn.add(sb, o_row, memory_config=ttnn.DRAM_MEMORY_CONFIG))
            free(sb)
            if self.sca_hs:
                from .msda import compact_cols

                ctx["sca_base_c"] = compact_cols(ctx["sca_base_px"], self.sca_hs, reps=C.NUM_HEADS)
        if taps is not None:
            taps["enc.bev_queries"] = bev_q
            taps["enc.prev_bev_rot"] = rot
        def kept(t) -> bool:
            return t is bev_q or (taps is not None and any(t is v for v in taps.values()))

        query = bev_q
        vin = None
        if use_const == 1 and self.vin_shared:      # TSA_VIN=shared: [rot; bev_q] concatenated once per frame
            vin = ttnn.concat([rot, bev_q], dim=0, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        for i, L in enumerate(self.layers):
            if use_const == 1:
                v0, v1 = rot, bev_q
            else:
                v0 = self._blend(rot, query, use, omu)
                v1 = bev_q if i == 0 else self._blend(bev_q, query, use, omu)    # layer 0: query == bev_q
            t = L["tsa"](query, v0, v1, bev_pos, base, shift, vin=vin)
            if v0 is not rot:
                free(v0)
            if v1 is not bev_q:
                free(v1)
            x = L["norms"][0](t, query, l1=L["sca"].tile_gather)          # the SCA gather reads it in L1
            free(t)
            if not kept(query):
                free(query)
            s = L["sca"](x, values[i], ctx)
            y = L["norms"][1](s, x)
            free(s, x)
            # BEVFORMER_ENC_L1 (D20 round 2): the [22528, 512] FFN hidden map in L1 (no DRAM round trip); =all
            # also the FFN output
            h = _linear_to_l1(L["ffn1"], y) if kconf.FFN_L1 else L["ffn1"](y)
            f = _linear_to_l1(L["ffn2"], h) if kconf.ENC_L1_ALL else L["ffn2"](h)
            free(h)
            query = L["norms"][2](f, y)
            free(f, y)
            if taps is not None:
                taps[f"encoder.layer{i}"] = query
        if vin is not None:
            free(vin)
        if taps is None:
            free(rot, bev_q)
        if shift is not ctx["tsa_shift"]:
            free(shift)
        if "sca_base_px" in ctx:
            free(ctx["sca_base_px"])
        if "sca_base_c" in ctx:
            free(ctx["sca_base_c"])
        return query

    def linears(self) -> List[Any]:
        out = [self.cb0, self.cb2]
        for L in self.layers:
            out += L["tsa"].linears() + L["sca"].linears() + [L["ffn1"], L["ffn2"]]
        return out

    def release(self) -> None:
        for lin in self.linears():
            lin.release()
        for c in (self.bev_embedding, self.bev_pos, self.tsa_base, self.shift_scale, self.tsa_base_c, self.sca_aff):
            if c is not None:
                c.release()
        self.cb_norm.release()
        self.grid_msda.release()
        for L in self.layers:
            for n in L["norms"]:
                n.release()
            for site in (L["tsa"], L["sca"]):
                if site.mk is not None:
                    site.mk.release()