# SPDX-License-Identifier: Apache-2.0 """MB3: achieved DRAM bandwidth of the deepseek_v3_b1 ``DRAMStreamingMatmul`` micro-op at DiT sizes. Device probe for docs/research/device-stack-smoke.md (target 4). The in-tree test (``models/demos/deepseek_v3_b1/tests/unit_tests/test_dram_streaming_matmul.py``) checks PCC only and never prints a bandwidth, so this script re-creates its tensor layout (in0 replicated on the 8 DRAM-adjacent worker cores, in1 WIDTH_SHARDED over the 8 DRAM banks with column-major tile shuffle, tensor-backed working buffer) and times ``[1, K] x [K, N]`` streams with **bfp8** weights for the GR00T DiT shapes (arch-n15/16/17.md: hidden 1536 for N1.5, 2048 for N1.6/N1.7, FF = 4x). Bandwidth = bytes of in1 streamed from DRAM / device time, where the device time is measured with ``num_loop_iters`` in-kernel repetitions (so dispatch and host overhead amortise away) after a warm-up call that pays the JIT compile. bfp8_b tiles are 1088 B per 32x32 tile (1024 B mantissas + 64 B shared exponents). Run through the device lock:: bin/with-device.sh python models/experimental/gr00t/benchmarks/mb3_dram_streaming.py """ from __future__ import annotations import os import sys import time import torch import ttnn from models.demos.deepseek_v3_b1.micro_ops.dram_streaming_matmul.op import DRAMStreamingMatmul from models.demos.deepseek_v3_b1.tests.unit_tests.test_dram_streaming_matmul import ( pad_to_dram_banks, shuffle_tensor_tiles, ) TILE = 32 BFP8_TILE_BYTES = 1024 + 64 # (K, N) pairs: N1.5 DiT (1536) and N1.6/N1.7 DiT (2048) square and 4x FF projections, plus the # DeepSeek shape the in-tree test uses as a sanity anchor. SHAPES = [(1536, 1536), (1536, 6144), (6144, 1536), (2048, 2048), (2048, 8192), (8192, 2048), (7168, 2048)] LOOP_ITERS = 200 def _pcc(a: torch.Tensor, b: torch.Tensor) -> float: a = a.flatten().float() b = b.flatten().float() return torch.corrcoef(torch.stack([a, b]))[0, 1].item() def bench_shape(device, k: int, n: int, in1_dtype, fidelity, fp32_acc: bool) -> None: m = 1 in0_tile = ttnn.Tile([m, TILE]) out_tile = ttnn.Tile([m, TILE]) compute_cores = device.get_optimal_dram_bank_to_logical_worker_assignment(ttnn.NOC.NOC_0) num_cores = len(compute_cores) num_banks = device.dram_grid_size().x assert num_cores == num_banks, (num_cores, num_banks) n_padded = pad_to_dram_banks(n, TILE, TILE * num_banks) per_core_n = n_padded // num_banks core_grid = ttnn.CoreRangeSet( [ttnn.CoreRange(ttnn.CoreCoord(c.x, c.y), ttnn.CoreCoord(c.x, c.y)) for c in compute_cores] ) torch.manual_seed(0) in0 = torch.randn(1, 1, m, k).bfloat16().float() in1 = torch.randn(1, 1, k, n_padded).bfloat16().float() in0_spec = ttnn.ShardSpec(core_grid, [m, k], ttnn.ShardOrientation.ROW_MAJOR) in0_t = ttnn.from_torch( in0.repeat(1, 1, num_cores, 1), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, in0_spec), tile=in0_tile, ) dram_grid = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(num_banks - 1, 0))]) in1_spec = ttnn.ShardSpec(dram_grid, [k, per_core_n], ttnn.ShardOrientation.ROW_MAJOR) in1_t = ttnn.from_torch( shuffle_tensor_tiles(in1, TILE, num_banks), dtype=in1_dtype, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, in1_spec), ) out_spec = ttnn.ShardSpec(core_grid, (m, per_core_n), ttnn.ShardOrientation.ROW_MAJOR) out_t = ttnn.from_torch( torch.zeros(1, 1, m, n_padded).bfloat16().float(), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, out_spec), tile=out_tile, ) kt = k // TILE subblock_k = kt // 4 if kt % 4 == 0 else kt // 2 num_in1_buffers = 3 wb_cols = subblock_k * num_in1_buffers * TILE wb_spec = ttnn.ShardSpec(core_grid, (TILE, wb_cols), ttnn.ShardOrientation.ROW_MAJOR) wb_t = ttnn.from_torch( torch.zeros(1, 1, TILE, wb_cols * num_cores).bfloat16().float(), dtype=in1_dtype, layout=ttnn.TILE_LAYOUT, device=device, memory_config=ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, wb_spec), tile=ttnn.Tile([TILE, TILE]), ) def run(iters: int) -> float: t0 = time.perf_counter() DRAMStreamingMatmul.op( in0_t, in1_t, out_t, fp32_dest_acc_en=fp32_acc, math_fidelity=fidelity, math_approx_mode=False, subblock_k=subblock_k, fused_activation=None, num_loop_iters=iters, working_buf_tensor=wb_t, ) ttnn.synchronize_device(device) return time.perf_counter() - t0 t_cold = run(1) # JIT compile + first run t_warm1 = min(run(1) for _ in range(5)) # single-shot latency, program cache warm # num_loop_iters is a compile-time arg, so each distinct loop count JIT-compiles once: warm both variants, # then take the difference of two loop counts so fixed dispatch/launch overhead cancels. run(LOOP_ITERS) run(2 * LOOP_ITERS) t_loop_a = min(run(LOOP_ITERS) for _ in range(3)) t_loop_b = min(run(2 * LOOP_ITERS) for _ in range(3)) per_iter = (t_loop_b - t_loop_a) / LOOP_ITERS t_loop = t_loop_a tile_bytes = ( BFP8_TILE_BYTES if in1_dtype == ttnn.bfloat8_b else (TILE * TILE * 2 if in1_dtype == ttnn.bfloat16 else 512 + 64) ) in1_bytes = kt * (n_padded // TILE) * tile_bytes bw_loop = in1_bytes / per_iter / 1e9 bw_single = in1_bytes / t_warm1 / 1e9 ref = DRAMStreamingMatmul.golden(in0, in1, None) got = ttnn.to_torch(out_t) pcc = _pcc(ref, got) dtype_name = str(in1_dtype).split(".")[-1] print( f"RESULT K={k} N={n}(pad {n_padded}) in1={dtype_name} fid={str(fidelity).split('.')[-1]} fp32acc={fp32_acc} " f"subblock_k={subblock_k} per_core_N_tiles={per_core_n // TILE} in1_bytes={in1_bytes / 1e6:.2f}MB | " f"cold(compile+run)={t_cold * 1e3:.1f}ms single_warm={t_warm1 * 1e6:.0f}us ({bw_single:.0f} GB/s incl. dispatch) | " f"looped x{LOOP_ITERS}={t_loop * 1e3:.2f}ms, d/d_iter={per_iter * 1e6:.1f}us/iter = {bw_loop:.0f} GB/s | PCC={pcc:.5f}" ) for t in (in0_t, in1_t, out_t, wb_t): ttnn.deallocate(t) def main() -> int: print( f"ENV slow_dispatch={os.environ.get('TT_METAL_SLOW_DISPATCH_MODE')} TT_METAL_CACHE={os.environ.get('TT_METAL_CACHE', '')}" ) device = ttnn.open_device(device_id=0) try: grid = device.compute_with_storage_grid_size() print(f"FACT compute_grid={grid.x}x{grid.y} dram_banks={device.dram_grid_size().x}") for k, n in SHAPES: try: bench_shape(device, k, n, ttnn.bfloat8_b, ttnn.MathFidelity.LoFi, False) except Exception as exc: # noqa: BLE001 - record, keep going print(f"RESULT K={k} N={n} in1=bfloat8_b FAILED: {exc!r}"[:600]) # anchors: the in-tree test's bfp4 config and a bf16 weight stream, one DiT shape each for dtype, name in ((ttnn.bfloat4_b, "bfloat4_b"), (ttnn.bfloat16, "bfloat16")): try: bench_shape(device, 2048, 8192, dtype, ttnn.MathFidelity.LoFi, False) except Exception as exc: # noqa: BLE001 print(f"RESULT K=2048 N=8192 in1={name} FAILED: {exc!r}"[:600]) # does the op accept anything but LoFi? (custom_mm is documented LoFi-only) try: bench_shape(device, 2048, 2048, ttnn.bfloat8_b, ttnn.MathFidelity.HiFi2, False) except Exception as exc: # noqa: BLE001 print(f"RESULT K=2048 N=2048 in1=bfloat8_b HiFi2 FAILED: {exc!r}"[:600]) finally: ttnn.close_device(device) return 0 if __name__ == "__main__": sys.exit(main())