| """DRAM-sharded matmul for the ViT shapes (M=192).""" |
| import sys |
| from common import * |
|
|
| dev = open_dev() |
| torch.manual_seed(0) |
| L1 = ttnn.L1_MEMORY_CONFIG |
| nb = dev.dram_grid_size().x |
| print("dram banks", dev.dram_grid_size()) |
| SHAPES = {"qkv": (1280, 4608), "proj": (1536, 1280), "fc1": (1280, 5120), "fc2": (5120, 1280)} |
| for name, (K, N) in [(n, SHAPES[n]) for n in (sys.argv[1:] or SHAPES)]: |
| x_t = torch.randn(1, 1, 192, K); w_t = torch.randn(K, N) / K ** 0.5 |
| ref = x_t @ w_t |
| Np = -(-N // (32 * nb)) * 32 * nb |
| w_pad = torch.nn.functional.pad(w_t, (0, Np - N)) |
| for wd, wn in [(ttnn.bfloat16, "bf16"), (ttnn.bfloat8_b, "bfp8")]: |
| grid = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(nb - 1, 0))}) |
| wmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, |
| ttnn.ShardSpec(grid, [K, Np // nb], ttnn.ShardOrientation.ROW_MAJOR)) |
| w = ttnn.from_torch(w_pad, dtype=wd, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wmc) |
| for ncores in [8, 10, 16, 20, 32, 40, 64, 80]: |
| Kt = K // 32 |
| if Kt % ncores or (Np // 32) % ncores: |
| continue |
| core_grid = ttnn.num_cores_to_corerangeset(ncores, ttnn.CoreCoord(12, 10), row_wise=True) |
| in0mc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, |
| ttnn.ShardSpec(core_grid, [192, K // ncores], ttnn.ShardOrientation.ROW_MAJOR)) |
| x = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=in0mc) |
| outmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1) |
| for ibw_div in [1, 2, 4]: |
| ibw = Kt // ncores // ibw_div |
| if ibw < 1 or (Kt // ncores) % ibw_div: |
| continue |
| pc = ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig( |
| in0_block_w=ibw, per_core_M=6, per_core_N=Np // 32 // ncores, fused_activation=None) |
| for fid in ["HiFi2", "LoFi"]: |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=getattr(ttnn.MathFidelity, fid), |
| fp32_dest_acc_en=False, packer_l1_acc=True) |
| fn = lambda pc=pc, ck=ck, x=x: ttnn.matmul(x, w, program_config=pc, memory_config=outmc, compute_kernel_config=ck) |
| tag = f"{name} w={wn} cores{ncores} ibw{ibw} {fid}" |
| try: |
| o = fn(); p = pcc(ttnn.to_torch(o)[..., :N], ref); ttnn.deallocate(o) |
| us = time_op(dev, fn) |
| print(f"{tag:45s} {us:8.1f} us pcc={p:.5f}", flush=True) |
| except Exception as e: |
| print(f"{tag:45s} FAIL {str(e).splitlines()[0][:150]}", flush=True) |
| ttnn.deallocate(x) |
| ttnn.deallocate(w) |
| ttnn.close_device(dev) |
|
|
