# 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()