"""
Overlapping Accumulator Block-Scaled Matrix Multiplication
==========================================================

This example demonstrates overlapping accumulator storage in TMEM for 2CTA
block scale MMA by comparing double buffered overlapping BLOCK_N=256
accumulators against a single buffered BLOCK_N=256 accumulator.
"""

import argparse
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,
    tcgen05_copy,
    tcgen05_commit,
    tcgen05_mma_barrier_count,
    tcgen05_mma_scaled,
    mbarrier,
    tma,
)


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


# ---------------------------------------------------------------------------
# Fixed Configuration
# ---------------------------------------------------------------------------

FORMATS = [("mxfp8", "mxfp8"), ("mxfp8", "mxfp4"), ("mxfp4", "mxfp4"), ("nvfp4", "nvfp4")]
BLOCK_M = 256
BLOCK_N = 256
EPILOGUE_BLOCK_N = 32
NUM_BUFFERS = 5
NUM_CTAS = 2
CGA_LAYOUT = ((1, 0), )
OUT_DTYPE = torch.float16
TMEM_COLS = gl.constexpr(512)

# ---------------------------------------------------------------------------
# Overlapping Accumulator Layout
# ---------------------------------------------------------------------------

# yapf: disable
#
# With 512 TMEM columns, double buffered BLOCK_N=256 accumulators fully consume TMEM
# leaving no room for the A and B scales required for a block scale MMA workload.
#
# Gluon example 04 uses a single buffered BLOCK_N=256 accumulator, leaving
# some performance on the table.
#
# This example keeps the double buffered BLOCK_N=256 accumulators by overlapping them
# in TMEM by exactly the size of the A and B scales as shown here:
#
#   TMEM:                                     BLOCK_N                        TMEM_COLS
#                0                              256                             512
#                [---------------------------------------------------------------)
#   left accum   [-------------------------------)
#   right accum                            [-------------------------------)
#   A scales                                                               [-)
#   B scales                                                                 [---)
#
# The overlap path demonstrates three major enhancements over example 04:
#   1) Fully allocate TMEM and then slice it to hold accumulators and scales.
#   2) Epilogue drains overlapped columns first so the next MMA can proceed.
#   3) New ping-pong style barrier guards the overlapped columns.
#
# Tile 0 MMA starts immediately and writes the left accumulator:
#
#   left accum   [xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx)
#                start -> MMA
#
# Tile 0 epilogue drains the overlap first and releases the overlap barrier:
#
#                <--------- drain from high to low
#   left accum   [xxxxxxxxxxxxxxxxxxxxxxxxx------)
#                drain overlap -> arrive overlap_bar
#
# Tile 1 MMA waits for Tile 0 epilogue to release the overlap barrier, then writes the right accumulator:
#
#   left accum   [xxxxxxxxxxxxxxxxxxxxxxxxx------)
#   right accum                            [xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx)
#                                          wait overlap_bar -> MMA
#
# Tile 1 epilogue drains the overlap first and releases the overlap barrier for Tile 2:
#
#   left accum   [xxxxxxxxxxxxxxxxxxx------------)
#   right accum                            [-----xxxxxxxxxxxxxxxxxxxxxxxxxx)
#                                          drain from low to high -------->
#                                          drain overlap -> arrive overlap_bar
#
# The overlap barrier is now released for Tile 2, but Tile 2 MMA must wait for Tile 0 epilogue
# to fully drain the left accumulator. The legacy acc_empty / acc_ready barriers are used to
# communicate empty / full accumulator semantics between MMA and epilogue partitions.
#
# yapf: enable


class TmemOverlapStoragePlan:

    def __init__(self, left_accum_offset, left_accum_cols, right_accum_offset, right_accum_cols, a_scale_offset,
                 a_scale_cols, b_scale_offset, b_scale_cols, early_release_subtile):
        self.left_accum_offset = left_accum_offset
        self.left_accum_cols = left_accum_cols
        self.right_accum_offset = right_accum_offset
        self.right_accum_cols = right_accum_cols
        self.a_scale_offset = a_scale_offset
        self.a_scale_cols = a_scale_cols
        self.b_scale_offset = b_scale_offset
        self.b_scale_cols = b_scale_cols
        self.early_release_subtile = early_release_subtile


@gluon.constexpr_function
def tmem_overlap_storage_plan(BLOCK_K, EPILOGUE_BLOCK_N, VEC_SIZE, USE_OVERLAP: gl.constexpr):
    # Hard code BLOCK_N to 256 for this example.
    BLOCK_N: gl.constexpr = 256

    # Scale views have K / VEC_SIZE columns. A scales are M-sharded
    # across CTAs, so each CTA stores one TMEM column per scale-K column.
    # B scales are duplicated across CTAs by the scale layout; per CTA,
    # they store one scale-K column group for each 128-column N block.
    scale_k = BLOCK_K // VEC_SIZE
    a_scale_cols = scale_k
    b_scale_cols = scale_k * (BLOCK_N // 128)

    # The overlap path uses double buffered BLOCK_N=256 accumulators.
    if USE_OVERLAP:
        # The overlap is the combined size of the A and B scales
        # which cannot exceed the size of BLOCK_N.
        overlap_cols = a_scale_cols + b_scale_cols
        assert overlap_cols < BLOCK_N

        # The left accumulator begins at 0 and ends at BLOCK_N.
        # The right accumulator begins at BLOCK_N minus the overlap.
        right_accum_offset = BLOCK_N - overlap_cols

        # Scale views begin after the overlapping accumulators.
        a_scale_offset = TMEM_COLS - overlap_cols

        # Identify the zero-based index of the epilogue subtile containing the last overlap column.
        early_release_subtile = (overlap_cols - 1) // EPILOGUE_BLOCK_N

    # The comparison path uses a single buffered BLOCK_N=256 accumulator.
    else:
        right_accum_offset = 0
        a_scale_offset = BLOCK_N
        early_release_subtile = 0

    return TmemOverlapStoragePlan(0, BLOCK_N, right_accum_offset, BLOCK_N, a_scale_offset, a_scale_cols,
                                  a_scale_offset + a_scale_cols, b_scale_cols, early_release_subtile)


@gluon.constexpr_function
def tmem_overlap_storage_plan_from_partition(p):
    A_ELEM_PER_BYTE: gl.constexpr = 2 if p.a_desc.dtype == gl.uint8 else 1
    BLOCK_N: gl.constexpr = p.b_desc.block_shape[0]
    assert BLOCK_N == 256
    BLOCK_K: gl.constexpr = p.a_desc.block_shape[1] * A_ELEM_PER_BYTE
    EPILOGUE_BLOCK_N: gl.constexpr = p.c_desc.block_shape[1]
    VEC_SIZE: gl.constexpr = 32 if p.a_scale_desc.dtype == gl.uint8 else 16
    return tmem_overlap_storage_plan(BLOCK_K, EPILOGUE_BLOCK_N, VEC_SIZE, p.USE_OVERLAP)


def test_tmem_overlap_accum():
    plan = tmem_overlap_storage_plan(BLOCK_K=256, EPILOGUE_BLOCK_N=32, VEC_SIZE=16, USE_OVERLAP=True)
    assert (plan.left_accum_offset, plan.left_accum_cols) == (0, 256)
    assert (plan.right_accum_offset, plan.right_accum_cols) == (208, 256)
    assert (plan.a_scale_offset, plan.a_scale_cols) == (464, 16)
    assert (plan.b_scale_offset, plan.b_scale_cols) == (480, 32)
    assert plan.early_release_subtile == 1

    plan = tmem_overlap_storage_plan(BLOCK_K=128, EPILOGUE_BLOCK_N=32, VEC_SIZE=32, USE_OVERLAP=True)
    assert (plan.right_accum_offset, plan.right_accum_cols) == (244, 256)
    assert (plan.a_scale_offset, plan.a_scale_cols) == (500, 4)
    assert (plan.b_scale_offset, plan.b_scale_cols) == (504, 8)
    assert plan.early_release_subtile == 0

    plan = tmem_overlap_storage_plan(BLOCK_K=256, EPILOGUE_BLOCK_N=32, VEC_SIZE=16, USE_OVERLAP=False)
    assert (plan.left_accum_offset, plan.left_accum_cols) == (0, 256)
    assert (plan.right_accum_offset, plan.right_accum_cols) == (0, 256)
    assert (plan.a_scale_offset, plan.a_scale_cols) == (256, 16)
    assert (plan.b_scale_offset, plan.b_scale_cols) == (272, 32)


# ---------------------------------------------------------------------------
# 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()


@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, tmem_pool, tmem_plan, 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)

    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 = tmem_pool.slice(tmem_plan.a_scale_offset,
                                   tmem_plan.a_scale_cols).reinterpret(dtype=a_scale.dtype, shape=a_scale.shape,
                                                                          layout=a_scale_layout)
    b_scale_tmem = tmem_pool.slice(tmem_plan.b_scale_offset,
                                   tmem_plan.b_scale_cols).reinterpret(dtype=b_scale.dtype, shape=b_scale.shape,
                                                                          layout=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])


@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, tmem_pool,
              tmem_plan, 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, tmem_pool, tmem_plan, 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 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
    NUM_PID_M: gl.tensor
    NUM_PID_N: gl.tensor

    @gluon.jit
    def initialize(TILE_M: gl.constexpr, TILE_N: 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 = tile_id // NUM_PID_N
            pid_n = tile_id % NUM_PID_N
        return SpsTileScheduler(has_work, tile_id, pid_m, pid_n, TILE_M, TILE_N, 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):
        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 = next_tile_id // self.NUM_PID_N
            pid_n = next_tile_id % self.NUM_PID_N
        return SpsTileScheduler(has_work, next_tile_id, pid_m, pid_n, self.TILE_M, self.TILE_N, 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
    overlap_bar: gl.shared_memory_descriptor
    NUM_PID_M: gl.tensor
    NUM_PID_N: gl.tensor
    USE_OVERLAP: gl.constexpr

    @gluon.jit
    def get_sps_scheduler(self):
        return SpsTileScheduler.initialize(
            self.a_desc.block_shape[0],
            self.b_desc.block_shape[0],
            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])
    scheduler = p.get_sps_scheduler()
    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()


@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
    tmem_plan: gl.constexpr = tmem_overlap_storage_plan_from_partition(p)
    load_state = Counter.create(0, p.load_empty_bars.shape[0])
    acc_state = Counter.create(1, p.acc_empty_bars.shape[0])
    scheduler = p.get_sps_scheduler()
    i = 0
    while scheduler.has_work:
        # The overlap barrier only protects the shared columns. Before reusing an
        # accumulator window, wait until its epilogue has drained the entire window.
        # For example, Tile 2 cannot reuse the left window until Tile 0 releases it.
        mbarrier.wait(p.acc_empty_bars.index(acc_state.index), acc_state.phase)

        # MMA must wait for the preceding epilogue to drain the overlap.
        if p.USE_OVERLAP:
            mbarrier.wait(p.overlap_bar.index(0), (i + 1) % 2)

        # When overlap is disabled, always use the left (only) accumulator.
        # When overlap is enabled, even tiles use the left and odd tiles use the right accumulator.
        if p.USE_OVERLAP and i % 2 != 0:
            acc_tmem = p.acc_bufs.slice(tmem_plan.right_accum_offset, tmem_plan.right_accum_cols)
        else:
            acc_tmem = p.acc_bufs.slice(tmem_plan.left_accum_offset, tmem_plan.left_accum_cols)

        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_tmem, p.acc_bufs, tmem_plan, 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 += 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]
    tmem_plan: gl.constexpr = tmem_overlap_storage_plan_from_partition(p)
    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)
    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)

        # When overlap is disabled, always use the left (only) accumulator.
        # When overlap is enabled, even tiles use the left and odd tiles use the right accumulator.
        if p.USE_OVERLAP and i % 2 != 0:
            acc_tmem = p.acc_bufs.slice(tmem_plan.right_accum_offset, tmem_plan.right_accum_cols)
        else:
            acc_tmem = p.acc_bufs.slice(tmem_plan.left_accum_offset, tmem_plan.left_accum_cols)

        for s in gl.static_range(subtile_factor):
            # When overlap is disabled, always drain subtiles from low-to-high.
            # When overlap is enabled, drain the overlap first:
            #   Even tiles use the left accumulator and drain high-to-low.
            #   Odd tiles use the right accumulator and drain low-to-high.
            if p.USE_OVERLAP and i % 2 == 0:
                acc_sub = acc_tmem.slice(EPILOGUE_BLOCK_N * (subtile_factor - 1 - s), EPILOGUE_BLOCK_N)
                store_n = off_n + EPILOGUE_BLOCK_N * (subtile_factor - 1 - s)
            else:
                acc_sub = acc_tmem.slice(EPILOGUE_BLOCK_N * s, EPILOGUE_BLOCK_N)
                store_n = off_n + EPILOGUE_BLOCK_N * s

            acc_smem = acc_smems.index(sub_acc_state.index)
            acc = acc_sub.load().to(p.c_desc.dtype)

            # Signal the barrier once the epilogue drains the overlap so the next MMA can proceed.
            if p.USE_OVERLAP and s == tmem_plan.early_release_subtile:
                mbarrier.arrive(p.overlap_bar.index(0), count=1)

            tma.store_wait(pendings=subtile_stages - 1)
            acc_smem.store(acc)
            tma.async_store(p.c_desc, [off_m, store_n], 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 += 1
    tma.store_wait(0)


# ---------------------------------------------------------------------------
# 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,
                                       CGA_LAYOUT: gl.constexpr, USE_OVERLAP: 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)
    num_acc_buffers: gl.constexpr = 2 if USE_OVERLAP else 1

    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)

    # Allocate all of TMEM and then slice it (later) to accommodate accumulators and scales.
    acc_bufs = allocate_tensor_memory(gl.float32, [BLOCK_M, TMEM_COLS], 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=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)
    overlap_bar = mbarrier.allocate_mbarrier(batch=1, two_ctas=TWO_CTAS)
    mbarrier.init(overlap_bar.index(0), count=1)

    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, overlap_bar,
                      gl.cdiv(M, BLOCK_M), gl.cdiv(N, BLOCK_N), USE_OVERLAP)

    gl.warp_specialize([
        (mma_scaled_epilogue_partition, (p, )),
        (mma_scaled_mma_partition, (p, )),
        (mma_scaled_load_partition, (p, )),
    ], [1, 1], [24, 24])


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


def make_descriptors(A, B, A_scale, B_scale, M, N):
    """Create TMA descriptors for the fixed example configuration."""
    a_is_fp4 = A.dtype == torch.uint8
    b_is_fp4 = B.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

    BLOCK_K = 128 if torch.float8_e4m3fn in [A.dtype, B.dtype] else 256
    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_BLOCK_N]

    cga = tuple(tuple(x) for x in CGA_LAYOUT)
    a_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)
    b_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)
    a_desc = TensorDescriptor.from_tensor(A, a_block, a_layout)
    b_desc = TensorDescriptor.from_tensor(B, b_block, b_layout)

    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(c_block, C_dtype, cga_layout=cga)
    c_desc = TensorDescriptor.from_tensor(C, c_block, c_layout)

    is_nvfp4 = A_scale.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)
    a_scale_block = [1, rep_m, rep_k, 2, 256]
    b_scale_block = [1, rep_n, rep_k, 2, 256]

    cga_a_scale = [[0, 1, 0, 0, 0]]
    cga_b_scale = [[0, 0, 0, 0, 0]]
    a_scale_layout = gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5, cga_layout=cga_a_scale)
    b_scale_layout = gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=5, cga_layout=cga_b_scale)

    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)
    a_scale_desc = TensorDescriptor.from_tensor(A_scale_5d, a_scale_block, a_scale_layout)
    b_scale_desc = TensorDescriptor.from_tensor(B_scale_5d, b_scale_block, b_scale_layout)

    return a_desc, b_desc, c_desc, a_scale_desc, b_scale_desc, BLOCK_K


def mma_scaled_warp_specialized(A, B, A_scale, B_scale, use_overlap=True):
    """Warp-specialized block-scale MMA with optional overlapping accumulators."""
    M, N = A.shape[0], B.shape[0]
    A_ELEM_PER_BYTE = 2 if A.dtype == torch.uint8 else 1
    B_ELEM_PER_BYTE = 2 if B.dtype == torch.uint8 else 1
    K = A.shape[1] * A_ELEM_PER_BYTE
    assert K == B.shape[1] * B_ELEM_PER_BYTE
    A_desc, B_desc, C_desc, A_scale_desc, B_scale_desc, BLOCK_K = make_descriptors(A, B, A_scale, B_scale, M, N)

    num_pid = triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N)
    sm_count = torch.cuda.get_device_properties(A.device).multi_processor_count
    grid = (max(1, min(num_pid, sm_count // NUM_CTAS)), )
    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,
        CGA_LAYOUT,
        use_overlap,
        num_ctas=NUM_CTAS,
    )
    return C_desc.base


def make_problem(M, N, K, a_format, b_format):
    A, A_scale, A_ref = random_quantized_tensor(M, K, a_format)
    B, B_scale, B_ref = random_quantized_tensor(N, K, b_format)
    A_scale = swizzle_scales_packed_block(A_scale)
    B_scale = swizzle_scales_packed_block(B_scale)
    return A, B, A_scale, B_scale, A_ref @ B_ref.T


def benchmark(M, N, K, a_format, b_format, use_overlap):
    A, B, A_scale, B_scale, _ = make_problem(M, N, K, a_format, b_format)

    def run():
        return mma_scaled_warp_specialized(A, B, A_scale, B_scale, use_overlap=use_overlap)

    ms = triton.testing.do_bench_cudagraph(run)
    return 2.0 * M * N * K * 1.0e-9 / ms


@pytest.mark.parametrize("a_format, b_format", FORMATS)
@pytest.mark.parametrize("use_overlap", [False, True])
@pytest.mark.skipif(not is_blackwell(), reason="Requires Blackwell")
def test_mma_scaled_overlap_accumulator(a_format, b_format, use_overlap):
    torch.manual_seed(0)
    A, B, A_scale, B_scale, C_ref = make_problem(8192, 8192, 8192, a_format, b_format)
    C = mma_scaled_warp_specialized(A, B, A_scale, B_scale, use_overlap=use_overlap)
    torch.testing.assert_close(C_ref, C.to(torch.float32), atol=2e-3, rtol=1e-3)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--M", type=int, default=8192)
    parser.add_argument("--N", type=int, default=8192)
    parser.add_argument("--K", type=int, nargs="+", default=[512, 1024, 2048, 4096, 8192])
    args = parser.parse_args()

    print("format            K        off         on    speedup")
    for a_format, b_format in FORMATS:
        label = f"{a_format}-{b_format}"
        for K in args.K:
            torch.manual_seed(0)
            off = benchmark(args.M, args.N, K, a_format, b_format, use_overlap=False)
            torch.manual_seed(0)
            on = benchmark(args.M, args.N, K, a_format, b_format, use_overlap=True)
            speedup = (on / off - 1.0) * 100.0
            print(f"{label:14s} {K:7d} {off:9.1f} {on:9.1f} {speedup:8.2f}%")
        print()


if __name__ == "__main__":
    main()
