| |
| """TT-friendly equivalent forms of the deployed graph, proven on the CPU before device time is spent (the reference |
| bundles' ``test_fused_host.py`` discipline, REFERENCE_PATTERNS.md 2.12). Each one is exact in real arithmetic (most |
| are exact in float32 too, as noted) and differs from :mod:`.model` only by float32 summation order; |
| ``tests/test_reference_rewrites_host.py`` checks them against the literal modules with the real weights. |
| |
| | form | replaces | why (PLAN.md 2.5) | |
| |---|---|---| |
| | :func:`pfn_tile_form` | PFN layer 1 on ``concat[h, tile(max h)]`` | ``relu(max_slots(h @ W1a) + max_slots(h) @ W1b + b1)``: no (N, 20, 64) concat, the per-pillar term added after the slot max (exact: ``relu(. + c)`` is monotonic) | |
| | :func:`pad_pillar_slots` | 20 point slots per pillar | 32 slots (one tile row block per pillar): slots 20-31 repeat slot 0, so both slot maxima are unchanged (max is idempotent; zero rows there would add ``relu(b0)`` to full pillars) | |
| | :func:`k_split_weights` | ``shared_conv(concat[d0, d1, d2])`` | ``sum_i conv(d_i, W_i)``: the 56 MB concat never exists (PR:P14, S:transfusion:276) | |
| | :func:`deblock0_as_linear` | Conv 2x2 stride 2 (deblock0) | space-to-depth + linear 256 -> 128 (S:transfusion:271), in case ``conv2d`` loses at 384^2 | |
| | :func:`deconv_as_linear` | ConvTranspose 2x2 s2 (deblock2) | linear + depth-to-space (the k = 2 device form is ``conv_transpose2d``, PR:P5; this is its oracle) | |
| | :func:`local_max_masks` | MaxPool p0 + zero border + equality | C22's masked chain ``heat * (m_pass + (heat == maxpool3x3_p1(heat)) * m_eq)`` (exact, PR:P6) | |
| | :func:`query_pos_from_index` | ``bev_pos`` gather | ``(p % 192 + 0.5, p // 192 + 0.5)`` in float32 (bit-exact; never a bf16 table: S:transfusion:241) | |
| | class table (``HeadParams.class_table``) | Conv1d of a one-hot | a (5, 128) row gather (bit-exact) | |
| | :func:`cross_attention_kv_tables` | ``(lidar + pe(bev_pos)) @ Wk + bk`` | ``lidar @ Wk + CK`` with the constant ``CK = pe(bev_pos) @ Wk + bk`` (same for V; S:transfusion:279) | |
| | :func:`padded_head_attention` | 8 heads x 16 | heads zero-padded to 32 in the projection weights (``ttaw.ops.attention.pad_head_columns`` / ``pad_head_rows``) with ``scale = 1/4`` (SDPA rejects D = 16; exact: PR:P7) | |
| | :func:`merged_prediction_head_weights` | six Conv1d 128 -> 64 + ReLU -> C heads | one 128 -> 384 + ReLU and a block-diagonal 384 -> 15 (padded to 32) | |
| |
| numpy weight transforms; torch only inside the ``*_torch`` forwards (imported lazily). |
| """ |
| from __future__ import annotations |
|
|
| from typing import Dict, List, Sequence, Tuple |
|
|
| import numpy as np |
|
|
| from ..ttaw.ops.attention import pad_head_columns, pad_head_rows |
| from .weights import AttentionParams, HeadParams, Layer |
|
|
| __all__ = ["pfn_tile_form", "pad_pillar_slots", "PILLAR_TILE_SLOTS", "PFN_IN_PADDED", "k_split_weights", |
| "k_split_torch", "deblock0_as_linear", "space_to_depth_nhwc", "deconv_as_linear", "depth_to_space_nhwc", |
| "local_max_masks", "local_max_masked_torch", "query_pos_from_index", "cross_attention_kv_tables", |
| "padded_head_attention", "pad_head_dim_rows", "merged_prediction_head_weights", "merged_heads_torch"] |
|
|
| PILLAR_TILE_SLOTS = 32 |
| PFN_IN_PADDED = 16 |
|
|
|
|
| def _t(a): |
| """float32 torch tensor owning a copy (the ONNX-derived arrays may be read-only).""" |
| import torch |
|
|
| return torch.from_numpy(np.array(a, dtype=np.float32, copy=True)) |
|
|
|
|
| |
|
|
| def pad_pillar_slots(features: np.ndarray, slots: int = PILLAR_TILE_SLOTS, width: int = PFN_IN_PADDED) -> np.ndarray: |
| """(P, 20, F) decorated features -> (P, slots, width): features zero-padded to ``width`` columns, slots |
| ``20..slots-1`` copies of slot 0 (a real point of every kept pillar).""" |
| f = np.asarray(features, dtype=np.float32) |
| P, K, F = f.shape |
| if slots < K or width < F: |
| raise ValueError(f"cannot pad ({K}, {F}) to ({slots}, {width})") |
| out = np.zeros((P, slots, width), dtype=np.float32) |
| out[:, :K, :F] = f |
| out[:, K:, :F] = f[:, :1, :] |
| return out |
|
|
|
|
| def pfn_tile_form(features, w0, b0, w1a, w1b, b1): |
| """The device form of the PFN (torch tensors): ``h = relu(x @ w0 + b0)``, ``m = max_slots h``, |
| ``y = relu(max_slots(h @ w1a) + m @ w1b + b1)`` -> (P, 64). ``x`` may carry padded columns (zero rows in ``w0``) |
| and repeated slots (:func:`pad_pillar_slots`).""" |
| import torch |
|
|
| h = torch.relu(features @ w0 + b0) |
| m = h.amax(dim=1) |
| z = (h @ w1a).amax(dim=1) |
| return torch.relu(z + m @ w1b + b1) |
|
|
|
|
| |
|
|
| def k_split_weights(shared: Layer, split: Sequence[int]) -> List[np.ndarray]: |
| """The shared conv weight cut along input channels into the deblocks' channel groups: |
| ``conv(concat[d_i], W) == sum_i conv(d_i, W_i)``.""" |
| if sum(split) != shared.in_channels: |
| raise ValueError(f"split {list(split)} does not cover {shared.in_channels} input channels") |
| bounds = np.cumsum([0] + list(split)) |
| return [np.ascontiguousarray(shared.weight[:, a:b]) for a, b in zip(bounds[:-1], bounds[1:])] |
|
|
|
|
| def k_split_torch(deblocks, shared: Layer): |
| """``sum_i conv(d_i, W_i) + b`` (no activation: the TransFusion shared conv has none), torch NCHW.""" |
| import torch.nn.functional as F |
|
|
| parts = k_split_weights(shared, [int(d.shape[1]) for d in deblocks]) |
| acc = None |
| for d, w in zip(deblocks, parts): |
| y = F.conv2d(d, _t(w), None, stride=shared.stride, padding=shared.padding) |
| acc = y if acc is None else acc + y |
| return acc + _t(shared.bias)[None, :, None, None] |
|
|
|
|
| def space_to_depth_nhwc(x: np.ndarray, s: int) -> np.ndarray: |
| """(H, W, C) -> (H/s, W/s, s*s*C), channel ``(a * s + b) * C + c`` = pixel ``(s*i + a, s*j + b)``.""" |
| H, W, C = x.shape |
| return np.asarray(x).reshape(H // s, s, W // s, s, C).transpose(0, 2, 1, 3, 4).reshape(H // s, W // s, s * s * C) |
|
|
|
|
| def deblock0_as_linear(layer: Layer) -> Tuple[np.ndarray, np.ndarray]: |
| """Conv k x k stride k (no padding) as ``space_to_depth(x) @ Wl + b``: ``Wl[(a * k + b) * Cin + ci, co] = |
| W[co, ci, a, b]``.""" |
| if layer.kind != "conv" or layer.kernel != layer.stride or layer.padding != 0: |
| raise ValueError("deblock0_as_linear needs a conv with kernel == stride and no padding") |
| w = layer.weight |
| cout, cin, k, _ = w.shape |
| wl = np.ascontiguousarray(w.transpose(2, 3, 1, 0).reshape(k * k * cin, cout), dtype=np.float32) |
| return wl, layer.bias.astype(np.float32) |
|
|
|
|
| def deconv_as_linear(layer: Layer) -> Tuple[np.ndarray, np.ndarray]: |
| """ConvTranspose with kernel == stride s (no padding) as a linear ``Cin -> s*s*Cout`` on NHWC rows: column |
| ``(a * s + b) * Cout + co`` holds ``W[ci, co, a, b]``, bias tiled the same way.""" |
| if layer.kind != "conv_transpose" or layer.kernel != layer.stride or layer.padding != 0: |
| raise ValueError("deconv_as_linear needs a ConvTranspose with kernel == stride and no padding") |
| w = layer.weight |
| cin, cout, s, _ = w.shape |
| wl = np.ascontiguousarray(w.transpose(0, 2, 3, 1).reshape(cin, s * s * cout), dtype=np.float32) |
| return wl, np.tile(layer.bias, s * s).astype(np.float32) |
|
|
|
|
| def depth_to_space_nhwc(y: np.ndarray, H: int, W: int, s: int, cout: int) -> np.ndarray: |
| """Rows ``(i, j)`` of ``y`` (H*W, s*s*Cout) -> the (H*s, W*s, Cout) map: ``out[s*i + a, s*j + b, co] = |
| y[i*W + j, (a*s + b)*Cout + co]``.""" |
| return np.asarray(y).reshape(H, W, s, s, cout).transpose(0, 2, 1, 3, 4).reshape(H * s, W * s, cout) |
|
|
|
|
| |
|
|
| def local_max_masks(height: int, width: int, classes: int, kernel: int = 3) -> Tuple[np.ndarray, np.ndarray]: |
| """The constant masks of C22's chain for TransFusion's rule (``MaxPool`` k3 s1 **p0** + zero border, every class): |
| ``m_eq`` = 1 on interior cells (the pooled window is fully inside), 0 on the ``kernel // 2`` border (always |
| suppressed); ``m_pass`` = 0 everywhere (no class bypasses the local max). (classes, H, W) float32.""" |
| r = kernel // 2 |
| m_eq = np.zeros((classes, height, width), np.float32) |
| m_eq[:, r:height - r, r:width - r] = 1.0 |
| return np.zeros_like(m_eq), m_eq |
|
|
|
|
| def local_max_masked_torch(heat, m_pass, m_eq, kernel: int = 3): |
| """``heat * (m_pass + (heat == maxpool_k_s1_p(k//2)(heat)) * m_eq)`` (torch NCHW; the pool's padding is -inf).""" |
| import torch.nn.functional as F |
|
|
| pooled = F.max_pool2d(heat, kernel, 1, kernel // 2) |
| return heat * (_t(m_pass)[None] + (heat == pooled).to(heat.dtype) * _t(m_eq)[None]) |
|
|
|
|
| |
|
|
| def query_pos_from_index(pos: np.ndarray, feat_w: int) -> np.ndarray: |
| """(K,) flat BEV cells -> (K, 2) float32 ``(x + 0.5, y + 0.5)`` (exact: integers below 2^24 plus 0.5).""" |
| p = np.asarray(pos, dtype=np.int64) |
| return np.stack([(p % feat_w).astype(np.float32) + np.float32(0.5), |
| (p // feat_w).astype(np.float32) + np.float32(0.5)], axis=-1) |
|
|
|
|
| def cross_attention_kv_tables(attn: AttentionParams, key_pos_embed: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: |
| """``CK = pe(bev_pos) @ Wk + bk`` and ``CV = pe(bev_pos) @ Wv + bv`` (36864, 128), float64 then one rounding: |
| ``K = lidar^T @ Wk + CK``, ``V = lidar^T @ Wv + CV`` (the 0.6 GMAC position branch leaves the per-frame graph).""" |
| kpe = np.asarray(key_pos_embed, dtype=np.float64) |
| ck = kpe @ attn.wk.astype(np.float64) + attn.bk.astype(np.float64) |
| cv = kpe @ attn.wv.astype(np.float64) + attn.bv.astype(np.float64) |
| return ck.astype(np.float32), cv.astype(np.float32) |
|
|
|
|
| def pad_head_dim_rows(table: np.ndarray, num_heads: int, head_dim: int, padded: int) -> np.ndarray: |
| """(L, H*D) -> (L, H*padded) with each head's columns zero-padded (a constant K / V table in the padded layout).""" |
| t = np.asarray(table, dtype=np.float32).reshape(-1, num_heads, head_dim) |
| out = np.zeros((t.shape[0], num_heads, padded), np.float32) |
| out[..., :head_dim] = t |
| return out.reshape(t.shape[0], num_heads * padded) |
|
|
|
|
| def padded_head_attention(q_in, k_in, v_in, attn: AttentionParams, *, k_const=None, v_const=None, |
| padded: int = 32): |
| """Multi-head attention with the head dim zero-padded to ``padded`` in the weights (torch, (L, 128) inputs): |
| ``softmax(scale * Qp Kp^T) Vp`` per head with ``scale = 1 / q_divisor`` (the logical ``1/sqrt(16)``), then the |
| row-padded output projection. ``k_const`` / ``v_const`` (L_k, 128): the constant K / V tables |
| (:func:`cross_attention_kv_tables`), whose bias they include (then ``k_in @ Wk`` alone is projected).""" |
| import torch |
|
|
| H, D = attn.num_heads, attn.head_dim |
| wq, bq = pad_head_columns(attn.wq, attn.bq, H, D, padded) |
| wk, bk = pad_head_columns(attn.wk, attn.bk, H, D, padded) |
| wv, bv = pad_head_columns(attn.wv, attn.bv, H, D, padded) |
| wo = pad_head_rows(attn.wo, H, D, padded) |
| q = q_in @ _t(wq) + _t(bq) |
| if k_const is None: |
| k = k_in @ _t(wk) + _t(bk) |
| v = v_in @ _t(wv) + _t(bv) |
| else: |
| k = k_in @ _t(wk) + _t(pad_head_dim_rows(k_const, H, D, padded)) |
| v = v_in @ _t(wv) + _t(pad_head_dim_rows(v_const, H, D, padded)) |
| lq, lk = q.shape[0], k.shape[0] |
| qh = q.reshape(lq, H, padded).transpose(0, 1) |
| kh = k.reshape(lk, H, padded).transpose(0, 1) |
| vh = v.reshape(lk, H, padded).transpose(0, 1) |
| o = torch.softmax((qh @ kh.transpose(1, 2)) * attn.scale, dim=-1) @ vh |
| return o.transpose(0, 1).reshape(lq, H * padded) @ _t(wo) + _t(attn.bo) |
|
|
|
|
| |
|
|
| def merged_prediction_head_weights(head: HeadParams, *, pad_to: int = 32 |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[Tuple[str, int, int]]]: |
| """The six prediction heads as two matmuls on the (K, 128) decoder rows: ``hidden = relu(x @ Wh + bh)`` with |
| ``Wh`` (128, 6*64), and ``out = hidden @ Wo + bo`` with the block-diagonal ``Wo`` (6*64, 15 -> ``pad_to``). |
| Returns (Wh, bh, Wo, bo, slices); ``slices`` lists ``(name, start, stop)`` of each head's output channels.""" |
| wh = np.concatenate([hid.weight[:, :, 0].T for _, hid, _ in head.prediction_heads], axis=1).astype(np.float32) |
| bh = np.concatenate([hid.bias for _, hid, _ in head.prediction_heads]).astype(np.float32) |
| rows = sum(hid.out_channels for _, hid, _ in head.prediction_heads) |
| cols = sum(fin.out_channels for _, _, fin in head.prediction_heads) |
| width = -(-cols // pad_to) * pad_to |
| wo = np.zeros((rows, width), np.float32) |
| bo = np.zeros(width, np.float32) |
| slices, r, c = [], 0, 0 |
| for name, hid, fin in head.prediction_heads: |
| wo[r:r + hid.out_channels, c:c + fin.out_channels] = fin.weight[:, :, 0].T |
| bo[c:c + fin.out_channels] = fin.bias |
| slices.append((name, c, c + fin.out_channels)) |
| r += hid.out_channels |
| c += fin.out_channels |
| return wh, bh, wo, bo, slices |
|
|
|
|
| def merged_heads_torch(x, merged) -> Dict[str, object]: |
| """Apply :func:`merged_prediction_head_weights` to the (K, 128) decoder output -> {name: (1, C, K)}.""" |
| import torch |
|
|
| wh, bh, wo, bo, slices = merged |
| y = torch.relu(x @ _t(wh) + _t(bh)) @ _t(wo) + _t(bo) |
| return {name: y[:, a:b].T[None] for name, a, b in slices} |
|
|
