| |
| """Attention score scale + additive mask as one ``ttnn.generic_op`` (``ATTN_SMASK``, OPT round 2 item 4b). |
| |
| ``scale_mask(s, scale, mask)``: fp32 TILE scores ``[1, H, Sq, Sk]`` and a ``[1, 1, Sq, Sk]`` mask (bf16 or fp32, |
| broadcast over the heads) -> ``s * scale + mask`` fp32, the two stock ``binary_ng`` programs of |
| ``tt/attention.py`` (``ttnn.multiply(s, scale)``, ``ttnn.add(., mask)``) in one pass over the scores |
| (``kernels/smask_*.cpp``): bit-identical, half the DRAM traffic and one program less per attention. A bf16 mask |
| sends the stock fp32 + bf16 add down binary_ng's FPU path, which truncates the scores to TF32 (measured); the |
| kernel reproduces that truncation (the mask holds 0 / -inf only, so the add itself is exact). |
| The mask-plane tiles are split in contiguous ranges over the grid; each core streams the H score tiles of its mask |
| tiles. |
| """ |
| from __future__ import annotations |
|
|
| import os |
| import struct |
| from typing import Any |
|
|
| __all__ = ["scale_mask", "supported"] |
|
|
| _KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kernels") |
| TB = 4096 |
| N_RT = 2 |
|
|
|
|
| def supported(s: Any, mask: Any) -> bool: |
| import ttnn |
|
|
| try: |
| ss, ms = list(s.padded_shape), list(mask.padded_shape) |
| return (hasattr(ttnn, "generic_op") and s.dtype == ttnn.float32 and s.layout == ttnn.TILE_LAYOUT |
| and mask.layout == ttnn.TILE_LAYOUT and mask.dtype in (ttnn.bfloat16, ttnn.float32) |
| and not s.is_sharded() and not mask.is_sharded() and len(ss) == 4 and len(ms) == 4 |
| and ss[0] == 1 and ss[1] % 2 == 0 and ms[0] == 1 and ms[1] == 1 and ss[2:] == ms[2:]) |
| except Exception: |
| return False |
|
|
|
|
| def scale_mask(s: Any, scale: float, mask: Any): |
| import ttnn |
|
|
| dev = s.device() |
| _, H, Sq, Sk = list(s.padded_shape) |
| plane = (Sq // 32) * (Sk // 32) |
| out = ttnn.allocate_tensor_on_device(s.shape, ttnn.float32, ttnn.TILE_LAYOUT, dev, ttnn.DRAM_MEMORY_CONFIG) |
| g = dev.compute_with_storage_grid_size() |
| n = min(plane, g.x * g.y) |
| cs = [(i % g.x, i // g.x) for i in range(n)] |
| full, rem = divmod(n, g.x) |
| rs = [] |
| if full: |
| rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(g.x - 1, full - 1))) |
| if rem: |
| rs.append(ttnn.CoreRange(ttnn.CoreCoord(0, full), ttnn.CoreCoord(rem - 1, full))) |
| crs = ttnn.CoreRangeSet(set(rs)) |
| base, extra = divmod(plane, n) |
| rd, wr, cp = ttnn.RuntimeArgs(), ttnn.RuntimeArgs(), ttnn.RuntimeArgs() |
| m0 = 0 |
| for i, (cx, cy) in enumerate(cs): |
| k = base + (1 if i < extra else 0) |
| rd[cx][cy] = [m0, k] |
| wr[cx][cy] = [m0, k] |
| cp[cx][cy] = [k, 0] |
| m0 += k |
| mbf = mask.dtype == ttnn.bfloat16 |
| tbm = 2048 if mbf else 4096 |
|
|
| def cb(idx, pages, dt, tb): |
| return ttnn.CBDescriptor(total_size=pages * tb, core_ranges=crs, format_descriptors=[ |
| ttnn.CBFormatDescriptor(buffer_index=idx, data_format=dt, page_size=tb)]) |
|
|
| cbs = [cb(0, 2 * H, ttnn.float32, TB), cb(1, 2, mask.dtype, tbm), cb(2, 1, ttnn.float32, TB), |
| cb(16, 4, ttnn.float32, TB)] |
| um = [ttnn.UnpackToDestMode.Default] * 64 |
| um[0] = um[2] = ttnn.UnpackToDestMode.UnpackToDestFp32 |
| if not mbf: |
| um[1] = ttnn.UnpackToDestMode.UnpackToDestFp32 |
| ccfg = ttnn.ComputeConfigDescriptor(math_fidelity=ttnn.MathFidelity.HiFi4, fp32_dest_acc_en=True, |
| math_approx_mode=False) |
| ccfg.unpack_to_dest_mode = um |
| bits = int.from_bytes(struct.pack("<f", float(scale)), "little") |
| SRC = ttnn.KernelDescriptor.SourceType.FILE_PATH |
|
|
| def acc(t): |
| return list(ttnn.TensorAccessorArgs(t).get_compile_time_args()) |
|
|
| def kd(name, ct, rt, common, config): |
| return ttnn.KernelDescriptor(kernel_source=os.path.join(_KDIR, name), source_type=SRC, core_ranges=crs, |
| compile_time_args=ct, defines=[], runtime_args=rt, common_runtime_args=common, |
| config=config) |
|
|
| ks = [kd("smask_reader.cpp", [H, plane, bits, N_RT, tbm] + acc(s) + acc(mask), rd, |
| [s.buffer_address(), mask.buffer_address()], ttnn.ReaderConfigDescriptor()), |
| kd("smask_writer.cpp", [H, plane, N_RT] + acc(out), wr, [out.buffer_address()], |
| ttnn.WriterConfigDescriptor()), |
| kd("smask_compute.cpp", [H, N_RT, int(mbf)], cp, [], ccfg)] |
| ttnn.generic_op([s, mask, out], ttnn.ProgramDescriptor(kernels=ks, semaphores=[], cbs=cbs)) |
| return out |
|
|
