code/models/tt_dit/utils/matmul.py

# SPDX-FileCopyrightText: © 2025 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

import math
import os
from typing import NamedTuple

from loguru import logger

import ttnn
from models.common.utility_functions import is_blackhole

# Track unique warning signatures to avoid stdout spam
_warned_matmul_signatures = set()
_warned_1d_matmul_signatures = set()

# ---------------------------------------------------------------------------
# Known-best 1D mcast_in0 configs for specific (M, K, N, grid_x, grid_y).
# Each value is a tuple: (in0_block_w, per_core_N, out_subblock_w).
# Generated by scripts/generate_flux2_1d_mm_sweep_md.py from
# test_1d_matmul_sweep_bh4x8_ring (2026-05-27, flux2_1d_mm_sweep.md).
# ---------------------------------------------------------------------------
grid_1d_configs: dict[tuple[int, int, int, int, int], tuple[int, int, int]] = {
    # Flux2 M=32 shapes — BH 4×8 ring, 11×10 grid (2026-05-27, flux2_1d_mm_sweep.md)
    (32, 256, 6144, 11, 10): (2, 2, 1),  # 11.7 μs,  3.2% FLOP util
    (32, 6144, 6144, 11, 10): (2, 5, 1),  # 197.9 μs, 11.3% FLOP util
    (32, 6144, 4608, 11, 10): (2, 2, 2),  # 149.9 μs,  6.1% FLOP util
    (32, 6144, 2304, 11, 10): (4, 2, 2),  # 76.6 μs,  11.9% FLOP util
    (32, 6144, 1536, 11, 10): (6, 2, 2),  # 51.6 μs,  17.6% FLOP util
    # Flux2 SP1 (1024-res) shapes — BH 4×8 ring, 11×10 grid (2026-06-01, 2048_sp1.md)
    (32, 6144, 9216, 11, 10): (6, 9, 3),  # 290.7 μs
    (32, 6144, 3072, 11, 10): (12, 3, 3),  # 99.5 μs
    (64, 6144, 9216, 11, 10): (8, 5, 1),  # 297.4 μs
    (64, 15360, 1536, 11, 10): (8, 2, 2),  # 122.8 μs
}


# Known best blockings for 8x8 core grid for specific (M, K, N) shapes
# Each value is a tuple: (M_block_size, K_block_size, N_block_size)
grid_88_configs = {
    (32, 2432, 7296): (2, 4, 8),
    (1024, 2432, 3648): (4, 4, 16),
    (352, 2432, 3648): (2, 4, 8),
    (1024, 2432, 1216): (4, 4, 8),
    (352, 2432, 1216): (2, 4, 8),
    (1024, 2432, 4864): (4, 4, 8),
    (1024, 4864, 2432): (4, 4, 16),
    (352, 2432, 4864): (4, 8, 4),
    (352, 4864, 2432): (2, 8, 4),
    (32, 3072, 6144): (2, 8, 8),
    (32, 3072, 3072): (2, 4, 16),
    (11264, 3072, 4608): (16, 8, 4),
    (128, 1536, 4608): (2, 8, 4),
    (11264, 3072, 1536): (8, 8, 8),
    (128, 3072, 768): (2, 8, 4),
    (11264, 3072, 8192): (16, 8, 4),
    (11264, 4096, 3072): (16, 4, 4),
    (128, 1536, 4096): (2, 8, 4),
    (128, 2048, 1536): (2, 4, 8),
    (18944, 5120, 2560): (8, 8, 8),
    (128, 5120, 2560): (4, 16, 2),
    (18944, 5120, 6912): (16, 8, 4),
    (18944, 6912, 5120): (8, 8, 8),
    (32, 2432, 3648): (2, 4, 8),
    (1024, 2432, 1920): (4, 4, 8),
    (352, 2432, 1920): (2, 4, 8),
    (1024, 2560, 608): (4, 4, 4),
    (352, 2560, 608): (2, 5, 4),
    (1024, 2432, 2432): (4, 4, 16),
    (352, 2432, 2432): (2, 4, 16),
    (32, 3072, 1536): (2, 4, 8),
    (5632, 3072, 2304): (8, 12, 4),
    (128, 1536, 2304): (2, 6, 4),
    (5632, 3072, 768): (8, 8, 4),
    (128, 3072, 384): (2, 6, 2),
    (5632, 3072, 4096): (8, 6, 8),
    (5632, 2048, 3072): (8, 8, 8),
    (128, 1536, 2048): (2, 6, 4),
    (128, 1024, 1536): (2, 4, 4),
    (9472, 5120, 1280): (8, 10, 8),
    (128, 5120, 1280): (2, 8, 8),
    (9472, 5120, 3456): (10, 8, 8),
    (9472, 5120, 3840): (10, 8, 8),
    (9472, 3456, 5120): (8, 12, 4),
    (32, 4096, 640): (2, 8, 4),
    (32, 4096, 1280): (2, 16, 2),
    (32, 4096, 5120): (2, 16, 4),
    (32, 5120, 768): (4, 8, 4),
    (32, 5120, 1536): (2, 4, 8),
    (32, 5120, 4096): (2, 16, 4),
    (32, 5120, 8192): (2, 16, 4),
    (32, 5120, 16384): (2, 16, 4),
    (32, 5120, 32768): (4, 32, 4),
    (32, 8192, 5120): (2, 16, 4),
    (512, 4096, 640): (2, 8, 4),
    (512, 4096, 1280): (2, 8, 8),
    (512, 4096, 2560): (2, 8, 16),
    (512, 4096, 5120): (2, 16, 4),
    (512, 5120, 768): (2, 8, 4),
    (512, 5120, 1536): (2, 4, 8),
    (512, 5120, 3072): (2, 8, 16),
    (512, 5120, 4096): (2, 8, 8),
    (512, 5120, 6144): (2, 8, 8),
    (512, 5120, 8192): (2, 16, 8),
    (512, 5120, 16384): (4, 16, 4),
    (512, 5120, 32768): (2, 16, 8),
    (512, 8192, 5120): (4, 16, 4),
    (512, 16384, 5120): (2, 16, 8),
    (512, 32768, 5120): (4, 16, 8),
}


# Known best blockings for 8x9 core grid for specific (M, K, N) shapes
# Each value is a tuple: (M_block_size, K_block_size, N_block_size)
grid_89_configs = {
    (32, 2432, 3648): (2, 4, 8),
    (1024, 2432, 1920): (4, 4, 8),
    (352, 2432, 1920): (2, 4, 4),
    (1024, 2560, 608): (4, 4, 4),
    (352, 2560, 608): (2, 4, 4),
    (1024, 2432, 2432): (4, 4, 16),
    (352, 2432, 2432): (2, 4, 16),
    (32, 3072, 3072): (2, 4, 16),
    (32, 3072, 1536): (2, 4, 8),
    (5632, 3072, 2304): (8, 8, 8),
    (128, 1536, 2304): (2, 8, 2),
    (5632, 3072, 768): (8, 8, 4),
    (128, 3072, 384): (2, 4, 4),
    (5632, 3072, 4096): (8, 8, 8),
    (5632, 2048, 3072): (8, 8, 4),
    (128, 1536, 2048): (4, 4, 4),
    (128, 1024, 1536): (2, 4, 4),
    (9472, 5120, 1280): (8, 8, 8),
    (128, 5120, 1280): (2, 16, 2),
    (9472, 5120, 3456): (8, 8, 8),
    (9472, 3456, 5120): (8, 8, 8),
}

grid_13_9_configs = {
    (9472, 5120, 1280): (8, 8, 8, (2, 2)),
    (128, 5120, 1280): (2, 16, 4, (2, 2)),
    (9472, 5120, 3456): (8, 8, 4, (1, 2)),
    (9472, 3456, 5120): (8, 12, 4, (1, 2)),
}

grid_12_10_configs = {
    (9472, 5120, 1280): (16, 8, 4, (2, 2)),
    (128, 5120, 1280): (1, 16, 8, (1, 2)),
    (9472, 5120, 3456): (16, 8, 4, (1, 2)),
    (9472, 5120, 3840): (16, 8, 4, (1, 2)),
    (9472, 3456, 5120): (8, 4, 8, (1, 2)),
}

grid_11_10_configs = {
    (512, 5120, 2560): (2, 5, 8, (1, 4)),
    (512, 4096, 1024): (2, 8, 4, (1, 4)),
    (512, 4096, 2560): (2, 4, 8, (1, 4)),
    (512, 2560, 4096): (2, 4, 16, (1, 4)),
    (512, 4096, 5120): (2, 8, 8, (1, 4)),
    (512, 5120, 5120): (2, 5, 8, (1, 4)),
    (32, 256, 5120): (1, 2, 8, (1, 4)),
    (32, 5120, 1280): (1, 10, 4, (1, 4)),
    (32, 1280, 30720): (1, 20, 4, (1, 4)),
    (32, 192, 384): (1, 3, 2, (1, 2)),
    (64, 192, 384): (1, 3, 2, (1, 2)),
    (2368, 5120, 64): (4, 8, 2, (2, 2)),
    (2368, 5120, 3840): (4, 8, 12, (1, 4)),
    (2368, 5120, 1280): (8, 4, 4, (1, 4)),
    (2368, 5120, 3456): (8, 2, 12, (2, 2)),
    (2368, 3456, 5120): (4, 4, 8, (1, 4)),
    (9472, 5120, 64): (2, 32, 2, (2, 2)),
    (9472, 5120, 3840): (16, 4, 4, (1, 4)),
    (9472, 5120, 1280): (16, 8, 4, (1, 4)),
    (9472, 5120, 3456): (16, 8, 4, (1, 4)),
    (9472, 3456, 5120): (16, 3, 4, (1, 4)),
    (14400, 384, 1152): (16, 3, 8, (2, 2)),
    (14400, 384, 384): (8, 3, 4, (2, 2)),
    # LTX Gemma encoder + connectors (seq=1024, swept on 110 cores; M_block=4 beats the 8x8x8
    # default by ~10-13% on the FFN/proj shapes; the FE aggregate is compute-bound, ~flat).
    (1024, 3840, 3840): (4, 8, 8, (1, 4)),
    (1024, 4096, 4096): (4, 8, 12, (1, 4)),
    (1024, 3840, 2048): (4, 8, 8, (1, 4)),
    (1024, 4096, 960): (4, 8, 4, (1, 4)),
    (1024, 15360, 960): (4, 8, 4, (1, 4)),
    (1024, 4096, 1024): (4, 8, 12, (1, 4)),
    (1024, 188160, 4096): (4, 8, 12, (1, 4)),
    # FE aggregate after column-parallel sharding on TP=4: video 4096/4=1024, audio 2048/4=512.
    (1024, 188160, 1024): (4, 8, 8, (1, 4)),
    (1024, 188160, 512): (4, 8, 4, (1, 4)),
    (32, 6144, 12288): (1, 8, 24, (1, 4)),
    (128, 32, 32): (4, 1, 1, (4, 1)),
    (128, 64, 512): (1, 2, 2, (1, 2)),
    (512, 64, 256): (2, 2, 1, (2, 1)),
    (512, 3072, 6144): (2, 6, 12, (1, 4)),
    (1024, 32, 128): (3, 1, 1, (3, 1)),
    (1024, 6144, 128): (8, 8, 1, (4, 1)),
    (16384, 512, 64): (6, 8, 2, (2, 2)),
    # flux2 dual-grid sweep rank-1 — 11×10 grid (2026-05-27, flux2_dual_grid_mm_sweep.md)
    (1024, 128, 768): (12, 2, 5, (1, 1)),  # dual-grid sweep 11x10, 9.5 μs, 7.0% FLOP util
    (512, 15360, 768): (4, 15, 6, (2, 1)),  # dual-grid sweep 11x10, 146.2 μs, 27.2% FLOP util
    (1024, 6144, 768): (6, 8, 8, (3, 1)),  # dual-grid sweep 11x10, 80.2 μs, 39.6% FLOP util
    (512, 6144, 768): (2, 8, 3, (1, 3)),  # dual-grid sweep 11x10, 67.4 μs, 23.6% FLOP util
    (1024, 6144, 4608): (4, 6, 8, (4, 1)),  # dual-grid sweep 11x10, 311.5 μs, 61.2% FLOP util
    (512, 6144, 4608): (4, 8, 8, (1, 4)),  # dual-grid sweep 11x10, 244.1 μs, 39.0% FLOP util
    (1024, 6144, 2304): (4, 4, 8, (4, 1)),  # dual-grid sweep 11x10, 162.9 μs, 58.5% FLOP util
    (512, 6144, 2304): (2, 6, 8, (1, 4)),  # dual-grid sweep 11x10, 128.9 μs, 37.0% FLOP util
    # flux2 2048-res sweep rank-1 — 11×10 grid (2026-05-30, 2048.md)
    (4096, 128, 768): (14, 2, 3, (1, 3)),  # 32.3 μs,   2.8% util  x_embedder
    (128, 15360, 768): (2, 16, 3, (1, 3)),  # 132.0 μs,  2.6% util  context_embedder
    (4096, 6144, 2304): (6, 4, 4, (2, 2)),  # 540.5 μs, 24.4% util  single-block to_qkv spatial
    (128, 6144, 2304): (2, 6, 10, (2, 2)),  # 127.4 μs,  3.2% util  single-block to_qkv prompt
    (4096, 6144, 4608): (8, 4, 14, (2, 2)),  # 1054.6 μs,25.0% util  single-block proj_mlp spatial
    (128, 6144, 4608): (6, 12, 8, (2, 2)),  # 238.0 μs,  3.5% util  single-block proj_mlp prompt
    (4096, 6144, 128): (3, 16, 2, (3, 1)),  # 206.8 μs,  6.9% util  proj_out
    # flux2 SP1 (1024-res) sweep rank-1 — 11×10 grid (2026-06-01, 2048_sp1.md)
    (2048, 128, 1536): (6, 2, 5, (3, 1)),  # 23.9 μs   x_embedder / img_in
    (2048, 6144, 128): (16, 8, 2, (2, 2)),  # 94.0 μs   proj_out
    (2048, 6144, 4608): (4, 4, 15, (4, 1)),  # 462.3 μs  proj_mlp spatial
    (2048, 6144, 9216): (8, 8, 10, (2, 2)),  # 972.7 μs  ff1 / qkv spatial
    (64, 6144, 4608): (2, 8, 8, (2, 2)),  # 196.3 μs  proj_mlp prompt
}

grid_12_9_configs = {
    # flux2 dual-grid sweep rank-1 — 12×9 grid (2026-05-27, flux2_dual_grid_mm_sweep.md)
    (1024, 128, 768): (8, 2, 8, (2, 1)),  # dual-grid sweep 12x9, 9.7 μs, 7.0% FLOP util
    (512, 15360, 768): (2, 16, 4, (2, 2)),  # dual-grid sweep 12x9, 120.2 μs, 33.7% FLOP util
    (128, 15360, 768): (2, 16, 4, (2, 2)),  # M_tiles=4 same as (512,…); unswept estimate
    (1024, 6144, 768): (3, 8, 6, (1, 3)),  # 195.3 μs, −25.0% vs baseline — attn to_out (sweep rank-1, 2026-06-11)
    (512, 6144, 768): (2, 8, 8, (2, 2)),  # dual-grid sweep 12x9, 55.9 μs, 29.0% FLOP util
    (1024, 6144, 4608): (3, 8, 8, (1, 4)),  # 432.6 μs, −16.6% vs baseline — ff1 spatial (sweep rank-1, 2026-06-11)
    (1152, 6144, 4608): (
        3,
        8,
        8,
        (1, 4),
    ),  # 463.7 μs, −6.4% vs baseline — ff1 spatial xc-merged (sweep rank-1, 2026-06-11)
    (512, 6144, 4608): (8, 8, 8, (2, 2)),  # 377.8 μs — attn to_qkv TP4_SP8 (sweep rank-1, 2026-06-12)
    (1024, 6144, 2304): (3, 8, 8, (1, 4)),  # 269.9 μs, −21.4% vs baseline — attn to_qkv (sweep rank-1, 2026-06-11)
    (1152, 6144, 2304): (3, 8, 12, (1, 4)),  # 293.8 μs — SNG_attn_to_qkv 1152 tokens (sweep rank-1, 2026-06-21)
    (512, 6144, 2304): (2, 8, 12, (1, 3)),  # dual-grid sweep 12x9, 116.2 μs, 41.8% FLOP util
    (9472, 5120, 1280): (10, 8, 8, (2, 1)),
    (2368, 5120, 1280): (10, 8, 6, (2, 1)),
    (128, 5120, 1280): (1, 16, 8, (1, 2)),
    (9472, 5120, 3456): (9, 5, 12, (1, 2)),
    (2368, 5120, 3456): (7, 5, 12, (1, 2)),
    (9472, 5120, 3840): (7, 5, 16, (1, 2)),
    (2368, 5120, 3840): (7, 5, 16, (1, 2)),
    # flux2 2048-res sweep rank-1 — 12×9 AGMM grid (2026-05-31, 2048.md)
    (4096, 6144, 4608): (11, 8, 8, (1, 4)),  # 1180.7 μs, −12.3% vs baseline — ff1 spatial (sweep rank-1, 2026-06-11)
    (128, 6144, 4608): (2, 8, 8, (2, 2)),  # 358.6 μs, −11.7% vs baseline — ff1 prompt (sweep rank-1, 2026-06-11)
    # x_c_merged_fused: concat([spatial, prompt_sp_sharded], dim=1) → M=4096+128=4224
    # Same M_per_core as (4096,…) on 9-row grid (ceil(132/9)=ceil(128/9)=15); same blocking.
    (4224, 6144, 4608): (
        11,
        8,
        8,
        (1, 4),
    ),  # 1189.1 μs, −9.4% vs baseline — ff1 spatial xc-merged (sweep rank-1, 2026-06-11)
    (4096, 6144, 2304): (11, 8, 8, (1, 4)),  # 722.7 μs, −7.0% vs baseline — attn to_qkv (sweep rank-1, 2026-06-11)
    (4096, 6144, 768): (12, 6, 3, (4, 1)),  # 659.7 μs, −4.9% vs baseline — attn to_out (sweep rank-1, 2026-06-11)
    # M=2048: M_tiles=64, M_per_core=ceil(64/9)=8 — M_block=8 fills each core exactly.
    # Same K/N as (1024, 6144, 4608); reuse that blocking. Unverified estimate.
    (2048, 6144, 4608): (6, 6, 16, (2, 2)),  # 618.7 μs — attn to_qkv TP4_SP8 (sweep rank-1, 2026-06-12)
    (2048, 6144, 9216): (6, 8, 12, (2, 2)),  # 1104.2 μs — ff1 spatial TP4_SP8 (sweep rank-1, 2026-06-12)
    (2048, 6144, 1536): (8, 6, 6, (2, 2)),  # 413.3 μs — attn to_out TP4_SP8 (sweep rank-1, 2026-06-12)
    (512, 6144, 9216): (4, 8, 8, (2, 2)),  # 705.0 μs — ff1 spatial TP4_SP8 (sweep rank-1, 2026-06-12)
    (512, 6144, 1536): (2, 8, 6, (2, 2)),  # 166.7 μs — attn to_out TP4_SP8 (sweep rank-1, 2026-06-12)
    (576, 6144, 9216): (2, 8, 8, (2, 2)),  # 706.0 μs — ff1 xc-merged TP4_SP8 (sweep rank-1, 2026-06-12)
    (64, 6144, 9216): (6, 4, 16, (2, 2)),  # 689.6 μs — ff1 prompt TP4_SP8 (sweep rank-1, 2026-06-12)
    (64, 6144, 4608): (4, 4, 16, (2, 2)),  # 356.9 μs — attn to_qkv prompt TP4_SP8 (sweep rank-1, 2026-06-12)
    (64, 6144, 1536): (2, 6, 6, (2, 2)),  # 144.8 μs — attn to_out prompt TP4_SP8 (sweep rank-1, 2026-06-12)
    (2112, 6144, 9216): (8, 6, 12, (2, 2)),  # 1112.9 μs — ff1 xc-merged TP4_SP8 (sweep rank-1, 2026-06-12)
    (128, 6144, 768): (2, 12, 6, (2, 2)),  # 90.1 μs, −44.4% vs baseline — attn to_out prompt (sweep rank-1, 2026-06-12)
    # flux2 4096-res estimates — 12×9 AGMM grid.
    # AGMM overhead empirically varies with K and M_block (not fixed):
    #   K=4608, M_b=16 → overhead ≈ 848 KB (848,384 B); max CBs = 353 tiles
    #   K=6144, M_b=16 → overhead ≈ 1069 KB (1,094,144 B); max CBs = 233 tiles  ← tightest
    # M_block=16 for M=16384 (ceil(512/12)=43 → capped); subblocks from 2048-res swept entries.
    # K_block and N_block chosen to keep CBs ≤ 224 tiles (safe margin below 233 tile limit).
    (16384, 6144, 4608): (10, 6, 10, (2, 2)),  # 4194.6 μs, −41.1% vs baseline — ff1 spatial (sweep rank-1, 2026-06-11)
    # Prompt-only shapes (M=128 → AGMM path, same 12x9 grid as spatial).
    # Default 1x8x8 (152 tiles) crashes at 4096-res; use K_b=4 to halve in1 tiles.
    (128, 6144, 2304): (
        2,
        4,
        8,
        (2, 2),
    ),  # 189.5 μs, −23.9% vs baseline — attn add_qkv prompt (sweep rank-1, 2026-06-12)
    (16384, 6144, 2304): (12, 6, 8, (2, 2)),  # 2666.7 μs, −47.1% vs baseline — attn to_qkv (sweep rank-1, 2026-06-11)
    (16384, 6144, 768): (16, 8, 3, (4, 1)),  # 2534.4 μs, −16.0% vs baseline — attn to_out (sweep rank-1, 2026-06-11)
    # x_c_merged at 4096-res: M=16384+128=16512 → M_tiles=516; 516%16≠0, use M_block=12 (516%12=0)
    # ff2 (K=4608, overhead ≈ 848 KB for K_per_dev=18): max CBs = 354 tiles; K_b=4 → 224 t ✓
    (16384, 4608, 768): (16, 4, 4, (1, 4)),  # 224 tiles — ff2 spatial                       [UNSWEPT ESTIMATE]
    # x_c_merged at 4096-res: M=16384+128=16512 → M_tiles=516; 516%16≠0, use M_block=12 (516%12=0)
    (16512, 6144, 4608): (
        10,
        6,
        10,
        (2, 2),
    ),  # 4243.7 μs, −39.5% vs baseline — ff1 spatial xc-merged (sweep rank-1, 2026-06-11)
    (16512, 6144, 2304): (12, 2, 8, (4, 1)),  # 176 tiles  (K_b=3→216 reduced to K_b=2)       [UNSWEPT ESTIMATE]
    (16512, 4608, 768): (12, 4, 4, (1, 4)),  # 176 tiles — ff2 merged    [UNSWEPT ESTIMATE]
    # flux2 8192-token shapes (e.g. 4096-res with SP=8, or alt-parallelism).
    # (8192, 6144, 4608) heuristic (16,8,16)→768 tiles crashes at 3,256,832 B (2.07×L1).
    # AGMM overhead scales with M_block; using M_block=8 (half of capped 16) + K_b=2 is safe.
    # M_t=256; 256%8=0 ✓. Subblocks mirror M=16384 entries.
    (8192, 6144, 4608): (6, 6, 16, (2, 2)),  # 2126.3 μs — attn to_qkv TP4_SP8 (sweep rank-1, 2026-06-12)
    (8192, 6144, 2304): (8, 2, 8, (4, 1)),  # 128 tiles — ff1 chunk/ff2 [UNSWEPT ESTIMATE]
    (8192, 6144, 1536): (8, 12, 6, (2, 2)),  # 1540.5 μs — attn to_out TP4_SP8 (sweep rank-1, 2026-06-12)
    (8192, 6144, 768): (8, 4, 4, (1, 4)),  # 112 tiles — to_out        [UNSWEPT ESTIMATE]
    (8192, 4608, 768): (8, 4, 4, (1, 4)),  # 112 tiles — ff2 spatial   [UNSWEPT ESTIMATE]
    # M=16384 companion for N=1536 (K=6144, M_b=16): overhead~1069 KB, max 233 tiles; (16,2,8)→224✓
    (16384, 6144, 1536): (16, 2, 8, (1, 4)),  # 224 tiles — attn/proj   [UNSWEPT ESTIMATE]
    # x_c_merged for M=8192+128=8320 → M_tiles=260; 260%4=0 (260%8≠0), use M_block=4
    (8320, 6144, 4608): (4, 2, 8, (1, 4)),  #  64 tiles — ff1 merged    [UNSWEPT ESTIMATE]
    (8320, 6144, 2304): (4, 2, 8, (4, 1)),  #  64 tiles  [UNSWEPT ESTIMATE]
    (8320, 4608, 768): (4, 4, 4, (1, 4)),  #  48 tiles — ff2 merged    [UNSWEPT ESTIMATE]
    # x_c_merged for M=8192+64=8256 → M_tiles=258; 258%6=0 (258%16≠0 → M_block snaps to 6).
    # total_CBs ≈ (63/32)×matmul_CBs + 158 KB (empirical from compile-time arg analysis).
    # N=9216: use N_block=8 (not 16) — N_block=16 with (16,8,16) crashes at 3,256,832 B.
    # N_t=288; 288%8=0 ✓.
    (8256, 6144, 9216): (8, 6, 12, (2, 2)),  # 3940.7 μs — ff1 xc-merged TP4_SP8 (sweep rank-1, 2026-06-12)
    (8256, 6144, 4608): (6, 2, 8, (1, 4)),  # 104 tiles — ff1 merged     [UNSWEPT ESTIMATE]
    (8256, 6144, 2304): (6, 2, 8, (4, 1)),  # 104 tiles  [UNSWEPT ESTIMATE]
    (8256, 6144, 768): (6, 4, 4, (1, 4)),  #  88 tiles — to_out merged  [UNSWEPT ESTIMATE]
    # Spatial-only M=8192 with N=9216 — same N_block=8 constraint applies
    (8192, 6144, 9216): (8, 6, 12, (2, 2)),  # 3905.8 μs — ff1 spatial TP4_SP8 (sweep rank-1, 2026-06-12)
    # M=16384 with N=9216 (K=6144, M_b=16, formula: (63/32)×224t×2048+158KB≈1.06MB ✓)
    (16384, 6144, 9216): (16, 2, 8, (1, 4)),  # 224 tiles               [UNSWEPT ESTIMATE]
    # MMRS→AGMM: proj_out / forward_fused_addcmul run as AGMM.
    # K_global = K_shard × TP8 = 3072 × 8 = 24576. N = 6144 / TP8 = 768.
    # use_case="to_out" (fused addcmul). Swept 2026-06-17, matmulshapes_new.md Section 4.
    # (1152, 24576, 768): (3, 24, 3, (1, 3)),  # 704.1 μs — SNG_proj_out xc-merged 1024-tok (sweep rank-1, 2026-06-17)
    # (1024, 24576, 768): (6, 24, 3, (1, 3)),  # 647.2 μs — proj_out spatial-only 1024-tok (sweep rank-1, 2026-06-17)
    (1216, 4096, 32): (8, 8, 1, (4, 1)),
    (1216, 4096, 3072): (4, 8, 12, (1, 4)),
    (1216, 4096, 1024): (4, 8, 4, (1, 4)),
    (1216, 4096, 512): (4, 8, 2, (2, 2)),
    (1216, 2048, 1024): (4, 8, 4, (1, 4)),
    (1216, 4096, 4096): (4, 8, 16, (1, 4)),
    (4864, 4096, 32): (20, 8, 1, (4, 1)),
    (4864, 4096, 3072): (10, 4, 12, (1, 4)),
    (4864, 4096, 1024): (16, 8, 4, (4, 1)),
    (4864, 4096, 512): (8, 8, 2, (4, 1)),
    (4864, 2048, 1024): (5, 8, 4, (1, 4)),
    (4864, 4096, 4096): (5, 8, 16, (1, 4)),
    (256, 2048, 1024): (2, 8, 4, (1, 4)),
    (32, 2048, 32): (1, 8, 1, (1, 1)),
    (32, 2048, 1536): (1, 4, 16, (1, 4)),
    (32, 2048, 512): (1, 8, 2, (1, 2)),
    (32, 2048, 2048): (1, 4, 12, (1, 4)),
    # H3
    (3424, 5376, 5376): (4, 7, 14, (2, 2)),  # qkv
    (3424, 7168, 1344): (9, 8, 5, (3, 1)),  # to_out
    (3424, 5376, 7168): (6, 7, 12, (2, 2)),  # ff1
}


_BH_GALAXY_MIN_DEVICES = 32
_BH_GALAXY_MAX_CORE_GRID = (11, 10)


def get_matmul_core_grid(mesh_device):
    """Return the compute core grid, clamped to 11x10 on Blackhole Galaxy (power constraint)."""
    core_grid = mesh_device.compute_with_storage_grid_size()
    if ttnn.device.is_blackhole() and mesh_device.get_num_devices() >= _BH_GALAXY_MIN_DEVICES:
        core_grid = ttnn.CoreCoord(
            min(core_grid.x, _BH_GALAXY_MAX_CORE_GRID[0]),
            min(core_grid.y, _BH_GALAXY_MAX_CORE_GRID[1]),
        )
    return core_grid


def _compute_heuristic_blocking(M: int, K: int, N: int, grid_x: int, grid_y: int, tp_factor: int = -1):
    """Heuristic matmul blocking for shapes absent from the per-grid lookup tables.

    Rules derived empirically from BH GLX sweep data across 11×10 and 12×9 grids.

    ┌───────────┬────────────────────────────────────────────────────────────────┐
    │ M_block   │ ceil(M_tiles / grid_x), rounded up to the next sb_h multiple, │
    │           │ capped at 16.                                                  │
    ├───────────┼────────────────────────────────────────────────────────────────┤
    │ N_block   │ ceil(N_tiles / grid_y) rounded up to the next sb_w multiple   │
    │           │ that evenly divides N_tiles, capped at 16.                     │
    ├───────────┼────────────────────────────────────────────────────────────────┤
    │ K_block   │ min(8, K_tiles_per_device) snapped to a divisor of            │
    │           │ K_tiles_per_device. 8 tiles fits BH's 1.5 MB L1 well.         │
    │           │ K_tiles_per_device = K/(tp_factor*32) when tp_factor is set,  │
    │           │ otherwise K//32.                                               │
    ├───────────┼────────────────────────────────────────────────────────────────┤
    │ subblock  │ (1, 4) default.                                                │
    │           │ (4, 1) when N_tiles ≤ 4 (tiny/narrow N → maximise M reuse).  │
    │           │ (1, 3) when N_tiles % 4 ≠ 0, N_tiles % 3 == 0, N_tiles ≤ 24. │
    │           │ (1, 2) when N_tiles % 4 ≠ 0, N_tiles % 2 == 0.               │
    │           │ (1, 1) otherwise.                                              │
    └───────────┴────────────────────────────────────────────────────────────────┘

    Examples (validated against sweep data):
        (512,  5120, 2560), 11×10 → (2, 8,  8, (1, 4))
        (512,  6144, 2304), 11×10 → (2, 8,  8, (1, 4))
        (1024, 6144, 4608), 11×10 → (3, 8, 16, (1, 4))
        (128,   512,  192), 11×10 → (1, 8,  3, (1, 3))
        (512,  6144,  768), 12× 9 → (2, 8,  4, (1, 4))  [sweep-optimal: (2,16,3,(1,3))]
        (1024, 6144,  128), 11×10 → (4, 8,  1, (4, 1))  [sweep-optimal: (8,8,1,(4,1))]
    """
    M_t = M // 32
    K_t = K // 32
    N_t = N // 32

    m_raw = max(1, math.ceil(M_t / grid_x))  # target M-blocks per column of cores
    n_raw = N_t / grid_y  # target N-tiles per row of cores

    # ── subblock (sb_h, sb_w) ─────────────────────────────────────────────
    if N_t <= 4:
        sb_h, sb_w = 4, 1  # tiny/narrow N: maximise M-direction reuse
    elif N_t % 4 == 0:
        sb_h, sb_w = 1, 4  # wide output (dominant case)
    elif N_t % 3 == 0 and N_t <= 24:
        sb_h, sb_w = 1, 3
    elif N_t % 2 == 0:
        sb_h, sb_w = 1, 2
    else:
        sb_h, sb_w = 1, 1

    # ── M_block ───────────────────────────────────────────────────────────
    M_block = max(sb_h, math.ceil(m_raw / sb_h) * sb_h)
    M_block = min(M_block, 16)

    # ── N_block ───────────────────────────────────────────────────────────
    N_block = max(sb_w, math.ceil(n_raw / sb_w) * sb_w)
    N_block = min(N_block, 16)
    # Snap down to the nearest multiple of sb_w that divides N_t exactly
    while N_t % N_block != 0 and N_block > sb_w:
        N_block -= sb_w
    N_block = max(sb_w, N_block)

    # ── K_block ───────────────────────────────────────────────────────────
    K_t_per_device = K // (tp_factor * 32) if tp_factor > 0 else K_t
    K_block = min(8, K_t_per_device)
    while K_t_per_device % K_block != 0 and K_block > 1:
        K_block -= 1

    return M_block, K_block, N_block, (sb_h, sb_w)


grid_12_8_configs = {
    # Concat-shapes grid sweep — bh_4x8_sp0_tp1 (TP8_SP4), K_global=6144
    # Source: 1024_TP8_X_both_06_21_26.csv.  Swept 2026-06-22.
    # All ops: AllGatherMinimalMatmulAsyncOp (agmm).
    (1152, 6144, 4608): (3, 8, 10, (3, 1)),  # fused swiglu optimized — ff1 xc-merged plain
    (1024, 6144, 4608): (3, 6, 10, (3, 1)),  # fused swiglu optimized — ff1 spatial plain
    (128, 6144, 4608): (2, 8, 6, (2, 2)),  # fused swiglu optimized — 12×9 preferred
    (1152, 6144, 2304): (3, 8, 10, (3, 1)),  # 273.0 µs, −7.1% vs 12×9 — attn to_qkv xc-merged
    (1024, 6144, 2304): (3, 8, 10, (3, 1)),  # 252.6 µs, −6.4% vs 12×9 — attn to_qkv spatial
    (1024, 6144, 768): (8, 8, 3, (4, 1)),  # 194.5 µs, ≈flat vs 12×9 (194.3) — attn to_out spatial
    (128, 6144, 2304): (2, 12, 6, (2, 2)),  # 209.1 µs (+11.2% vs 12×9 winner) — 12×9 preferred
    (128, 6144, 768): (4, 12, 4, (2, 2)),  # 87.4 µs, −2.2% vs 12×9 — attn to_out prompt
}

_grid_config_lookup = {
    (8, 8): grid_88_configs,
    (8, 9): grid_89_configs,
    (13, 9): grid_13_9_configs,
    (12, 10): grid_12_10_configs,
    (11, 10): grid_11_10_configs,
    (12, 9): grid_12_9_configs,
    (12, 8): grid_12_8_configs,
}


def _lookup_nearest_blocking(M: int, K: int, N: int, grid_x: int, grid_y: int, tp_factor: int):
    """Return the blocking config from the grid's lookup table whose (M, K, N) is
    closest to the query shape.

    Distance is measured in log-ratio space so that relative differences matter equally
    at all scales (e.g. 64→128 is the same distance as 1024→2048):

        d = |log(M/M_ref)| + |log(K/K_ref)| + |log(N/N_ref)|

    Block sizes borrowed from the nearest neighbour are snapped down to the nearest
    divisor of the query's tile counts so the config is always valid for the query shape.
    When tp_factor is provided, K_block is snapped against K_tiles_per_device
    (K / (tp_factor * 32)) so the config is valid under tensor-parallel K-sharding.

    Returns the config tuple for the nearest shape, or None if the grid has no table.
    """
    grid_dict = _grid_config_lookup.get((grid_x, grid_y))
    if not grid_dict:
        return None

    lM, lK, lN = math.log(M), math.log(K), math.log(N)
    best_key, best_dist = None, float("inf")
    for ref_M, ref_K, ref_N in grid_dict:
        dist = abs(lM - math.log(ref_M)) + abs(lK - math.log(ref_K)) + abs(lN - math.log(ref_N))
        if dist < best_dist:
            best_dist = dist
            best_key = (ref_M, ref_K, ref_N)

    cfg = grid_dict[best_key]
    if len(cfg) == 4:
        Mb, Kb, Nb, sb = cfg
    else:
        Mb, Kb, Nb = cfg
        sb = None

    M_t, N_t = M // 32, N // 32
    K_t_per_device = K // (tp_factor * 32)

    def snap(val, tiles):
        """Snap val down to the nearest divisor of tiles."""
        val = min(val, tiles)
        while tiles % val != 0 and val > 1:
            val -= 1
        return val

    Mb = snap(Mb, M_t)
    Kb = snap(Kb, K_t_per_device)
    Nb = snap(Nb, N_t)

    return (Mb, Kb, Nb, sb) if sb is not None else (Mb, Kb, Nb)


def get_matmul_config(M, K, N, core_grid, default_block_size=None, use_heuristic=False, num_k_shards=-1):
    """Get the matmul config for the given shape and core grid.

    Args:
        M: The number of rows in the matrix.
        K: The number of columns in the matrix.
        N: The number of elements in the matrix.
        core_grid: The core grid to use.
        default_block_size: The default block size to use if no config is found.
        use_heuristic: Whether to use the heuristic blocking algorithm.
        num_k_shards: The number of shards. If set, we will use the nearest blocking algorithm.

        Algorithm order:
            - Default block size
            - Heuristic blocking
            - Nearest blocking based on num_k_shards provided num_k_shards

    Returns:
        The matmul config.
    """
    # Default to 8x8x8 with subblock 2x2 when unknown
    subblock_h = 2
    subblock_w = 2

    # Fallback core grid if not provided
    if core_grid is None:
        core_grid = ttnn.CoreCoord(8, 8)

    config_tuple = None
    grid_x = getattr(core_grid, "x", None)
    grid_y = getattr(core_grid, "y", None)
    grid_dict = _grid_config_lookup.get((grid_x, grid_y))
    if grid_dict is not None:
        config_tuple = grid_dict.get((M, K, N))

    # Unpack: 3-tuple (M_block_size, K_block_size, N_block_size) or
    # 4-tuple (M_block_size, K_block_size, N_block_size, (sub_h, sub_w))
    if config_tuple is not None and len(config_tuple) == 4:
        subblock_h, subblock_w = config_tuple[3]
        config_tuple = config_tuple[:3]

    if config_tuple is None:
        if default_block_size is not None:
            assert (
                not use_heuristic and num_k_shards == -1
            ), "Default block size cannot be used with heuristic or nearest blocking"
            M_block_size, K_block_size, N_block_size = default_block_size
        elif use_heuristic:
            M_block_size, K_block_size, N_block_size, (subblock_h, subblock_w) = _compute_heuristic_blocking(
                M, K, N, grid_x, grid_y
            )
        elif num_k_shards > 0:
            nearest = _lookup_nearest_blocking(M, K, N, grid_x, grid_y, num_k_shards) or (8, 8, 8)
            if len(nearest) == 4:
                subblock_h, subblock_w = nearest[3]
                nearest = nearest[:3]
            M_block_size, K_block_size, N_block_size = nearest
        else:
            M_block_size, K_block_size, N_block_size = 8, 8, 8

        M_tiles = math.ceil(M / 32)
        N_tiles = math.ceil(N / 32)

        if M_tiles < M_block_size:
            M_block_size = subblock_h
        if N_tiles < N_block_size:
            N_block_size = subblock_w

        signature = (M, K, N, grid_x, grid_y)
        if signature not in _warned_matmul_signatures:
            logger.warning(
                f"No known best blocking for (M, K, N) = ({M}, {K}, {N}) on {grid_x}x{grid_y} core grid; using default {M_block_size}x{K_block_size}x{N_block_size}"
            )
            _warned_matmul_signatures.add(signature)
    else:
        M_block_size, K_block_size, N_block_size = config_tuple

    return ttnn.MinimalMatmulConfig(
        M_block_size=M_block_size,
        K_block_size=K_block_size,
        N_block_size=N_block_size,
        subblock_h=subblock_h,
        subblock_w=subblock_w,
        compute_with_storage_grid_size=core_grid,
    )


def get_1d_matmul_config(
    M: int,
    K: int,
    N: int,
    core_grid: ttnn.CoreCoord,
) -> ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig:
    """Return a 1D mcast_in0 program config for small-M matmuls (M ≤ 64).

    Broadcasts the activation to all cores; each core handles a contiguous
    shard of N tiles from DRAM.  This eliminates the M-padding waste of the
    2D minimal_matmul distribution, where (grid_y - 1)/grid_y cores process
    zero-padded rows when M_tiles < grid_y.

    Parameters are computed as safe defaults suitable for correctness; run
    the 1D sweep in test_sweep_mm.py to find the optimal in0_block_w /
    per_core_N / out_subblock_w for each shape.
    """
    num_cores = core_grid.x * core_grid.y
    M_tiles = M // 32
    K_tiles = K // 32
    N_tiles = N // 32
    per_core_M = max(1, M_tiles)

    signature = (M, K, N, core_grid.x, core_grid.y)
    config_tuple = grid_1d_configs.get(signature)

    if config_tuple is not None:
        in0_block_w, per_core_N, out_subblock_w = config_tuple
    else:
        if signature not in _warned_1d_matmul_signatures:
            logger.warning(
                f"1D matmul: no swept config for (M, K, N) = ({M}, {K}, {N}) on "
                f"{core_grid.x}x{core_grid.y} grid; using default blocking — "
                f"run test_1d_matmul_sweep_bh4x8_ring to find optimal params"
            )
            _warned_1d_matmul_signatures.add(signature)

        per_core_N = max(1, math.ceil(N_tiles / num_cores))

        # in0_block_w: K-tiles per inner loop step. Must divide K_tiles.
        # 4 tiles is a conservative default that fits all target K shapes.
        in0_block_w = 4
        while K_tiles % in0_block_w != 0 and in0_block_w > 1:
            in0_block_w -= 1

        # out_subblock_h = 1 satisfies the 1D mcast constraint:
        #   out_subblock_w == per_core_N  OR  out_subblock_h == 1
        out_subblock_w = min(8, per_core_N)
        while per_core_N % out_subblock_w != 0 and out_subblock_w > 1:
            out_subblock_w -= 1

    out_subblock_h = 1

    return ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(
        compute_with_storage_grid_size=core_grid,
        in0_block_w=in0_block_w,
        out_subblock_h=out_subblock_h,
        out_subblock_w=out_subblock_w,
        per_core_M=per_core_M,
        per_core_N=per_core_N,
        fuse_batch=True,
        fused_activation=None,
        mcast_in0=True,
    )


class FusedMMRSConfig(NamedTuple):
    compute_with_storage_grid_size: ttnn.CoreCoord
    M_block_size: int
    K_block_size: int
    N_block_size: int
    subblock_h: int
    subblock_w: int
    num_buffers_per_channel: int | None
    chunk_width_in_mm_blocks: int
    # Optional explicit reduce-scatter worker count
    num_workers_per_link: int | None = None

    def get_params(self, core_grid, num_links):
        config_dict = self._asdict()
        num_buffers_per_channel = config_dict.pop("num_buffers_per_channel")
        chunk_width_in_mm_blocks = config_dict.pop("chunk_width_in_mm_blocks")
        num_workers_override = config_dict.pop("num_workers_per_link")

        if num_workers_override is not None:
            num_workers_per_link = num_workers_override
        else:
            rs_zone_capacity = (core_grid.y - self.compute_with_storage_grid_size.y) * core_grid.x
            num_workers_per_link = rs_zone_capacity // (2 * num_links) - 1

        # Order is important. Guaranteed for python 3.7+
        return {
            "reduce_scatter_core_grid_offset": ttnn.CoreCoord(0, self.compute_with_storage_grid_size.y),
            "num_links": num_links,
            "config": ttnn.MinimalMatmulConfig(**config_dict),
            "num_buffers_per_channel": num_buffers_per_channel,
            "chunk_width_in_mm_blocks": chunk_width_in_mm_blocks,
            "num_workers_per_link": num_workers_per_link,
        }


if is_blackhole():
    default_fused_mmrs_config = FusedMMRSConfig(ttnn.CoreCoord(12, 8), 8, 4, 8, 2, 1, None, 1)
else:
    default_fused_mmrs_config = FusedMMRSConfig(ttnn.CoreCoord(8, 7), 2, 8, 8, 1, 1, None, 1)

# core_grid: {MKN: mm_core_grid, M, K, N, sub_h, sub_w, num_buffers_per_channel, chunk_width_in_mm_blocks, num_workers_per_link}
fused_mmrs_configs = {
    ttnn.CoreCoord(8, 9): {
        (9472, 5120, 1280): FusedMMRSConfig(ttnn.CoreCoord(8, 7), 8, 8, 8, 2, 2, None, 1),
    },
    ttnn.CoreCoord(12, 10): {
        (9472, 3456, 5120): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 8, 4, 8, 2, 1, None, 1),
        (9472 // 4, 3456, 5120): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 4, 8, 2, 2, None, 1),
        # LTX video FFN ff2 (RowParallel): per-device [4864,4096]@[4096,4096]
        (4864, 4096, 4096): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 7, 5, 6, 1, 3, None, 1, 3),
        (1152, 3072, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 12, 4, 8, 1, 1, None, 1, 5),
        (512, 3072, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 8, 4, 8, 2, 2, None, 1),
        (512, 2304, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 4, 8, 2, 2, None, 1),
        (1024, 3072, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 8, 4, 8, 2, 1, None, 1),
        (1024, 2304, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 4, 8, 1, 2, None, 1),
        # 2048-resolution shapes — BH 4×8 ring sweep (2026-05-30, 2048.md)
        (4096, 2304, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 3, 8, 1, 2, None, 1),  # 1546.0 μs
        (128, 2304, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 6, 8, 1, 2, None, 1),  # 595.7 μs
        (4096, 3072, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 3, 8, 2, 2, None, 1),  # 1408.3 μs
        (128, 3072, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 6, 8, 1, 2, None, 1),  # 626.2 μs
        # x_c_merged: concat([spatial, prompt_sp_sharded]) → M=4096+128=4224; same K/N/grid as (4096,…). This is an unconfirmed estimate.
        (4224, 3072, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 3, 8, 2, 2, None, 1),
        (4224, 2304, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 3, 8, 1, 2, None, 1),
    },
}


def get_fused_mmrs_config(M, K, N, device_core_grid, num_links):
    config = fused_mmrs_configs.get(device_core_grid, {})
    if len(config) == 0:
        logger.warning(f"No known fused MM/RS config for {device_core_grid} core grid, using default")
    elif (M, K, N) not in config:
        logger.warning(
            f"No known fused MM/RS config for (M, K, N) = ({M}, {K}, {N}) on {device_core_grid} core grid, using default"
        )
    elif (M, K, N) not in config:
        # Worth a warning even though the grid is known: the default puts the matmul on an 8x7 grid at
        # subblock 1x1, so on a larger device it is drastically slower than the unfused
        # matmul + reduce-scatter it is meant to replace, and the fusion reads as a regression.
        logger.warning(
            f"No known best MM/RS blocking for (M, K, N) = ({M}, {K}, {N}) on {device_core_grid} core grid; "
            "using default, which is likely slower than not fusing at all"
        )
    config = config.get((M, K, N), default_fused_mmrs_config)
    return config.get_params(device_core_grid, num_links)


def register_matmul_configs(configs: dict) -> None:
    """Register additional matmul block-size configs from external models.

    Args:
        configs: Mapping from grid key string to dict of (M,K,N) -> config tuples.
            Grid keys: ``"11x10"``, ``"12x10"``, ``"12x9"``, ``"13x9"``, ``"8x8"``, ``"8x9"``.
            Config tuple format: ``(M_block, K_block, N_block)`` or
            ``(M_block, K_block, N_block, (sub_h, sub_w))``.
            When subblock is omitted, the default ``(2, 2)`` is used.

    Example::

        register_matmul_configs({
            "11x10": {
                (14400, 384, 384): (9, 12, 3, (3, 1)),
                (14400, 5120, 3456): (15, 20, 1, (3, 1)),
            },
        })
    """
    grid_map = {
        "8x8": grid_88_configs,
        "8x9": grid_89_configs,
        "11x10": grid_11_10_configs,
        "12x10": grid_12_10_configs,
        "12x9": grid_12_9_configs,
        "13x9": grid_13_9_configs,
    }
    for grid_key, entries in configs.items():
        target = grid_map.get(grid_key)
        if target is None:
            msg = f"Unknown grid key {grid_key!r}, expected one of {list(grid_map)}"
            raise ValueError(msg)
        target.update(entries)


def register_fused_mmrs_configs(configs: dict) -> None:
    """Register additional fused matmul+reduce-scatter configs.

    Args:
        configs: Mapping from ``ttnn.CoreCoord`` to dict of
            ``(M,K,N)`` -> :class:`FusedMMRSConfig`.

    Example::

        register_fused_mmrs_configs({
            ttnn.CoreCoord(12, 10): {
                (14400, 3456, 5120): FusedMMRSConfig(
                    ttnn.CoreCoord(12, 8), 8, 4, 8, 2, 1, None, 1
                ),
            },
        })
    """
    for core_grid, entries in configs.items():
        fused_mmrs_configs.setdefault(core_grid, {}).update(entries)


# ===================================================================== Fabric-bound all-gather-matmul
class FabricAGMMConfig(NamedTuple):
    mm_core_grid: ttnn.CoreCoord
    ag_core_grid_offset: tuple
    M_block_size: int
    K_block_size: int
    N_block_size: int
    subblock_h: int
    subblock_w: int
    num_workers_per_link: int
    num_buffers_per_channel: int


# Keyed by device core-grid (``ttnn.CoreCoord``) then ``(M, K, N, chunks)``. M is the per-SP-device
# sequence length, K the full (gathered) contraction dim, N the per-TP-device WEIGHT width, chunks
# the output-split count. For a fused-SwiGLU layer the weight is the packed [gate|up] matrix, so N
# is twice the layer's output width (that is also the N the device op blocks on).
#
# ``M=-1`` is a wildcard: ``get_fabric_agmm_config`` tries an exact ``(M, K, N, chunks)`` match
# first, then ``(-1, K, N, chunks)``. LTX entries use the wildcard because in-model M varies with
# resolution and SP padding. Flux entries are exact-M; an exact-match miss is harmless (the caller
# falls back to ``all_gather_minimal_matmul_async``).
#
# Every entry must be measured on device before it lands here: a guessed entry silently ships an
# untuned blocking, and a bad M/N block pairing can deadlock the split output write outright.
fabric_agmm_configs: dict[ttnn.CoreCoord, dict[tuple, FabricAGMMConfig]] = {
    ttnn.CoreCoord(12, 10): {
        # (M, K, N, chunks) -> FabricAGMMConfig(mm_core_grid, ag_core_grid_offset,
        #                                       M_block, K_block, N_block, sub_h, sub_w, num_w/link, num_buf/chan)
        # LTX: M-independent (wildcard). Same blocking as the previous (K, N, chunks) keys.
        (-1, 4096, 1024, 1): FabricAGMMConfig(ttnn.CoreCoord(12, 8), (0, 8), 16, 8, 4, 2, 2, 3, 8),
        # video to_gate_logits: per-device N = num_heads/tp = 8, tile-padded to 32 (1 tile)
        (-1, 4096, 32, 1): FabricAGMMConfig(ttnn.CoreCoord(12, 8), (0, 8), 16, 8, 1, 2, 1, 3, 8),
        # audio_to_video_attn.to_q (a2v cross): K = video_dim = 4096, N = audio_dim/tp = 512
        (-1, 4096, 512, 1): FabricAGMMConfig(ttnn.CoreCoord(12, 8), (0, 8), 16, 8, 2, 2, 2, 3, 8),
        # -----------------------------------------------------------------------------------------
        # flux2 feed-forward shapes, K_global = 6144, packed SwiGLU weight width N = 4608. Every M
        # here is smaller than N, so the factory does NOT transpose the grid: N splits across
        # grid.x=12 (12 weight tiles/core) and M across grid.y=8 (1 / 4 / 5 tiles/core).
        #
        # N_block=12 is the full per-core width everywhere, which keeps N_blocks_per_core == 1.
        # Two N blocks per core combined with a single M block deadlocks the split output write
        # (device-confirmed hang on the LTX a2v shape), so full-width N is the safe choice as well
        # as the fastest one in the sweep. N_block is even, as fused SwiGLU requires.
        #
        # Swept on device via sweep_mm_block_sizes.py (op_kind "sagmm"); durations below are the
        # best device_kernel_duration for that M. Per-M keying buys 9.0% at M=128 and 6.2% at
        # M=1024 over the single M-independent entry this table used to carry.
        #
        # ff / ff_context context stream (8 layers each).
        (128, 6144, 4608, 1): FabricAGMMConfig(ttnn.CoreCoord(12, 8), (0, 8), 2, 8, 12, 2, 2, 3, 8),  # 288.0 us
        # ff / ff_context image stream (8 layers each).
        (1024, 6144, 4608, 1): FabricAGMMConfig(ttnn.CoreCoord(12, 8), (0, 8), 4, 8, 12, 1, 4, 3, 8),  # 336.4 us
        # single-stream proj_mlp (48 layers): M = 1024 image + 128 context.
        (1152, 6144, 4608, 1): FabricAGMMConfig(ttnn.CoreCoord(12, 8), (0, 8), 5, 4, 12, 1, 4, 3, 8),  # 382.9 us
    },
}


def get_fabric_agmm_config(M, K, N, chunks, device_core_grid) -> FabricAGMMConfig | None:
    """Return the tuned fabric-bound strided-AGMM config for this shape, or ``None``.

    ``None`` means the shape is not (known to be) fabric-bound; the caller keeps the current
    ``all_gather_minimal_matmul_async`` path. Keyed on ``(M, K, N, chunks)``, where ``N`` is the
    per-device weight width (the packed [gate|up] width for fused SwiGLU). ``M=-1`` in the table
    is a wildcard: an exact match is tried first, then ``(-1, K, N, chunks)``.

    A/B switch: set ``DISABLE_FABRIC_AGMM=1`` to force a miss for every shape, routing the whole
    model back onto the old ``all_gather_minimal_matmul_async`` op. Used to get an apples-to-apples
    old-agmm-vs-strided-sagmm baseline under the same trace/fabric config.
    """
    if os.environ.get("DISABLE_FABRIC_AGMM") in ("1", "true", "True"):
        return None
    table = fabric_agmm_configs.get(device_core_grid, {})
    cfg = table.get((M, K, N, chunks))
    if cfg is None and M != -1:
        cfg = table.get((-1, K, N, chunks))
    return cfg


def register_fabric_agmm_configs(configs: dict) -> None:
    """Register additional fabric-bound strided-AGMM configs from external models.

    Args:
        configs: Mapping from ``ttnn.CoreCoord`` (device core-grid) to dict of
            ``(M, K, N, chunks)`` -> :class:`FabricAGMMConfig`. Use ``M=-1`` for an
            M-independent wildcard entry.

    Example::

        register_fabric_agmm_configs({
            ttnn.CoreCoord(12, 10): {
                (1024, 6144, 4608, 1): FabricAGMMConfig(
                    ttnn.CoreCoord(12, 8), (0, 8), 4, 8, 12, 1, 4, 3, 8
                ),
            },
        })
    """
    for core_grid, entries in configs.items():
        fabric_agmm_configs.setdefault(core_grid, {}).update(entries)