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 agmm_worker_grid(full_grid, transpose):
    """all_gather_minimal_matmul_async matmul worker grid, sized to leave the mux axis free.

    The op places its input muxes on the device's last ROW when the core grid is transposed and its
    last COLUMN when it is not (see the program factory's `in0_mux_in_column` logic). The matmul
    workers must therefore avoid that row/column: reserve a row when transposed -> (x, y-1); reserve a
    column when not -> (x-1, y). `transpose` should be the op's own decision, i.e.
    `force_transpose or (M > N)`.
    """
    return ttnn.CoreCoord(full_grid.x, full_grid.y - 1) if transpose else ttnn.CoreCoord(full_grid.x - 1, full_grid.y)


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,
    )


_logged_agmm_v3_signatures = set()


def get_agmm_config(
    M,
    K,
    N,
    full_grid,
    cluster_size,
    num_links,
    core_grid=None,
    default_block_size=None,
    use_heuristic=False,
    fuse_swiglu=False,
    use_addcmul=False,
    force_transpose=True,
):
    """Resolve (core_grid, MinimalMatmulConfig, num_workers_per_link) for
    `all_gather_minimal_matmul_async`.

    Precedence -- anything measured or explicit wins, bit-for-bit identical to the legacy path:
      1. An explicit `core_grid` or `default_block_size` from the caller.
      2. A swept `(M, K, N)` table entry for the legacy grid (`_grid_config_lookup`, i.e.
         `grid_12_9_configs` and entries added via `register_matmul_configs`).
      3. The v3 rule engine (`utils/agmm_rules.py`) -- picks the worker grid (orientation-aware)
         and blocking; blind-validated within 5% of swept optimum on ~96% of shapes. Blackhole
         only; the rules were fitted against Blackhole L1 and fabric parameters. Skipped when
         `force_transpose` overrides the op's own `M > N` orientation, since the rules derive
         layout and blocking from that same test and have no fit for the forced orientation.
      4. The legacy warned generic fallback.

    `force_transpose` mirrors the op parameter (whose default is likewise `True`) and must be the
    value the caller actually passes to the op, since the worker grid has to reserve the mux axis
    the op will use.

    On a v3 hit, an info log prints the config as a paste-able table line: sweep the shape with
    `sweep_mm_block_sizes.py` and paste the winner into `grid_12_9_configs` (or a model table)
    to override the rules permanently.
    """
    # The op transposes when the caller forces it or the output is wide, and places its in0 muxes
    # on the axis the worker grid leaves free -- reserve a row when transposed, a column when not.
    transpose_core_grid = force_transpose or M > N
    legacy_grid = core_grid or agmm_worker_grid(full_grid, transpose_core_grid)
    # in0 senders group along the M-parallel axis (grid.x transposed, grid.y not); split it into
    # exactly num_links groups, matching the op's `ceil(in0_axis / workers) == num_links` assert.
    # For the default transposed grid this is the same 6 the old `full_grid.x // num_links` gave.
    legacy_workers = math.ceil((legacy_grid.x if transpose_core_grid else legacy_grid.y) / num_links)
    table_hit = _grid_config_lookup.get((legacy_grid.x, legacy_grid.y), {}).get((M, K, N)) is not None
    if core_grid is not None or default_block_size is not None or use_heuristic or table_hit:
        config = get_matmul_config(M, K, N, legacy_grid, default_block_size, use_heuristic)
        return legacy_grid, config, legacy_workers

    v3 = None
    if is_blackhole() and transpose_core_grid == (M > N):
        from .agmm_rules import pick_v3

        v3 = pick_v3(
            M,
            K,
            N,
            cluster_size=cluster_size,
            full_grid=(full_grid.x, full_grid.y),
            fuse_swiglu=fuse_swiglu,
            use_addcmul=use_addcmul,
        )
    if v3 is None:
        config = get_matmul_config(M, K, N, legacy_grid)  # legacy warned generic fallback
        return legacy_grid, config, legacy_workers

    grid_x, grid_y = v3["core_grid"]
    m_blk, k_blk, n_blk = v3["blocks"]
    sub_h, sub_w = v3["subblock"]
    signature = (M, K, N, grid_x, grid_y)
    if signature not in _logged_agmm_v3_signatures:
        logger.info(
            f"AGMM v3 rule config for (M, K, N) = ({M}, {K}, {N}) on {grid_x}x{grid_y}"
            f"{' transposed' if v3['transposed'] else ''}: "
            f"({M}, {K}, {N}): ({m_blk}, {k_blk}, {n_blk}, ({sub_h}, {sub_w}))  "
            f"# paste into grid config table after sweeping to override"
        )
        _logged_agmm_v3_signatures.add(signature)
    config = ttnn.MinimalMatmulConfig(
        M_block_size=m_blk,
        K_block_size=k_blk,
        N_block_size=n_blk,
        subblock_h=sub_h,
        subblock_w=sub_w,
        compute_with_storage_grid_size=ttnn.CoreCoord(grid_x, grid_y),
    )
    in0_axis = grid_x if v3["transposed"] else grid_y
    return ttnn.CoreCoord(grid_x, grid_y), config, math.ceil(in0_axis / num_links)


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
    # Rolling L1 window over the MM output, in M blocks. The op requires it whenever the MM output
    # lands in L1: unwindowed, the resident shard is Mt_per_core * Nt_per_core tiles per matmul
    # core, which for a large M crowds out later programs' circular buffers (issue #52863). 2 is
    # the shallowest depth that still lets the matmul run a block ahead of the RS readers, and
    # measured perf is flat in this knob, so the default leaves the most L1 to circular buffers.
    # get_params clamps it to the blocks a core actually has.
    mm_window_blocks: int | None = 2

    def get_params(self, core_grid, num_links, M=None):
        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")
        mm_window_blocks = config_dict.pop("mm_window_blocks")
        if mm_window_blocks is not None and M is not None:
            # Clamp to the M blocks a core actually walks (fused MMRS never transposes: M on grid.y).
            mt_per_core = math.ceil(math.ceil(M / 32) / self.compute_with_storage_grid_size.y)
            blocks_per_core = math.ceil(mt_per_core / self.M_block_size)
            if blocks_per_core <= 1:
                # A single block per core makes the window degenerate twice over: there is no slot
                # rotation, so the pipelining the window exists for cannot happen; and the windowed
                # tensor's height is quantized to WHOLE M blocks, so when M_block > Mt_per_core the
                # "window" is TALLER than the real output (e.g. flux2 1024px: M_block=12 over 5 real
                # rows -> a 384 KB/core resident shard, 2.4x full residency, which clashed with the
                # blocking's own CBs). Fall back to the DRAM handoff instead.
                mm_window_blocks = None
            else:
                mm_window_blocks = min(mm_window_blocks, blocks_per_core)

        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,
            "mm_window_blocks": mm_window_blocks,
        }


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): {
        # Wan2.2 720p ff2, single galaxy, swept 2026-08-24 (windowed): Mt_per_core=37, M_block=8
        # leaves 5 blocks.
        (9472, 3456, 5120): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 8, 3, 8, 2, 2, None, 1, 5),  # 1744.0 us
        # Entries marked "swept 2026-08-24 (windowed)" were swept under the windowed L1 handoff
        # with the runner that windows combos leaving >= 2 M blocks per core (DRAM otherwise), on
        # a binary verified to place the RS intermediate in DRAM (an earlier stale build parked it
        # in L1, flattering times ~3-4%; those numbers were restated from clean re-sweeps).
        # M_block must stay <= ceil(Mt_per_core / 2) for the window to rotate.
        #
        # Wan2.2 720p ff2 on the quad-galaxy config (M = 9472 / 4), swept 2026-08-24 (windowed):
        # 557.4 us windowed vs 722.2 us best-DRAM. Mt_per_core=10, M_block=6 leaves 2 blocks.
        (9472 // 4, 3456, 5120): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 6, 3, 8, 2, 2, None, 1, 5),  # 557.4 us
        # LTX ff2 @stage_1 (M = 9728/sp8), swept 2026-08-24 (windowed): 362.9 us vs 376.1 us
        # best-DRAM. Previously absent, so it fell to the default config whose M_block=8 exceeds
        # the 5 rows per core -> DRAM fallback; this entry puts stage_1 on the windowed handoff.
        (1216, 4096, 4096): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 8, 6, 2, 2, None, 1, 5),  # 362.9 us
        # LTX ff2 @stage_2 (M = 38912/sp8), swept 2026-08-24 (windowed): 973.0 us vs ~1088 us for
        # the previous DRAM-era M7/K5/N6 blocking.
        (4864, 4096, 4096): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 8, 4, 6, 2, 2, None, 1, 5),  # 973.0 us
        # Aang ff2 (same K/N family as Wan). Windowed beats the best DRAM blocking on both:
        # a2v 601.9 us vs 746.99 us DRAM-swept @ M6/K4/N8 (-19%); SR 2084.2 us vs 2191.75 us
        # DRAM-swept @ M6/K4/N7 (-5%).
        (2656, 3456, 5120): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 6, 3, 8, 2, 2, None, 1, 5),  # a2v, 601.9 us
        (11520, 3456, 5120): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 6, 4, 8, 2, 2, None, 1, 5),  # SR, 2084.2 us
        # Aang @1080p variants, swept 2026-08-24 (windowed). a2v-1080p has only 7 M tile-rows per
        # core, so M_block drops to 4 to keep the window rotating.
        (1664, 3456, 5120): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 4, 6, 8, 2, 2, None, 1, 5),  # a2v 1080p, 428.3 us
        (7200, 3456, 5120): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 6, 4, 8, 2, 2, None, 1, 5),  # SR 1080p, 1378.2 us
        # Flux2 @1024px, swept 2026-08-24 (windowed). Windowed won both shapes outright:
        # 407.5 us vs 443.8 us best-DRAM for (1152,...), 339.9 us vs 345.4 us for (1024,...).
        # The pre-window (1152,...) M_block=12 entry was degenerate under the window (12-tile
        # block over 5 rows/core -> 384 KB resident shard, CB clash).
        (1152, 3072, 6144): FusedMMRSConfig(ttnn.CoreCoord(12, 8), 3, 6, 8, 1, 4, None, 1, 5),  # 407.5 us
        (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), 2, 4, 8, 2, 2, None, 1, 5),  # 339.9 us, see above
        # 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),
    },
}


_logged_mmrs_rule_signatures = set()


def get_fused_mmrs_config(M, K, N, device_core_grid, num_links):
    """Resolve the fused MM+RS parameter dict for a shape.

    Precedence -- anything measured or explicit wins:
      1. A swept ``(M, K, N)`` table entry (``fused_mmrs_configs``, including entries added via
         ``register_fused_mmrs_configs``).
      2. The v2.3 rule engine (`utils/mmrs_rules.py`) -- always the 12x8 matmul grid with the
         reduce-scatter above it; blind-validated within 5% of the swept optimum on ~90% of
         shapes across five ff2 families. Blackhole 12x10 only; the rules were fitted against
         Blackhole L1, and unaligned M stays off them (the RS output height must be tile-aligned).
      3. The warned generic default, which puts the matmul on an 8x7 grid at subblock 1x1 -- on a
         larger device drastically slower than the unfused matmul + reduce-scatter it replaces,
         so the fusion reads as a regression: sweep the shape instead.

    On a rule hit, an info log prints the config as a paste-able table entry: sweep the shape
    with `sweep_mm_block_sizes.py` (use case ``mmrs``) and paste the winner into
    ``fused_mmrs_configs`` (or register it from a model table) to override the rules permanently.
    """
    table = fused_mmrs_configs.get(device_core_grid, {})
    config = table.get((M, K, N))

    if config is None and is_blackhole() and M % 32 == 0:
        from .mmrs_rules import pick_v23

        v23 = pick_v23(M, K, N, full_grid=(device_core_grid.x, device_core_grid.y))
        if v23 is not None:
            m_blk, k_blk, n_blk = v23["blocks"]
            sub_h, sub_w = v23["subblock"]
            signature = (M, K, N, device_core_grid.x, device_core_grid.y)
            if signature not in _logged_mmrs_rule_signatures:
                logger.info(
                    f"MMRS v2.3 rule config for (M, K, N) = ({M}, {K}, {N}): "
                    f"FusedMMRSConfig(ttnn.CoreCoord{v23['mm_grid']}, "
                    f"{m_blk}, {k_blk}, {n_blk}, {sub_h}, {sub_w}, None, 1)  "
                    f"# paste into fused_mmrs_configs after sweeping to override"
                )
                _logged_mmrs_rule_signatures.add(signature)
            config = FusedMMRSConfig(ttnn.CoreCoord(*v23["mm_grid"]), m_blk, k_blk, n_blk, sub_h, sub_w, None, 1)

    if config is None:
        logger.warning(
            f"No fused MM/RS config for (M, K, N) = ({M}, {K}, {N}) on {device_core_grid} core grid "
            "and the rule engine does not cover it; using default, which is likely slower than not "
            "fusing at all"
        )
        config = default_fused_mmrs_config
    return config.get_params(device_core_grid, num_links, M=M)


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)