File size: 1,548 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 | """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)
|