File size: 1,233 Bytes
be62f78 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | # SPDX-License-Identifier: Apache-2.0
"""The swept 2-D multicast configs fit the 12x10 ETH grid and fall back to the auto config on the 11x10 WORKER grid
(republish 2026-10-11: ``dispatch="worker"`` failed at capture with "Num output blocks along x (12) must be smaller
than or equal to the number of columns in compute grid (11)")."""
from types import SimpleNamespace
from tt_diffusion_planner.tt.layers import DEC_MM_CONFIGS, DEC_ROWS, KCAT_CONFIGS, fit_pcm, mcast_fits
ETH, WORKER = SimpleNamespace(x=12, y=10), SimpleNamespace(x=11, y=10)
def test_every_swept_config_fits_the_eth_grid():
for rows in (DEC_ROWS, 32, 64, 96, 128, 192):
for (k, n, _ps), (tr, pcm, pcn, _kb, _sw) in DEC_MM_CONFIGS.items():
assert mcast_fits(tr, fit_pcm(tr, pcm, rows, ETH), pcn, rows, n, ETH), (rows, k, n)
for (rows, _ktp, n), (tr, pcm, pcn, _kb, _sw) in KCAT_CONFIGS.items():
assert mcast_fits(tr, pcm, pcn, rows, n, ETH), (rows, n)
def test_the_qkv_config_does_not_fit_the_worker_grid():
tr, pcm, pcn, _kb, _sw = KCAT_CONFIGS[(352, 25, 768)]
assert not mcast_fits(tr, pcm, pcn, 352, 768, WORKER) # 24 N tiles / 2 = 12 blocks on x > 11
assert mcast_fits(False, 2, 1, 352, 256, WORKER)
|