| |
| """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 |
| 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() |
| |
| |
| |
| self.attn_l1 = bool(attn_l1) |
| self.qk_cfg = bool(qk_cfg) |
| 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") |
| |
| |
| 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)] |
|
|
| |
| 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 |
|
|
| |
| 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: |
| |
| |
| |
| 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) |
| 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: |
| 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) |
| 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) |
| |
| lazy = self.msda_mode == "kernel" |
|
|
| from . import kconf |
|
|
| def make_value(L): |
| if kconf.DEC_VALUE_FUSED and lazy and self.tile: |
| |
| |
| 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() |
|
|
