File size: 3,027 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """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)
|