| |
| """The fp32 matmul attention of the fusion transformer and the decoder (C20 ``attention_matmul`` math) with stock-op |
| program choices tuned for the planner's shapes (``ATTN_FAST``, OPT round 1 item 4a). |
| |
| ``softmax(scale * Q K^T + mask) V`` as in ``ttaw.ops.attention.attention_matmul``, with: |
| |
| - (``mode >= 1``) ``P V`` on a ``MatmulMultiCoreReuseProgramConfig`` with one output tile per core (8 heads x |
| Sq / 32 = 88 cores at Sq = 352) instead of the auto config's 11-18 cores; the K-sum in one block |
| (``in0_block_w`` = Sk / 32): bit-identical to the auto config in the device sweep (73 -> 21 us self, 71 -> 32 us |
| cross, 81 -> 42 us fusion; ``logs/diffusion-planner/opt_r1/attn_sweep.log``); |
| - (``mode == 2``) the scale applied to Q (``[1, 8, Sq, 32]``, 11x smaller than the scores) before ``Q K^T`` |
| instead of to the scores: one small multiply instead of the large one; a precision change (rounding order); |
| - (``mode == 3``) the scale as an SFPU pre-activation of the mask add (``input_tensor_a_activations``): one |
| program instead of two over the scores (75 -> 42 us), but NOT bit-identical (final_x0 rel 4e-3; e2e ego mean |
| 0.143 -> 0.173 m over the 99 scenes): measured and rejected (OPT_REPORT round 1). |
| - ``Q K^T`` keeps the auto config (the reuse configs of the sweep were slower). |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Optional |
|
|
| __all__ = ["attention_matmul", "pv_config"] |
|
|
| TILE = 32 |
| |
| MAX_KB_TILES = 18 |
|
|
|
|
| def pv_config(device: Any, heads: int, sq: int, sk: int): |
| """``MatmulMultiCoreReuseProgramConfig`` for ``P [1, H, Sq, Sk] @ V [1, H, Sk, 32]``: one ``[32, 32]`` output |
| tile per core (``per_core_M = 1`` when ``H * Sq / 32`` fits the grid, else 2).""" |
| import ttnn |
|
|
| grid = device.compute_with_storage_grid_size() |
| mt, kt = sq // TILE, sk // TILE |
| pcm = 1 if heads * mt <= grid.x * grid.y else 2 |
| kb = max(d for d in range(1, min(kt, MAX_KB_TILES) + 1) if kt % d == 0) |
| return ttnn.MatmulMultiCoreReuseProgramConfig(compute_with_storage_grid_size=grid, in0_block_w=kb, |
| out_subblock_h=1, out_subblock_w=1, per_core_M=pcm, per_core_N=1) |
|
|
|
|
| def attention_matmul(q: Any, k: Any, v: Any, *, scale: float, attn_mask: Any = None, |
| compute_kernel_config: Any = None, mode: int = 1, pv_pc: Optional[Any] = None, |
| smask: bool = False, smsm: int = 0): |
| """``q`` ``[1, H, Sq, D]``, ``k`` / ``v`` ``[1, H, Sk, D]`` (tile-aligned ``Sk``) -> ``[1, H, Sq, D]``. |
| ``mode=0``: exactly ``ttaw.ops.attention.attention_matmul``. ``smask`` (``ATTN_SMASK``, mode 1 with a mask): |
| the scale multiply and the mask add as one bit-identical generic_op (``tt/smask_kernel.py``); ``smsm`` |
| (``ATTN_SMSM`` 1 / 2, with ``smask``): the scale, the mask and the softmax as one generic_op |
| (``tt/smsm_kernel.py``; 2 = the scale multiply by an immediate).""" |
| import ttnn |
|
|
| from ..ttaw.ops import attention as A |
| from ..ttaw.precision import compute_kernel_config as ckc |
|
|
| cfg = compute_kernel_config if compute_kernel_config is not None else ckc("HiFi4", fp32_acc=True) |
| dev = q.device() if callable(getattr(q, "device", None)) else None |
| if mode and (dev is None or not hasattr(dev, "compute_with_storage_grid_size") |
| or not hasattr(getattr(ttnn, "UnaryOpType", None), "MUL_UNARY_SFPU")): |
| mode = 0 |
| if not mode: |
| return A.attention_matmul(q, k, v, scale=scale, attn_mask=attn_mask, compute_kernel_config=cfg) |
| sk = int(k.shape[-2]) |
| if sk % TILE: |
| raise ValueError(f"attention_matmul needs a tile-aligned Sk (got {sk})") |
| if mode == 2: |
| q = ttnn.multiply(q, float(scale)) |
| scores = ttnn.matmul(q, k, transpose_b=True, compute_kernel_config=cfg) |
| fused = False |
| if smsm and smask and mode == 1 and attn_mask is not None: |
| from .smsm_kernel import scale_mask_softmax |
| from .smsm_kernel import supported as smsm_supported |
|
|
| if smsm_supported(scores, attn_mask): |
| probs = scale_mask_softmax(scores, scale, attn_mask, scale_mode=int(smsm) - 1) |
| if pv_pc is None: |
| pv_pc = pv_config(dev, int(q.shape[1]), int(q.shape[-2]), sk) |
| return ttnn.matmul(probs, v, compute_kernel_config=cfg, program_config=pv_pc) |
| if smask and mode == 1 and attn_mask is not None: |
| from .smask_kernel import scale_mask, supported |
|
|
| if supported(scores, attn_mask): |
| scores, fused = scale_mask(scores, scale, attn_mask), True |
| if fused: |
| pass |
| elif mode == 3 and attn_mask is not None: |
| act = [ttnn.UnaryWithParam(ttnn.UnaryOpType.MUL_UNARY_SFPU, float(scale))] |
| scores = ttnn.add(scores, attn_mask, input_tensor_a_activations=act) |
| else: |
| if mode != 2: |
| scores = ttnn.multiply(scores, float(scale)) |
| if attn_mask is not None: |
| scores = ttnn.add(scores, attn_mask) |
| probs = ttnn.softmax(scores, dim=-1, numeric_stable=True, compute_kernel_config=cfg) |
| if pv_pc is None: |
| pv_pc = pv_config(dev, int(q.shape[1]), int(q.shape[-2]), sk) |
| return ttnn.matmul(probs, v, compute_kernel_config=cfg, program_config=pv_pc) |
|
|
