# SPDX-License-Identifier: Apache-2.0 """The DETR decoder and the heads on the device: 6 x (self-attention -> LN -> deformable cross-attention over the BEV -> LN -> FFN -> LN -> box refinement) and the last layer's classification / box outputs. The 900 object queries run as 928 rows (``tt.params.DEC_ROWS``: a whole number of tiles; the 28 pad rows are masked as keys in the self-attention and never read back). Per layer: - self-attention: ``[x + pos | x] @ W_qkv`` (Q, K from ``x + pos``, V from ``x``; the Q columns carry the ``1 / sqrt(32)`` scale) -> ``nlp_create_qkv_heads`` -> scores ``Q K^T`` in fp32 (HiFi4, fp32 accumulation) + the pad-key mask (-inf) -> softmax (fp32) -> bf16 -> ``@ V`` -> ``nlp_concat_heads`` -> ``out_proj`` -> LN(+ x). The fused SDPA is avoided on purpose: its bf16 scores cost accuracy (probe P7: PCC 0.9997 vs fp64) and 900 x 928 x 8 is small; - cross-attention (``CustomMSDeformableAttention``, 150 x 150, 4 points): the value maps of the layer (``value_proj(bev_embed)``, 8 heads) are built from ``bev_embed`` at the start of the decoder; the grid base ``2 ref - 1`` comes from the replicated reference (``tt.params.REF_COLS``) by two element-wise ops; - the box refinement in fp32: ``ref = sigmoid(t + inverse_sigmoid(ref))`` on the replicated columns (``t`` = the ``reg_branches[i]`` output written into the same columns, ``tt.params.reg_branch``), ``inverse_sigmoid(x) = log(x') - log(1 - x')``, ``x' = clamp(x, 1e-5, 1 - 1e-5)`` (DH ``decoder.py:24-37``; the head's mmdet form gives the same value on the clamped range, so the last layer's coords are ``ref * range + min`` of the refined reference). Outputs of the production variant: the LAST layer's class logits ``[928, 32]`` (10 used) and coords ``[928, 32]`` (``tt.params.coords_from_device`` picks the network's 10 columns), fp32. With ``taps`` every layer's output, class logits and coords are kept too. """ from __future__ import annotations import math 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.precision import Precision from .encoder import DEFAULT_PRECISION, TtLayerNorm, _lin from .image import _Const from .msda import GridSampleMSDA, attend, free, site_modes, value_maps from .params import DEC_ROWS, PX_SCALE, BEVFormerParams, LinearP __all__ = ["TtDecoder"] E = C.EMBED class _Layer: def __init__(self, params: BEVFormerParams, i: int, precision: Any, px: bool = False): lp = params.decoder_layer(i) pre = f"decoder.layers.{i}" mha = lp["mha"] w = mha["qkv"].w.copy() b = mha["qkv"].b.copy() scale = np.float32(1.0 / math.sqrt(C.HEAD_DIM)) w[:, :E] *= scale # Q columns pre-scaled b[:E] *= scale self.qkv = _lin(LinearP(w, b), f"{pre}.self_attn.qkv", precision) self.attn_out = _lin(mha["out_proj"], f"{pre}.self_attn.out_proj", precision) cr = lp["cross"] self.value_proj = _lin(cr["value_proj"], f"{pre}.cross.value_proj", precision) self.sampling = _lin(cr["sampling_px" if px else "sampling"], f"{pre}.cross.sampling", precision, out="float32") self.attention = _lin(cr["attention"], f"{pre}.cross.attention", precision, out="float32") self.cross_out = _lin(cr["output_proj"], f"{pre}.cross.output_proj", precision) self.ffn1 = _lin(lp["ffn1"], f"{pre}.ffn.0", precision, activation="relu") self.ffn2 = _lin(lp["ffn2"], f"{pre}.ffn.1", precision) self.norms = [TtLayerNorm(n, precision=precision) for n in lp["norms"]] r0, r1, r2 = params.reg_branch(i) self.reg = [_lin(r0, f"reg_branches.{i}.0", precision, activation="relu"), _lin(r1, f"reg_branches.{i}.2", precision, activation="relu"), _lin(r2, f"reg_branches.{i}.4", precision, out="float32")] (c0, c1, c2), (n0, n1) = params.cls_branch(i) self.cls = [_lin(c0, f"cls_branches.{i}.0", precision), _lin(c1, f"cls_branches.{i}.3", precision), _lin(c2, f"cls_branches.{i}.6", precision, out="float32")] self.cls_norms = [TtLayerNorm(n0, precision=precision), TtLayerNorm(n1, precision=precision)] def linears(self) -> List[RowLinear]: return [self.qkv, self.attn_out, self.value_proj, self.sampling, self.attention, self.cross_out, self.ffn1, self.ffn2] + self.reg + self.cls def norms_all(self) -> List[TtLayerNorm]: return self.norms + self.cls_norms class TtDecoder: """6 decoder layers + heads (module docstring). ``__call__(bev_embed, taps=None) -> (cls, coords)`` of the last layer (fp32 ``[1, 1, 928, 32]``).""" def __init__(self, params: BEVFormerParams, *, precision: Any = DEFAULT_PRECISION, msda: str = "grid", msda_px: bool = False, msda_tile: bool = False, msda_skip: bool = False, attn_l1: bool = False, qk_cfg: bool = False, attn: str = "matmul"): self.cfg = Precision.parse(precision).compute_kernel_config() # DEC_ATTN_L1: the [8, 928, 928] scores / masked scores / probabilities (27.5 MB fp32 each) and the bf16 # probabilities live in L1 instead of DRAM (the same ops and configs: bit-identical; 0.78 -> 0.53 ms per layer, # logs/bevformer/opt_r3/job_p2.out) self.attn_l1 = bool(attn_l1) self.qk_cfg = bool(qk_cfg) # DEC_QK=grid: QK^T on a 10 x 10 grid (K = one tile: same bits) self.msda_mode = site_modes(msda)["dec"] self.px = bool(msda_px) and self.msda_mode == "kernel" self.tile = bool(msda_tile) and self.msda_mode == "kernel" self.skip = bool(msda_skip) self.grid_msda = GridSampleMSDA() pos, x, ref = params.query_init() self.pos = _Const(pos.reshape(1, 1, DEC_ROWS, E)) self.x0 = _Const(x.reshape(1, 1, DEC_ROWS, E)) self.ref0 = _Const(ref, "float32") self.mask = _Const(params.mha_mask(), "float32") # DEC_ATTN=sdpa (PLAN D20): the self-attention core as ttnn's flash SDPA (bf16 scores, streaming softmax, # heads concatenated by the op) with the pad-key mask in bf16; not bit-identical if attn not in ("matmul", "sdpa"): raise ValueError(f"dec attn must be 'matmul' or 'sdpa', got {attn!r}") self.sdpa = attn == "sdpa" self.mask16 = _Const(params.mha_mask(), "bfloat16") a, b, m = params.coords_affine() self.coords_a, self.coords_b, self.coords_m = (_Const(v, "float32") for v in (a, b, m)) self.layers = [_Layer(params, i, precision, px=self.px) for i in range(C.DEC_LAYERS)] # ---- pieces --------------------------------------------------------------------------------------------------- def self_attention(self, L: _Layer, x, pos): import ttnn dev = x.device() xp = ttnn.add(x, pos, memory_config=ttnn.DRAM_MEMORY_CONFIG) cat = ttnn.concat([xp, x], dim=3, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(xp) qkv = L.qkv(cat) free(cat) q, k, v = ttnn.experimental.nlp_create_qkv_heads(qkv, num_heads=C.NUM_HEADS, num_kv_heads=C.NUM_HEADS, transpose_k_heads=False, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(qkv) if self.sdpa: from ..ttaw.ops.attention import sdpa # the 1 / sqrt(32) scale is folded into the Q projection (params): scale 1 merged = sdpa(q, k, v, scale=1.0, attn_mask=self.mask16(dev), concat_heads=True, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(q, k, v) y = L.attn_out(merged) free(merged) return y amc = ttnn.L1_MEMORY_CONFIG if self.attn_l1 else ttnn.DRAM_MEMORY_CONFIG kw = {} if self.qk_cfg and int(q.shape[-1]) == 32: # BEVFORMER_DEC_QK=grid: the default 2D config runs on 4 x 4 cores (per_core 8 x 8 tiles); 3 x 3 tiles # per core on a 10 x 10 grid. K is one tile, so every config accumulates the same products in the same # order: bit-identical (logs/bevformer/opt_r6/qk_probe.out: 0.115 -> 0.056 ms) mt = -(-int(q.shape[-2]) // 32) g = dev.compute_with_storage_grid_size() gx, gy = -(-mt // 3), -(-mt // 3) if gx <= int(g.x) and gy <= int(g.y): kw["program_config"] = ttnn.MatmulMultiCoreReuseMultiCastProgramConfig( compute_with_storage_grid_size=(gx, gy), in0_block_w=1, out_subblock_h=3, out_subblock_w=1, out_block_h=3, out_block_w=3, per_core_M=3, per_core_N=3, transpose_mcast=False, fused_activation=None, fuse_batch=False) scores = ttnn.matmul(q, k, transpose_b=True, compute_kernel_config=self.cfg, dtype=ttnn.float32, memory_config=amc, **kw) # [1, 8, 928, 928] fp32 free(q, k) masked = ttnn.add(scores, self.mask(dev), memory_config=amc) free(scores) probs = ttnn.softmax(masked, dim=-1, numeric_stable=True, compute_kernel_config=self.cfg, memory_config=amc) free(masked) pb = ttnn.typecast(probs, ttnn.bfloat16, memory_config=amc) if self.attn_l1 else ttnn.typecast( probs, ttnn.bfloat16) free(probs) o = ttnn.matmul(pb, v, compute_kernel_config=self.cfg, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(pb, v) merged = ttnn.experimental.nlp_concat_heads(o, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(o) y = L.attn_out(merged) free(merged) return y def cross_attention(self, L: _Layer, x, pos, value, ref): import ttnn qc = ttnn.add(x, pos, memory_config=ttnn.DRAM_MEMORY_CONFIG) if self.px: # (ref size - 0.5) 2^16, size 150 x 150 (the x / y scales agree) r2 = ttnn.multiply(ref, float(C.BEV_W * PX_SCALE), memory_config=ttnn.DRAM_MEMORY_CONFIG) base = ttnn.subtract(r2, 0.5 * PX_SCALE, memory_config=ttnn.DRAM_MEMORY_CONFIG) else: r2 = ttnn.multiply(ref, 2.0, memory_config=ttnn.DRAM_MEMORY_CONFIG) base = ttnn.subtract(r2, 1.0, memory_config=ttnn.DRAM_MEMORY_CONFIG) # 2 ref - 1 (replicated) free(r2) cols, logits = L.sampling(qc), L.attention(qc) free(qc) merged = attend(value, cols, logits, num_heads=C.NUM_HEADS, points=C.DEC_POINTS, base=base, 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) free(base, cols, logits) y = L.cross_out(merged) free(merged) return y @staticmethod def refine(t, ref): """``sigmoid(t + log(x') - log(1 - x'))``, ``x' = clamp(ref, eps, 1 - eps)`` (fp32).""" import ttnn eps = float(C.INV_SIG_EPS) xc = ttnn.clamp(ref, min=eps, max=1.0 - eps, memory_config=ttnn.DRAM_MEMORY_CONFIG) l1 = ttnn.log(xc, memory_config=ttnn.DRAM_MEMORY_CONFIG) om = ttnn.rsub(xc, 1.0, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(xc) l2 = ttnn.log(om, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(om) inv = ttnn.subtract(l1, l2, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(l1, l2) z = ttnn.add(t, inv, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(inv) out = ttnn.sigmoid(z, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(z) return out def heads(self, L: _Layer, x, ref_new, t): """Classification logits and coords of one layer (fp32 ``[1, 1, 928, 32]`` each).""" import ttnn dev = x.device() h = x for lin, ln in zip(L.cls[:2], L.cls_norms): a = lin(h) if h is not x: free(h) n = ln(a) free(a) h = ttnn.relu(n, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(n) cls = L.cls[2](h) free(h) ca = ttnn.multiply(ref_new, self.coords_a(dev), memory_config=ttnn.DRAM_MEMORY_CONFIG) cb = ttnn.add(ca, self.coords_b(dev), memory_config=ttnn.DRAM_MEMORY_CONFIG) free(ca) tm = ttnn.multiply(t, self.coords_m(dev), memory_config=ttnn.DRAM_MEMORY_CONFIG) coords = ttnn.add(cb, tm, memory_config=ttnn.DRAM_MEMORY_CONFIG) free(cb, tm) return cls, coords def __call__(self, bev, taps: Optional[Dict[str, Any]] = None): dev = bev.device() pos = self.pos(dev) x = self.x0(dev) ref = self.ref0(dev) # msda_k gathers from L1: each layer's value maps (11.5 MB) are made right before that layer, in L1 lazy = self.msda_mode == "kernel" from . import kconf def make_value(L): if kconf.DEC_VALUE_FUSED and lazy and self.tile: # BEVFORMER_DEC_VALUE=fused (D20 round 3): the projection straight to L1, read by msda_k head- # interleaved (tile (row / 32) * H + head): no nlp_create_qkv_heads, no DRAM round trip from .encoder import _linear_to_l1 return _linear_to_l1(L.value_proj, bev) vp = L.value_proj(bev) v = value_maps(vp, num_heads=C.NUM_HEADS, hw=(C.BEV_H, C.BEV_W), l1=lazy, tile=self.tile) free(vp) return v values = [None if lazy else make_value(L) for L in self.layers] last = len(self.layers) - 1 out = None for i, L in enumerate(self.layers): a = self.self_attention(L, x, pos) x1 = L.norms[0](a, x) free(a) if values[i] is None: values[i] = make_value(L) c = self.cross_attention(L, x1, pos, values[i], ref) free(values[i]) x2 = L.norms[1](c, x1) free(c, x1) h = L.ffn1(x2) f = L.ffn2(h) free(h) x3 = L.norms[2](f, x2) free(f, x2) if i > 0 and (taps is None or not any(x is v for v in taps.values())): free(x) x = x3 t = x for lin in L.reg: t2 = lin(t) if t is not x: free(t) t = t2 ref_new = self.refine(t, ref) if taps is not None or i == last: cls, coords = self.heads(L, x, ref_new, t) if taps is not None: taps[f"decoder.layer{i}"] = x taps[f"head.cls{i}"] = cls taps[f"head.coords{i}"] = coords if i == last: out = (cls, coords) free(t) if i > 0 or ref is not self.ref0(dev): free(ref) ref = ref_new if taps is None: free(x) free(ref) return out def release(self) -> None: for L in self.layers: for lin in L.linears(): lin.release() for n in L.norms_all(): n.release() for c in (self.pos, self.x0, self.ref0, self.mask, self.mask16, self.coords_a, self.coords_b, self.coords_m): c.release() self.grid_msda.release()