code/tt_bevformer/tt/decoder.py

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