"""
2CTA Block-Scaled Matrix Multiplication
=======================================

High-performance 2CTA warp-specialized block-scaled MMA.
Two CTAs cooperate per output tile, sharing operands to
increase arithmetic intensity and reduce the per-CTA SMEM
footprint.
"""

import argparse
import itertools
import pytest
import torch

import triton
import triton.experimental.gluon as gluon
import triton.experimental.gluon.language as gl
from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor

from triton.experimental.gluon.nvidia.blackwell import TensorDescriptor
from triton.experimental.gluon.language.nvidia.blackwell import (
    TensorMemoryLayout,
    TensorMemoryScalesLayout,
    allocate_tensor_memory,
    tensor_memory_descriptor,
    clc,
    tcgen05_copy,
    tcgen05_commit,
    tcgen05_mma_barrier_count,
    tcgen05_mma_scaled,
    mbarrier,
    tma,
)

# ---------------------------------------------------------------------------
# Tile scheduler
# ---------------------------------------------------------------------------


@gluon.jit
def _planar_snake(lin_idx, m_tiles, n_tiles, minor_dim: gl.constexpr, tile_width: gl.constexpr):
    major_size = n_tiles if minor_dim == 0 else m_tiles
    minor_size = m_tiles if minor_dim == 0 else n_tiles

    full_minor_tiles = minor_size // tile_width
    full_minor_size = full_minor_tiles * tile_width
    full_elements = full_minor_tiles * tile_width * major_size

    minor_tile_idx = lin_idx // (tile_width * major_size)

    full_minor_within = lin_idx % tile_width
    full_major_within = (lin_idx // tile_width) % major_size
    full_minor = minor_tile_idx * tile_width + full_minor_within
    full_major = gl.where((minor_tile_idx % 2) == 0, full_major_within, major_size - 1 - full_major_within)

    partial_width = minor_size - full_minor_size
    partial_width = gl.where(partial_width > 0, partial_width, 1)
    partial_lin = lin_idx - full_elements
    partial_minor_within = partial_lin % partial_width
    partial_major_within = (partial_lin // partial_width) % major_size
    partial_minor = minor_tile_idx * tile_width + partial_minor_within
    partial_major = gl.where((minor_tile_idx % 2) == 0, partial_major_within, major_size - 1 - partial_major_within)

    in_full_tile = lin_idx < full_elements
    minor = gl.where(in_full_tile, full_minor, partial_minor)
    major = gl.where(in_full_tile, full_major, partial_major)

    if minor_dim == 0:
        return minor, major
    return major, minor


def is_blackwell():
    if not torch.cuda.is_available():
        return False
    target = triton.runtime.driver.active.get_current_target()
    return target.backend == "cuda" and torch.cuda.get_device_capability()[0] == 10


@gluon.constexpr_function
def get_split_dim(cga_layout, dim):
    return 1 << sum(b[dim] != 0 for b in cga_layout)


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------


def random_quantized_tensor(MN, K, format):
    assert format in ["mxfp4", "mxfp8", "nvfp4"]
    VEC_SIZE = 16 if format == "nvfp4" else 32

    # Generate a random quantized tensor and its scale factors, assuming we are
    # scaling along the K dimension.
    base = MXFP4Tensor(size=(MN, K), device="cuda").random()
    scale = MXScaleTensor(size=(MN, K // VEC_SIZE), device="cuda").random(low=1 / 128, high=2.0)

    # Compute the dequantized tensor to use for testing.
    ref = base.to(torch.float32)
    scale_ref = scale.to(torch.float32)
    value = ref * scale_ref.repeat_interleave(VEC_SIZE, dim=1)

    if format == "mxfp8":
        # For mxfp8, convert the tensor to a regular float8 torch tensor.
        return ref.to(torch.float8_e4m3fn), scale.data, value
    elif format == "mxfp4":
        # For mxfp4, pack the elements along the K dimension.
        return base.to_packed_tensor(dim=1), scale.data, value
    else:
        # For nvfp4, pack the elements along the K dimension, and convert the
        # scale factors to float8_e4m3fn.
        return base.to_packed_tensor(dim=1), scale_ref.to(torch.float8_e4m3fn), value


def align_to(a, b):
    # Return next multiple of `b` greater than or equal to `a`.
    return triton.cdiv(a, b) * b


def swizzle_scales_packed_block(scales: torch.Tensor):
    # When the scale tensor is not an even multiple of [128, 4], we need to pad
    # the scale tensor so it can use the packed block format.
    PAD_MN = align_to(scales.shape[0], 128) - scales.shape[0]
    PAD_K = align_to(scales.shape[1], 4) - scales.shape[1]
    scales = torch.nn.functional.pad(scales, (0, PAD_K, 0, PAD_MN))

    MN, SCALE_K = scales.shape[0], scales.shape[1]
    REP_MN = MN // 128
    REP_K = SCALE_K // 4
    scales = scales.reshape(REP_MN, 4, 32, REP_K, 4)
    scales = scales.permute(0, 3, 2, 1, 4)
    return scales.contiguous()


# ---------------------------------------------------------------------------
# Autotuning configs and hook
# ---------------------------------------------------------------------------

SCHEDULER_CLC = gl.constexpr("cluster_launch_control")
SCHEDULER_SPS = gl.constexpr("static_persistent")


def mma_scaled_get_configs(pre_hook=None, cga_layouts=None):
    if cga_layouts is None:
        cga_layouts = [(), ((1, 0), )]
    return [
        triton.Config(
            {
                "BLOCK_M": BM,
                "BLOCK_N": BN,
                "BLOCK_K": BK,
                "EPILOGUE_BLOCK_N": epilogue_n,
                "num_buffers": stages,
                "num_acc_buffers": acc_buffers,
                "GRID_MINOR_DIM": minor_dim,
                "GRID_TILE_WIDTH": grid_tile_width,
                "CGA_LAYOUT": cga_layout,
                "scheduler": scheduler,
            },
            num_warps=4,
            num_ctas=2**len(cga_layout),
            pre_hook=pre_hook,
        )
        for BM in (128, 256)
        for BN in (128, 256)
        for BK in (128, 256)
        for epilogue_n in (64, BN)
        for minor_dim in (0, 1)
        for grid_tile_width in (4, 8, 16)
        for stages in (3, 4, 5)
        for acc_buffers in (1, 2)
        for cga_layout in cga_layouts
        for scheduler in (SCHEDULER_CLC, SCHEDULER_SPS)
        # tcgen05_mma_scaled requires BLOCK_M_PER_CTA == 128
        if BM // (2**len(cga_layout)) == 128 and epilogue_n <= BN
    ]


def mma_scaled_tma_set_block_size_hook(nargs):
    block_m = nargs["BLOCK_M"]
    block_n = nargs["BLOCK_N"]
    block_k = nargs["BLOCK_K"]
    epilogue_n = nargs["EPILOGUE_BLOCK_N"]
    cga_layout = nargs["CGA_LAYOUT"]

    a_base = nargs["a_desc"].base
    b_base = nargs["b_desc"].base
    a_is_fp4 = a_base.dtype == torch.uint8
    b_is_fp4 = b_base.dtype == torch.uint8
    mixed_prec = a_is_fp4 != b_is_fp4
    a_elem_per_byte = 2 if a_is_fp4 else 1
    b_elem_per_byte = 2 if b_is_fp4 else 1

    a_block = [block_m, block_k // a_elem_per_byte]
    b_block = [block_n, block_k // b_elem_per_byte]
    c_block = [block_m, epilogue_n]

    nargs["a_desc"].block_shape = a_block
    nargs["b_desc"].block_shape = b_block
    nargs["c_desc"].block_shape = c_block

    cga = tuple(tuple(x) for x in cga_layout) if cga_layout else None
    nargs["a_desc"].layout = gl.NVMMASharedLayout.get_default_for(a_block, gl.uint8 if a_is_fp4 else gl.float8e4nv,
                                                                  fp4_padded=(mixed_prec and a_is_fp4), cga_layout=cga)
    nargs["b_desc"].layout = gl.NVMMASharedLayout.get_default_for(b_block, gl.uint8 if b_is_fp4 else gl.float8e4nv,
                                                                  fp4_padded=(mixed_prec and b_is_fp4), cga_layout=cga)
    c_dtype = getattr(gl, str(nargs["c_desc"].base.dtype).split('.')[1])
    nargs["c_desc"].layout = gl.NVMMASharedLayout.get_default_for(c_block, c_dtype, cga_layout=cga)

    a_scale_base = nargs["a_scale_desc"].base
    is_nvfp4 = a_scale_base.dtype == torch.float8_e4m3fn
    vec_size = 16 if is_nvfp4 else 32
    rep_m = block_m // 128
    rep_n = block_n // 128
    rep_k = block_k // (vec_size * 4)
    nargs["a_scale_desc"].block_shape = [1, rep_m, rep_k, 2, 256]
    nargs["b_scale_desc"].block_shape = [1, rep_n, rep_k, 2, 256]

    if cga_layout:
        cga_a_scale = [[0, 1, 0, 0, 0]]
        cga_b_scale = [[0, 0, 0, 0, 0]]
        nargs["a_scale_desc"].layout = gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5,
                                                            cga_layout=cga_a_scale)
        nargs["b_scale_desc"].layout = gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5,
                                                            cga_layout=cga_b_scale)
    else:
        no_swizzle = gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5)
        nargs["a_scale_desc"].layout = no_swizzle
        nargs["b_scale_desc"].layout = no_swizzle


@gluon.jit
def unswizzle_scales_shared_memory(smem, BLOCK_MN: gl.constexpr, BLOCK_K: gl.constexpr, VEC_SIZE: gl.constexpr):
    smem = smem.reshape((smem.shape[1], smem.shape[2], 32, 4, 4))
    smem = smem.permute((0, 3, 2, 1, 4))
    return smem.reshape((BLOCK_MN, BLOCK_K // VEC_SIZE))


@gluon.jit
def async_mma_scaled_impl(a_smem, b_smem, a_scale_smem, b_scale_smem, acc_tmem, mma_bar, use_acc, pred):
    A_ELEM_PER_BYTE: gl.constexpr = 2 if a_smem.dtype == gl.uint8 else 1
    BLOCK_M: gl.constexpr = a_smem.shape[0]
    BLOCK_N: gl.constexpr = b_smem.shape[0]
    BLOCK_K: gl.constexpr = a_smem.shape[1] * A_ELEM_PER_BYTE
    # Recall we use `uint8` to represent fp4 elements.
    VEC_SIZE: gl.constexpr = 32 if a_scale_smem.dtype == gl.uint8 else 16

    a_scale = unswizzle_scales_shared_memory(a_scale_smem, BLOCK_M, BLOCK_K, VEC_SIZE)
    b_scale = unswizzle_scales_shared_memory(b_scale_smem, BLOCK_N, BLOCK_K, VEC_SIZE)

    # We don't need to hoist the scales tensor memory allocations outside of the loop,
    # so we can pull them into this helper function.
    two_ctas: gl.constexpr = acc_tmem.type.layout.two_ctas
    a_scale_layout: gl.constexpr = TensorMemoryScalesLayout(cga_layout=[[1, 0]] if two_ctas else [])
    b_scale_layout: gl.constexpr = TensorMemoryScalesLayout(cga_layout=[[0, 0]] if two_ctas else [])
    a_scale_tmem = allocate_tensor_memory(a_scale.dtype, a_scale.type.shape, a_scale_layout)
    b_scale_tmem = allocate_tensor_memory(b_scale.dtype, b_scale.type.shape, b_scale_layout)
    tcgen05_copy(a_scale, a_scale_tmem)
    tcgen05_copy(b_scale, b_scale_tmem)

    a_format: gl.constexpr = "e2m1" if a_smem.dtype == gl.uint8 else "e4m3"
    b_format: gl.constexpr = "e2m1" if b_smem.dtype == gl.uint8 else "e4m3"
    tcgen05_mma_scaled(a_smem, b_smem.permute((1, 0)), acc_tmem, a_scale_tmem, b_scale_tmem, a_format, b_format,
                       use_acc=use_acc, pred=pred, multicast=True, mbarriers=[mma_bar])


# This helper function computes all the load indexing and issues the async loads
# based on the current `pid_m`, `pid_n`, and `k` indices. The compiler will run
# loop-invariant code motion to hoist code that does not depend on `k`, like
# `pid_m * BLOCK_M`, outside of the inner loop, so we can safely abstract the
# load indexing without performance loss.
#
# Encapsulating the load indexing logic will help keep our pipelined kernel code
# clean, as pipelining can get messy.
@gluon.jit
def issue_loads(producer, pid_m, pid_n, k, a_desc, b_desc, a_scale_desc, b_scale_desc, a_bufs, b_bufs, a_scale_bufs,
                b_scale_bufs, bars, pred):
    A_ELEM_PER_BYTE: gl.constexpr = 2 if a_desc.dtype == gl.uint8 else 1
    B_ELEM_PER_BYTE: gl.constexpr = 2 if b_desc.dtype == gl.uint8 else 1
    BLOCK_M: gl.constexpr = a_desc.block_shape[0]
    BLOCK_N: gl.constexpr = b_desc.block_shape[0]
    BLOCK_K: gl.constexpr = a_desc.block_shape[1] * A_ELEM_PER_BYTE
    REP_M: gl.constexpr = a_scale_desc.block_shape[1]
    REP_N: gl.constexpr = b_scale_desc.block_shape[1]
    A_REP_K: gl.constexpr = a_scale_desc.block_shape[2]
    B_REP_K: gl.constexpr = b_scale_desc.block_shape[2]

    off_m = pid_m * BLOCK_M
    off_n = pid_n * BLOCK_N
    off_m_a_scale = pid_m * REP_M
    off_n_b_scale = pid_n * REP_N
    off_k_a = k // A_ELEM_PER_BYTE
    off_k_b = k // B_ELEM_PER_BYTE
    off_k_a_scale = (k // BLOCK_K) * A_REP_K
    off_k_b_scale = (k // BLOCK_K) * B_REP_K

    index = producer.index
    bar = bars.index(index)
    mbarrier.expect(
        bar, a_desc.nbytes_per_cta + b_desc.nbytes_per_cta + a_scale_desc.nbytes_per_cta + b_scale_desc.nbytes_per_cta,
        pred)
    tma.async_load(a_desc, [off_m, off_k_a], bar, a_bufs.index(index), pred, multicast=True)
    tma.async_load(b_desc, [off_n, off_k_b], bar, b_bufs.index(index), pred, multicast=True)
    tma.async_load(a_scale_desc, [0, off_m_a_scale, off_k_a_scale, 0, 0], bar, a_scale_bufs.index(index), pred,
                   multicast=True)
    tma.async_load(b_scale_desc, [0, off_n_b_scale, off_k_b_scale, 0, 0], bar, b_scale_bufs.index(index), pred,
                   multicast=True)
    return producer.next(pred)


@gluon.jit
def issue_mma(consumer, c_bars, a_bufs, b_bufs, a_scale_bufs, b_scale_bufs, producer, p_bars, acc_tmem, use_acc, pred):
    c_index = consumer.index
    mbarrier.wait(c_bars.index(c_index), consumer.phase, pred)
    async_mma_scaled_impl(a_bufs.index(c_index), b_bufs.index(c_index), a_scale_bufs.index(c_index),
                          b_scale_bufs.index(c_index), acc_tmem, p_bars.index(producer.index), use_acc, pred)
    return consumer.next(pred), producer.next(pred)


@gluon.aggregate
class Counter:
    index: gl.tensor
    phase: gl.tensor
    num_barriers: gl.constexpr

    @gluon.jit
    def create(phase, num_barriers: gl.constexpr):
        return Counter(gl.to_tensor(0), gl.to_tensor(phase), num_barriers)

    @gluon.must_use_result
    @gluon.jit
    def next(self, pred=True):
        incr = self.index + gl.where(pred, 1, 0)
        rollover = incr == self.num_barriers
        index = gl.where(rollover, 0, incr)
        phase = gl.where(rollover, self.phase ^ 1, self.phase)
        return Counter(index, phase, self.num_barriers)


@gluon.aggregate
class ClcTileSchedulerConsumer:
    has_work: gl.tensor
    tile_id: gl.tensor
    pid_m: gl.tensor
    pid_n: gl.tensor
    num_pid_m: gl.tensor
    num_pid_n: gl.tensor
    TILE_M: gl.constexpr
    TILE_N: gl.constexpr
    MINOR_DIM: gl.constexpr
    GRID_TILE_WIDTH: gl.constexpr
    clc_result_buffers: gl.shared_memory_descriptor
    clc_barriers: gl.shared_memory_descriptor
    clc_planar_pid_buffers: gl.shared_memory_descriptor
    clc_planar_ready_bars: gl.shared_memory_descriptor
    clc_consumed_bars: gl.shared_memory_descriptor
    counter: Counter
    consumed_counter: Counter

    @gluon.jit
    def initialize(M, N, TILE_M: gl.constexpr, TILE_N: gl.constexpr, MINOR_DIM: gl.constexpr,
                   GRID_TILE_WIDTH: gl.constexpr, clc_result_buffers, clc_barriers, clc_planar_pid_buffers,
                   clc_planar_ready_bars, clc_consumed_bars):
        tile_id = gl.program_id(axis=0)
        num_pid_m = gl.cdiv(M, TILE_M)
        num_pid_n = gl.cdiv(N, TILE_N)
        pid_m, pid_n = _planar_snake(tile_id, num_pid_m, num_pid_n, MINOR_DIM, GRID_TILE_WIDTH)
        has_work = gl.to_tensor(True)
        counter = Counter.create(0, clc_barriers.shape[0])
        consumed_counter = Counter.create(0, clc_barriers.shape[0])
        return ClcTileSchedulerConsumer(
            has_work,
            tile_id,
            pid_m,
            pid_n,
            num_pid_m,
            num_pid_n,
            TILE_M,
            TILE_N,
            MINOR_DIM,
            GRID_TILE_WIDTH,
            clc_result_buffers,
            clc_barriers,
            clc_planar_pid_buffers,
            clc_planar_ready_bars,
            clc_consumed_bars,
            counter,
            consumed_counter,
        )

    @gluon.jit
    def get_offsets(self):
        return self.pid_m * self.TILE_M, self.pid_n * self.TILE_N

    @gluon.jit
    def step(self, iteration):
        # The 0-th iteration uses the program_id as the tile_id.
        # At the end of each iteration we prefetch the next tile.
        # As such we must signal the consumed slot at the end of
        # each iteration skipping the first one.
        consumed_counter = self.consumed_counter
        if iteration > 0:
            mbarrier.arrive(self.clc_consumed_bars.index(consumed_counter.index))
            consumed_counter = consumed_counter.next()
        counter = self.counter
        barrier = self.clc_barriers.index(counter.index)
        result = self.clc_result_buffers.index(counter.index)
        mbarrier.wait(barrier, counter.phase)
        clc_res = clc.load_result(result)
        mbarrier.wait(self.clc_planar_ready_bars.index(counter.index), counter.phase)
        planar_slot = self.clc_planar_pid_buffers.index(counter.index)
        planar_layout: gl.constexpr = gl.BlockedLayout([1], [32], [gl.num_warps()], [0],
                                                       [[0]] * (gl.num_ctas().bit_length() - 1))
        packed_pid = planar_slot.load(planar_layout).reshape([])
        pid_m = ((packed_pid >> 32) & 0xFFFFFFFF).to(gl.int32)
        pid_n = (packed_pid & 0xFFFFFFFF).to(gl.int32)
        has_work = clc_res.is_canceled()
        tile_id = self.tile_id
        if has_work:
            tile_id = clc_res.program_id(0)
        return ClcTileSchedulerConsumer(
            has_work,
            tile_id,
            pid_m,
            pid_n,
            self.num_pid_m,
            self.num_pid_n,
            self.TILE_M,
            self.TILE_N,
            self.MINOR_DIM,
            self.GRID_TILE_WIDTH,
            self.clc_result_buffers,
            self.clc_barriers,
            self.clc_planar_pid_buffers,
            self.clc_planar_ready_bars,
            self.clc_consumed_bars,
            counter.next(),
            consumed_counter,
        )


@gluon.aggregate
class SpsTileScheduler:
    has_work: gl.tensor
    tile_id: gl.tensor
    pid_m: gl.tensor
    pid_n: gl.tensor
    TILE_M: gl.constexpr
    TILE_N: gl.constexpr
    MINOR_DIM: gl.constexpr
    GRID_TILE_WIDTH: gl.constexpr
    NUM_PID_M: gl.tensor
    NUM_PID_N: gl.tensor

    @gluon.jit
    def initialize(TILE_M: gl.constexpr, TILE_N: gl.constexpr, MINOR_DIM: gl.constexpr, GRID_TILE_WIDTH: gl.constexpr,
                   NUM_PID_M, NUM_PID_N):
        tile_id = gl.program_id(axis=0)
        has_work = tile_id < NUM_PID_M * NUM_PID_N
        pid_m = gl.to_tensor(0)
        pid_n = gl.to_tensor(0)
        if has_work:
            pid_m, pid_n = _planar_snake(tile_id, NUM_PID_M, NUM_PID_N, MINOR_DIM, GRID_TILE_WIDTH)
        return SpsTileScheduler(has_work, tile_id, pid_m, pid_n, TILE_M, TILE_N, MINOR_DIM, GRID_TILE_WIDTH, NUM_PID_M,
                                NUM_PID_N)

    @gluon.jit
    def get_offsets(self):
        return self.pid_m * self.TILE_M, self.pid_n * self.TILE_N

    @gluon.jit
    def step(self, iteration):
        next_tile_id = self.tile_id + gl.num_programs(axis=0)
        has_work = next_tile_id < self.NUM_PID_M * self.NUM_PID_N
        pid_m = gl.to_tensor(0)
        pid_n = gl.to_tensor(0)
        if has_work:
            pid_m, pid_n = _planar_snake(next_tile_id, self.NUM_PID_M, self.NUM_PID_N, self.MINOR_DIM,
                                         self.GRID_TILE_WIDTH)
        return SpsTileScheduler(has_work, next_tile_id, pid_m, pid_n, self.TILE_M, self.TILE_N, self.MINOR_DIM,
                                self.GRID_TILE_WIDTH, self.NUM_PID_M, self.NUM_PID_N)


# ---------------------------------------------------------------------------
# Partitions
# ---------------------------------------------------------------------------


@gluon.aggregate
class PartitionArgs:
    a_desc: tma.tensor_descriptor
    b_desc: tma.tensor_descriptor
    c_desc: tma.tensor_descriptor
    a_scale_desc: tma.tensor_descriptor
    b_scale_desc: tma.tensor_descriptor
    a_bufs: gl.shared_memory_descriptor
    b_bufs: gl.shared_memory_descriptor
    a_scale_bufs: gl.shared_memory_descriptor
    b_scale_bufs: gl.shared_memory_descriptor
    load_empty_bars: gl.shared_memory_descriptor
    load_ready_bars: gl.shared_memory_descriptor
    acc_bufs: tensor_memory_descriptor
    acc_empty_bars: gl.shared_memory_descriptor
    acc_ready_bars: gl.shared_memory_descriptor
    clc_result_buffers: gl.shared_memory_descriptor
    clc_barriers: gl.shared_memory_descriptor
    clc_planar_pid_buffers: gl.shared_memory_descriptor
    clc_planar_ready_bars: gl.shared_memory_descriptor
    clc_consumed_bars: gl.shared_memory_descriptor
    MINOR_DIM: gl.constexpr
    GRID_TILE_WIDTH: gl.constexpr
    NUM_PID_M: gl.tensor
    NUM_PID_N: gl.tensor
    scheduler: gl.constexpr

    @gluon.jit
    def get_clc_consumer(self):
        return ClcTileSchedulerConsumer.initialize(
            self.c_desc.shape[0],
            self.c_desc.shape[1],
            self.a_desc.block_shape[0],
            self.b_desc.block_shape[0],
            self.MINOR_DIM,
            self.GRID_TILE_WIDTH,
            self.clc_result_buffers,
            self.clc_barriers,
            self.clc_planar_pid_buffers,
            self.clc_planar_ready_bars,
            self.clc_consumed_bars,
        )

    @gluon.jit
    def get_sps_scheduler(self):
        return SpsTileScheduler.initialize(
            self.a_desc.block_shape[0],
            self.b_desc.block_shape[0],
            self.MINOR_DIM,
            self.GRID_TILE_WIDTH,
            self.NUM_PID_M,
            self.NUM_PID_N,
        )


@gluon.jit
def mma_scaled_load_partition(p):
    A_ELEM_PER_BYTE: gl.constexpr = 2 if p.a_desc.dtype == gl.uint8 else 1
    BLOCK_K: gl.constexpr = p.a_desc.block_shape[1] * A_ELEM_PER_BYTE
    K = p.a_desc.shape[1] * A_ELEM_PER_BYTE
    state = Counter.create(1, p.load_empty_bars.shape[0])
    if p.scheduler == SCHEDULER_CLC:
        scheduler = p.get_clc_consumer()
    else:
        scheduler = p.get_sps_scheduler()
    i = 0
    while scheduler.has_work:
        for k in range(0, K, BLOCK_K):
            mbarrier.wait(p.load_empty_bars.index(state.index), state.phase)
            state = issue_loads(state, scheduler.pid_m, scheduler.pid_n, k, p.a_desc, p.b_desc, p.a_scale_desc,
                                p.b_scale_desc, p.a_bufs, p.b_bufs, p.a_scale_bufs, p.b_scale_bufs, p.load_ready_bars,
                                pred=True)
        scheduler = scheduler.step(i)
        i += 1


@gluon.jit
def mma_scaled_mma_partition(p):
    A_ELEM_PER_BYTE: gl.constexpr = 2 if p.a_desc.dtype == gl.uint8 else 1
    BLOCK_K: gl.constexpr = p.a_desc.block_shape[1] * A_ELEM_PER_BYTE
    K = p.a_desc.shape[1] * A_ELEM_PER_BYTE
    load_state = Counter.create(0, p.load_empty_bars.shape[0])
    acc_state = Counter.create(1, p.acc_empty_bars.shape[0])
    if p.scheduler == SCHEDULER_CLC:
        scheduler = p.get_clc_consumer()
    else:
        scheduler = p.get_sps_scheduler()
    i = 0
    while scheduler.has_work:
        mbarrier.wait(p.acc_empty_bars.index(acc_state.index), acc_state.phase)
        acc_buf = p.acc_bufs.index(acc_state.index)
        use_acc = False
        for k in range(0, K, BLOCK_K):
            _, load_state = issue_mma(load_state, p.load_ready_bars, p.a_bufs, p.b_bufs, p.a_scale_bufs, p.b_scale_bufs,
                                      load_state, p.load_empty_bars, acc_buf, use_acc, pred=True)
            use_acc = True
        tcgen05_commit(p.acc_ready_bars.index(acc_state.index))
        acc_state = acc_state.next()
        scheduler = scheduler.step(i)
        i += 1


@gluon.jit
def mma_scaled_epilogue_partition(p):
    tile_m: gl.constexpr = p.c_desc.block_shape[0]
    BLOCK_N: gl.constexpr = p.b_desc.block_shape[0]
    EPILOGUE_BLOCK_N: gl.constexpr = p.c_desc.block_shape[1]
    subtile_factor: gl.constexpr = BLOCK_N // EPILOGUE_BLOCK_N
    subtile_stages: gl.constexpr = 1 if subtile_factor == 1 else 2
    acc_state = Counter.create(0, p.acc_empty_bars.shape[0])
    acc_smems = gl.allocate_shared_memory(p.c_desc.dtype, [subtile_stages, tile_m, EPILOGUE_BLOCK_N], p.c_desc.layout)
    sub_acc_state = Counter.create(0, subtile_stages)
    if p.scheduler == SCHEDULER_CLC:
        scheduler = p.get_clc_consumer()
    else:
        scheduler = p.get_sps_scheduler()
    i = 0
    while scheduler.has_work:
        off_m, off_n = scheduler.get_offsets()
        mbarrier.wait(p.acc_ready_bars.index(acc_state.index), acc_state.phase)
        acc_buf = p.acc_bufs.index(acc_state.index)

        for s in gl.static_range(subtile_factor):
            acc_sub = acc_buf.slice(EPILOGUE_BLOCK_N * s, EPILOGUE_BLOCK_N)
            acc_smem = acc_smems.index(sub_acc_state.index)
            acc = acc_sub.load().to(p.c_desc.dtype)
            tma.store_wait(pendings=subtile_stages - 1)
            acc_smem.store(acc)
            tma.async_store(p.c_desc, [off_m, off_n + EPILOGUE_BLOCK_N * s], acc_smem)
            sub_acc_state = sub_acc_state.next()
        mbarrier.arrive(p.acc_empty_bars.index(acc_state.index), count=1)
        acc_state = acc_state.next()
        scheduler = scheduler.step(i)
        i += 1
    tma.store_wait(0)


@gluon.jit
def mma_scaled_clc_partition(p):
    TILE_M: gl.constexpr = p.a_desc.block_shape[0]
    TILE_N: gl.constexpr = p.b_desc.block_shape[0]
    has_work = gl.to_tensor(True)
    num_pid_m = gl.cdiv(p.c_desc.shape[0], TILE_M)
    num_pid_n = gl.cdiv(p.c_desc.shape[1], TILE_N)
    state = Counter.create(0, p.clc_barriers.shape[0])
    consumed_state = Counter.create(1, p.clc_barriers.shape[0])
    ACC_STAGES: gl.constexpr = p.clc_barriers.shape[0]
    i = 0
    while has_work:
        # Reuse the slot only after all consumer partitions signaled consumed.
        mbarrier.wait(p.clc_consumed_bars.index(consumed_state.index), consumed_state.phase, pred=(i >= ACC_STAGES))
        barrier = p.clc_barriers.index(state.index)
        result = p.clc_result_buffers.index(state.index)
        # 16: clc.try_cancel has a `.b128` modifier
        mbarrier.expect(barrier, 16, from_cta=0x0)
        clc.try_cancel(result, barrier)
        mbarrier.wait(barrier, state.phase)
        clc_res = clc.load_result(result)
        has_work = clc_res.is_canceled()
        pid_m = gl.to_tensor(0)
        pid_n = gl.to_tensor(0)
        if has_work:
            tile_id = clc_res.program_id(0)
            pid_m, pid_n = _planar_snake(tile_id, num_pid_m, num_pid_n, p.MINOR_DIM, p.GRID_TILE_WIDTH)
        packed_pid = (pid_m.to(gl.int64) << 32) | (pid_n.to(gl.int64) & 0xFFFFFFFF)
        planar_slot = p.clc_planar_pid_buffers.index(state.index)
        planar_layout: gl.constexpr = gl.BlockedLayout([1], [32], [gl.num_warps()], [0],
                                                       [[0]] * (gl.num_ctas().bit_length() - 1))
        planar_slot.store(gl.full([1], packed_pid, gl.int64, layout=planar_layout))
        mbarrier.arrive(p.clc_planar_ready_bars.index(state.index))
        state = state.next()
        consumed_state = consumed_state.next()
        i += 1


# ---------------------------------------------------------------------------
# Kernel
# ---------------------------------------------------------------------------


@gluon.jit
def mma_scaled_warp_specialized_kernel(a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, M, N, K, A_ELEM_PER_BYTE,
                                       num_buffers: gl.constexpr, BLOCK_M: gl.constexpr, BLOCK_N: gl.constexpr,
                                       BLOCK_K: gl.constexpr, EPILOGUE_BLOCK_N: gl.constexpr,
                                       num_acc_buffers: gl.constexpr, GRID_MINOR_DIM: gl.constexpr,
                                       GRID_TILE_WIDTH: gl.constexpr, CGA_LAYOUT: gl.constexpr,
                                       scheduler: gl.constexpr):
    NUM_CTAS: gl.constexpr = gl.num_ctas()
    TWO_CTAS: gl.constexpr = NUM_CTAS > 1
    BLOCK_M_PER_CTA: gl.constexpr = BLOCK_M // NUM_CTAS
    gl.static_assert(BLOCK_M_PER_CTA == 64 or BLOCK_M_PER_CTA == 128)
    N_PARTITIONS: gl.constexpr = 4 if scheduler == SCHEDULER_CLC else 3

    a_bufs = gl.allocate_shared_memory(a_desc.dtype, [num_buffers] + a_desc.block_shape, a_desc.layout)
    b_bufs = gl.allocate_shared_memory(b_desc.dtype, [num_buffers] + b_desc.block_shape, b_desc.layout)
    a_scale_bufs = gl.allocate_shared_memory(a_scale_desc.dtype, [num_buffers] + a_scale_desc.block_shape,
                                             a_scale_desc.layout)
    b_scale_bufs = gl.allocate_shared_memory(b_scale_desc.dtype, [num_buffers] + b_scale_desc.block_shape,
                                             b_scale_desc.layout)

    tmem_layout: gl.constexpr = TensorMemoryLayout([BLOCK_M_PER_CTA, BLOCK_N], col_stride=1, cga_layout=CGA_LAYOUT,
                                                   two_ctas=TWO_CTAS)
    acc_bufs = allocate_tensor_memory(gl.float32, [num_acc_buffers, BLOCK_M, BLOCK_N], tmem_layout)

    mma_barrier_count: gl.constexpr = tcgen05_mma_barrier_count(
        [a_bufs.index(0), b_bufs.index(0),
         a_scale_bufs.index(0), b_scale_bufs.index(0)], multicast=True, two_ctas=acc_bufs.index(0).type.layout.two_ctas)

    load_empty_bars = mbarrier.allocate_mbarrier(batch=num_buffers)
    load_ready_bars = mbarrier.allocate_mbarrier(batch=num_buffers, two_ctas=TWO_CTAS)
    for i in gl.static_range(num_buffers):
        mbarrier.init(load_empty_bars.index(i), count=mma_barrier_count)
        mbarrier.init(load_ready_bars.index(i), count=1)

    acc_empty_bars = mbarrier.allocate_mbarrier(batch=num_acc_buffers, two_ctas=TWO_CTAS)
    acc_ready_bars = mbarrier.allocate_mbarrier(batch=num_acc_buffers)
    for i in gl.static_range(num_acc_buffers):
        mbarrier.init(acc_empty_bars.index(i), count=1)
        mbarrier.init(acc_ready_bars.index(i), count=1)

    clc_barriers = mbarrier.allocate_mbarrier(batch=num_acc_buffers)
    clc_planar_ready_bars = mbarrier.allocate_mbarrier(batch=num_acc_buffers)
    cga_layout_clc: gl.constexpr = [[0]] * (gl.num_ctas().bit_length() - 1)
    clc_layout: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, [0], cga_layout=cga_layout_clc)
    clc_consumed_bars = gl.allocate_shared_memory(gl.int64, [num_acc_buffers, 1], clc_layout)
    if scheduler == SCHEDULER_CLC:
        for i in gl.static_range(num_acc_buffers):
            mbarrier.init(clc_barriers.index(i), count=1)
            mbarrier.init(clc_planar_ready_bars.index(i), count=1)
            mbarrier.init(clc_consumed_bars.index(i), count=N_PARTITIONS - 1)

    clc_result_buffers = gl.allocate_shared_memory(gl.int64, [clc_barriers.shape[0], 2], clc_layout)
    clc_planar_pid_buffers = gl.allocate_shared_memory(gl.int64, [clc_barriers.shape[0], 1], clc_layout)

    p = PartitionArgs(a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, a_bufs, b_bufs, a_scale_bufs, b_scale_bufs,
                      load_empty_bars, load_ready_bars, acc_bufs, acc_empty_bars, acc_ready_bars, clc_result_buffers,
                      clc_barriers, clc_planar_pid_buffers, clc_planar_ready_bars, clc_consumed_bars, GRID_MINOR_DIM,
                      GRID_TILE_WIDTH, gl.cdiv(M, BLOCK_M), gl.cdiv(N, BLOCK_N), scheduler)

    if scheduler == SCHEDULER_CLC:
        gl.warp_specialize([
            (mma_scaled_epilogue_partition, (p, )),
            (mma_scaled_mma_partition, (p, )),
            (mma_scaled_load_partition, (p, )),
            (mma_scaled_clc_partition, (p, )),
        ], [1, 1, 1], [24, 24, 24])
    else:
        gl.warp_specialize([
            (mma_scaled_epilogue_partition, (p, )),
            (mma_scaled_mma_partition, (p, )),
            (mma_scaled_load_partition, (p, )),
        ], [1, 1], [24, 24])


mma_scaled_kernel = triton.autotune(
    configs=mma_scaled_get_configs(pre_hook=mma_scaled_tma_set_block_size_hook),
    key=["M", "N", "K", "A_ELEM_PER_BYTE"],
)(mma_scaled_warp_specialized_kernel)

mma_scaled_1cta_kernel = triton.autotune(
    configs=mma_scaled_get_configs(pre_hook=mma_scaled_tma_set_block_size_hook, cga_layouts=[()]),
    key=["M", "N", "K", "A_ELEM_PER_BYTE"],
)(mma_scaled_warp_specialized_kernel)

mma_scaled_2cta_kernel = triton.autotune(
    configs=mma_scaled_get_configs(pre_hook=mma_scaled_tma_set_block_size_hook, cga_layouts=[((1, 0), )]),
    key=["M", "N", "K", "A_ELEM_PER_BYTE"],
)(mma_scaled_warp_specialized_kernel)

# ---------------------------------------------------------------------------
# Wrapper
# ---------------------------------------------------------------------------


def make_dummy_descriptors(A, B, A_scale, B_scale, out_dtype, M, N):
    """Create TMA descriptors with dummy block shapes; the hook sets the real ones."""
    dummy_block_2d = [1, 1]
    dummy_layout_2d = gl.NVMMASharedLayout.get_default_for(dummy_block_2d, gl.float8e4nv)
    a_desc = TensorDescriptor.from_tensor(A, dummy_block_2d, dummy_layout_2d)
    b_desc = TensorDescriptor.from_tensor(B, dummy_block_2d, dummy_layout_2d)

    C = torch.empty(M, N, device="cuda", dtype=out_dtype)
    C_dtype = getattr(gl, str(out_dtype).split('.')[1])
    c_layout = gl.NVMMASharedLayout.get_default_for(dummy_block_2d, C_dtype)
    c_desc = TensorDescriptor.from_tensor(C, dummy_block_2d, c_layout)

    A_scale_5d = A_scale.reshape(1, A_scale.shape[0], A_scale.shape[1], 2, 256)
    B_scale_5d = B_scale.reshape(1, B_scale.shape[0], B_scale.shape[1], 2, 256)
    dummy_block_5d = [1, 1, 1, 2, 256]
    dummy_layout_5d = gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5)
    a_scale_desc = TensorDescriptor.from_tensor(A_scale_5d, dummy_block_5d, dummy_layout_5d)
    b_scale_desc = TensorDescriptor.from_tensor(B_scale_5d, dummy_block_5d, dummy_layout_5d)

    return a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc


def mma_scaled_warp_specialized_grid(M, N, BLOCK_M, BLOCK_N, num_ctas, scheduler, device):
    num_pid = triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N)
    if scheduler == SCHEDULER_CLC:
        return (num_pid, )
    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
    return (max(1, min(num_pid, sm_count // num_ctas)), )


def mma_scaled_warp_specialized(A, B, A_scale, B_scale, VEC_SIZE, scheduler, GRID_MINOR_DIM=0, GRID_TILE_WIDTH=4,
                                out_dtype=torch.float16, BLOCK_M=128, BLOCK_N=256, BLOCK_K=None, EPILOGUE_BLOCK_N=None,
                                num_buffers=3, acc_buffers=None, num_ctas=1):
    """Warp-specialized block-scale MMA (supports 1CTA and 2CTA)."""
    if BLOCK_K is None:
        BLOCK_K = 128 if torch.float8_e4m3fn in [A.dtype, B.dtype] else 256
    if EPILOGUE_BLOCK_N is None:
        EPILOGUE_BLOCK_N = BLOCK_N
    if acc_buffers is None:
        acc_buffers = 2 if BLOCK_N < 256 else 1

    M, N = A.shape[0], B.shape[0]
    IS_FP4_A = A.dtype == torch.uint8
    K = A.shape[1] * (2 if IS_FP4_A else 1)
    cga_layout = ((1, 0), ) if num_ctas > 1 else ()

    A_desc, B_desc, C_desc, A_scale_desc, B_scale_desc = make_dummy_descriptors(A, B, A_scale, B_scale, out_dtype, M, N)

    mma_scaled_tma_set_block_size_hook({
        "a_desc": A_desc,
        "b_desc": B_desc,
        "c_desc": C_desc,
        "a_scale_desc": A_scale_desc,
        "b_scale_desc": B_scale_desc,
        "BLOCK_M": BLOCK_M,
        "BLOCK_N": BLOCK_N,
        "BLOCK_K": BLOCK_K,
        "EPILOGUE_BLOCK_N": EPILOGUE_BLOCK_N,
        "CGA_LAYOUT": cga_layout,
    })

    grid = mma_scaled_warp_specialized_grid(M, N, BLOCK_M, BLOCK_N, num_ctas, scheduler, A.device)
    A_ELEM_PER_BYTE = 2 if IS_FP4_A else 1
    mma_scaled_warp_specialized_kernel[grid](
        A_desc,
        B_desc,
        C_desc,
        A_scale_desc,
        B_scale_desc,
        M,
        N,
        K,
        A_ELEM_PER_BYTE,
        num_buffers,
        BLOCK_M,
        BLOCK_N,
        BLOCK_K,
        EPILOGUE_BLOCK_N,
        acc_buffers,
        GRID_MINOR_DIM,
        GRID_TILE_WIDTH,
        cga_layout,
        scheduler=scheduler,
        num_ctas=num_ctas,
    )
    return C_desc.base


def mma_scaled_matmul(A, B, A_scale, B_scale, VEC_SIZE, out_dtype=torch.float16, num_ctas=None):
    """Autotuned block-scaled matmul.

    Args:
        num_ctas: None = autotune across all configs (1CTA and 2CTA),
                  1 = autotune 1CTA configs only,
                  2 = autotune 2CTA configs only.
    """
    M, N = A.shape[0], B.shape[0]
    IS_FP4_A = A.dtype == torch.uint8
    A_ELEM_PER_BYTE = 2 if IS_FP4_A else 1
    K = A.shape[1] * A_ELEM_PER_BYTE

    a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc = make_dummy_descriptors(A, B, A_scale, B_scale, out_dtype, M, N)

    def grid(meta):
        return mma_scaled_warp_specialized_grid(M, N, meta["BLOCK_M"], meta["BLOCK_N"], 2**len(meta["CGA_LAYOUT"]),
                                                meta["scheduler"], A.device)

    kernel = {None: mma_scaled_kernel, 1: mma_scaled_1cta_kernel, 2: mma_scaled_2cta_kernel}[num_ctas]
    kernel[grid](a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, M, N, K, A_ELEM_PER_BYTE)
    return c_desc.base


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("K", [128, 640, 704, 1152, 4096])
@pytest.mark.parametrize("M, N", [(2048, 2048), (500, 600), (256, 256), (128, 128), (8192, 8192)])
@pytest.mark.parametrize("a_format, b_format",
                         list(itertools.product(["mxfp8", "mxfp4"], repeat=2)) + [("nvfp4", "nvfp4")])
@pytest.mark.parametrize("num_ctas, BLOCK_N, EPILOGUE_BLOCK_N, num_buffers", [
    (2, 256, 256, 4),
    (2, 256, 64, 5),
    (2, 128, 64, 6),
    (1, 256, 256, 3),
    (1, 256, 64, 3),
    (1, 128, 64, 5),
])
@pytest.mark.parametrize("scheduler", [SCHEDULER_CLC, SCHEDULER_SPS])
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
@pytest.mark.enable_warmup(min_capability=10)
def test_mma_scaled_warp_specialized(M, N, K, a_format, b_format, num_ctas, BLOCK_N, EPILOGUE_BLOCK_N, num_buffers,
                                     scheduler):
    if a_format != b_format and K % 128 != 0:
        pytest.skip("fp4 packed tensor descriptor requires K to be a multiple of 128")
    BLOCK_M = 256 if num_ctas > 1 else 128
    torch.manual_seed(0)
    A, A_scale, A_ref = random_quantized_tensor(M, K, a_format)
    B, B_scale, B_ref = random_quantized_tensor(N, K, b_format)
    VEC_SIZE = 16 if a_format == "nvfp4" else 32
    A_scale = swizzle_scales_packed_block(A_scale)
    B_scale = swizzle_scales_packed_block(B_scale)
    C_ref = A_ref @ B_ref.T
    C = mma_scaled_warp_specialized(A, B, A_scale, B_scale, VEC_SIZE, scheduler=scheduler, BLOCK_M=BLOCK_M,
                                    BLOCK_N=BLOCK_N, EPILOGUE_BLOCK_N=EPILOGUE_BLOCK_N, num_buffers=num_buffers,
                                    num_ctas=num_ctas)
    torch.testing.assert_close(C_ref, C.to(torch.float32), atol=1e-3, rtol=1e-3)


# ---------------------------------------------------------------------------
# Benchmark
# ---------------------------------------------------------------------------

CUBLAS_FORMATS = {"mxfp8", "nvfp4"}


def cublas_block_scaled_matmul(A, B, A_scale_flat, B_scale_flat, fmt):
    """cuBLAS block-scaled matmul. Supports mxfp8 and nvfp4 (mxfp4 not supported by cuBLAS)."""
    M, N = A.shape[0], B.shape[0]
    output = torch.empty((M, N), dtype=torch.float16, device="cuda")
    if fmt == "mxfp8":
        triton.testing.cublas().block_scaled_matmul_mxfp8(A, B, output, A_scale_flat, B_scale_flat)
    elif fmt == "nvfp4":
        triton.testing.cublas().block_scaled_matmul_nvfp4(A, B, output, A_scale_flat, B_scale_flat)
    else:
        raise ValueError(f"cuBLAS does not support format: {fmt}")
    return output


ALL_FORMATS = [("mxfp8", "mxfp8"), ("nvfp4", "nvfp4"), ("mxfp8", "mxfp4"), ("mxfp4", "mxfp4")]

MNK_VALS = [8192, 16384, 32768]

BEST_1CTA_CONFIG = dict(BLOCK_M=128, BLOCK_N=256, EPILOGUE_BLOCK_N=64, num_buffers=3, num_ctas=1, GRID_MINOR_DIM=1,
                        GRID_TILE_WIDTH=8)
BEST_2CTA_CONFIG = dict(BLOCK_M=256, BLOCK_N=256, EPILOGUE_BLOCK_N=64, num_buffers=5, num_ctas=2, GRID_MINOR_DIM=0,
                        GRID_TILE_WIDTH=8)


def scheduler_for_k(K):
    return SCHEDULER_SPS if K <= 8192 else SCHEDULER_CLC


def fixed_config_for_mnk(config, MNK):
    config = dict(config)
    config["scheduler"] = scheduler_for_k(MNK)
    return config


def make_fn(variant, A, B, A_scale, B_scale, VEC_SIZE, a_format, MNK, use_autotuned=False):
    """Build the callable for a given variant (1cta, 2cta, or cublas)."""
    if variant == "2cta":
        if use_autotuned:
            return lambda: mma_scaled_matmul(A, B, A_scale, B_scale, VEC_SIZE, num_ctas=2)
        return lambda: mma_scaled_warp_specialized(A, B, A_scale, B_scale, VEC_SIZE,
                                                   **fixed_config_for_mnk(BEST_2CTA_CONFIG, MNK))
    elif variant == "1cta":
        if use_autotuned:
            return lambda: mma_scaled_matmul(A, B, A_scale, B_scale, VEC_SIZE, num_ctas=1)
        return lambda: mma_scaled_warp_specialized(A, B, A_scale, B_scale, VEC_SIZE,
                                                   **fixed_config_for_mnk(BEST_1CTA_CONFIG, MNK))
    elif variant == "cublas":
        A_scale_flat = A_scale.contiguous().flatten()
        B_scale_flat = B_scale.contiguous().flatten()

        def cublas_fn():
            return cublas_block_scaled_matmul(A, B, A_scale_flat, B_scale_flat, a_format)

        return cublas_fn
    else:
        raise ValueError(f"Unknown variant: {variant}")


def make_tensors(MNK, a_format, b_format):
    """Allocate and prepare input tensors for a given size and format."""
    M = N = K = MNK
    torch.manual_seed(0)
    A, A_scale, _ = random_quantized_tensor(M, K, a_format)
    B, B_scale, _ = random_quantized_tensor(N, K, b_format)
    A_scale = swizzle_scales_packed_block(A_scale)
    B_scale = swizzle_scales_packed_block(B_scale)
    VEC_SIZE = 16 if a_format == "nvfp4" else 32
    return A, B, A_scale, B_scale, VEC_SIZE


def get_variants(a_format, b_format):
    """Return the list of variants available for a given format pair."""
    has_cublas = a_format == b_format and a_format in CUBLAS_FORMATS
    return ["1cta", "2cta", "cublas"] if has_cublas else ["1cta", "2cta"]


def print_table(label, variants, mnk_vals, results):
    """Print a formatted benchmark table with optional ratio columns."""
    has_cublas = "cublas" in variants
    col_w = 16
    header = f"{'MNK':>8}"
    header += f"  {'1cta (TFLOPS)':>{col_w}}"
    header += f"  {'2cta (TFLOPS)':>{col_w}}"
    header += f"  {'2cta/1cta':>{col_w}}"
    if has_cublas:
        header += f"  {'cublas (TFLOPS)':>{col_w}}"
        header += f"  {'2cta/cublas':>{col_w}}"
    print(f"block-scale-matmul-{label}:")
    print(header)
    for MNK in mnk_vals:
        t1 = results.get((label, "1cta", MNK))
        t2 = results.get((label, "2cta", MNK))
        ratio_2v1 = t2 / t1 if t1 and t2 else 0.0
        row = f"{MNK:>8}"
        row += f"  {t1:>{col_w}.1f}" if t1 else f"  {'--':>{col_w}}"
        row += f"  {t2:>{col_w}.1f}" if t2 else f"  {'--':>{col_w}}"
        row += f"  {ratio_2v1:>{col_w}.2f}"
        if has_cublas:
            tc = results.get((label, "cublas", MNK))
            ratio_2vc = t2 / tc if t2 and tc else 0.0
            row += f"  {tc:>{col_w}.1f}" if tc else f"  {'--':>{col_w}}"
            row += f"  {ratio_2vc:>{col_w}.2f}"
        print(row)
    print()


def format_config(cfg):
    """Format an autotuner Config as a concise string."""
    if cfg is None:
        return "(none)"
    kw = cfg.kwargs
    parts = [
        f"BM={kw['BLOCK_M']}", f"BN={kw['BLOCK_N']}", f"BK={kw['BLOCK_K']}", f"epilogue_N={kw['EPILOGUE_BLOCK_N']}",
        f"bufs={kw['num_buffers']}", f"acc_bufs={kw['num_acc_buffers']}", f"minor={kw['GRID_MINOR_DIM']}",
        f"tile_w={kw['GRID_TILE_WIDTH']}", f"scheduler={kw['scheduler']}", f"cga={kw['CGA_LAYOUT']}"
    ]
    return ", ".join(parts)


def run_benchmark(use_autotuned=False):
    results = {}
    best_configs = {}
    for a_format, b_format in ALL_FORMATS:
        label = f"{a_format}-{b_format}"
        variants = get_variants(a_format, b_format)
        for MNK in MNK_VALS:
            A, B, A_scale, B_scale, VEC_SIZE = make_tensors(MNK, a_format, b_format)
            for variant in variants:
                if use_autotuned:
                    print(f"  {label} {variant} MNK={MNK}: ...", end="", flush=True)
                fn = make_fn(variant, A, B, A_scale, B_scale, VEC_SIZE, a_format, MNK, use_autotuned=use_autotuned)
                ms = triton.testing.do_bench(fn)
                tflops = 2.0 * MNK**3 * 1e-12 / (ms * 1e-3)
                results[(label, variant, MNK)] = tflops
                if use_autotuned:
                    print(f"\r  {label} {variant} MNK={MNK}: {tflops:.1f} TFLOPS")
                    if variant == "1cta":
                        best_configs[(label, "1cta", MNK)] = mma_scaled_1cta_kernel.best_config
                    elif variant == "2cta":
                        best_configs[(label, "2cta", MNK)] = mma_scaled_2cta_kernel.best_config

    if use_autotuned:
        largest_mnk = MNK_VALS[-1]
        print(f"\nBest autotuned configs (MNK={largest_mnk}):")
        for a_format, b_format in ALL_FORMATS:
            label = f"{a_format}-{b_format}"
            c1 = best_configs.get((label, "1cta", largest_mnk))
            c2 = best_configs.get((label, "2cta", largest_mnk))
            print(f"  {label}:")
            print(f"    1cta: {format_config(c1)}")
            print(f"    2cta: {format_config(c2)}")
        print()

    for a_format, b_format in ALL_FORMATS:
        label = f"{a_format}-{b_format}"
        variants = get_variants(a_format, b_format)
        print_table(label, variants, MNK_VALS, results)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Block-scaled matmul benchmark")
    parser.add_argument(
        "--use-autotuned",
        action="store_true",
        help="Use autotuned mma_scaled_matmul() instead of mma_scaled_warp_specialized().",
    )
    args = parser.parse_args()
    run_benchmark(use_autotuned=args.use_autotuned)
