| """Round 15: the MANO-head K/V projection (library 1D matmul, [192, 1280] x [1280, 6144] bfp8, 12x10, in0_block_w 5, |
| per_core_N 2, HiFi2 + fp32 acc) with its weights in DRAM (served) vs L1 interleaved: is it weight-stream-bound?""" |
| import sys |
| from pathlib import Path |
| from common import * |
| dev = open_dev() |
| torch.manual_seed(0) |
| L1, DR = ttnn.L1_MEMORY_CONFIG, ttnn.DRAM_MEMORY_CONFIG |
| G = dev.compute_with_storage_grid_size() |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, |
| packer_l1_acc=True) |
| try: |
| N = 6144 |
| x = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1) |
| wt = torch.randn(1280, N) / 1280 ** 0.5 |
| for ibw in (5, 10, 4): |
| pc = ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig( |
| compute_with_storage_grid_size=ttnn.CoreCoord(G.x, G.y), in0_block_w=ibw, out_subblock_h=1, out_subblock_w=2, |
| per_core_M=6, per_core_N=2, fuse_batch=True, fused_activation=None, mcast_in0=True) |
| for wloc in ("dram", "l1"): |
| w = ttnn.from_torch(wt, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR if wloc == "dram" else L1) |
| fn = lambda w=w: ttnn.linear(x, w, program_config=pc, memory_config=L1, compute_kernel_config=ck) |
| print(f"[kvw] ibw {ibw} w={wloc:4s} {time_op(dev, fn):7.2f} us", flush=True) |
| ttnn.deallocate(w) |
| finally: |
| ttnn.close_device(dev) |
|
|
