#!POPCORN benchmark 
#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpu MI300

"""
Matrix Multiplication
=====================
In this tutorial, you will write a very short high-performance FP16 matrix multiplication kernel that achieves
performance on par with cuBLAS or rocBLAS.

You will specifically learn about:

* Block-level matrix multiplications.

* Multi-dimensional pointer arithmetic.

* Program re-ordering for improved L2 cache hit rate.

* Automatic performance tuning.

"""

# %%
# Motivations
# -----------
#
# Matrix multiplications are a key building block of most modern high-performance computing systems.
# They are notoriously hard to optimize, hence their implementation is generally done by
# hardware vendors themselves as part of so-called "kernel libraries" (e.g., cuBLAS).
# Unfortunately, these libraries are often proprietary and cannot be easily customized
# to accommodate the needs of modern deep learning workloads (e.g., fused activation functions).
# In this tutorial, you will learn how to implement efficient matrix multiplications by
# yourself with Triton, in a way that is easy to customize and extend.
#
# Roughly speaking, the kernel that we will write will implement the following blocked
# algorithm to multiply a (M, K) by a (K, N) matrix:
#
#  .. code-block:: python
#
#    # Do in parallel
#    for m in range(0, M, BLOCK_SIZE_M):
#      # Do in parallel
#      for n in range(0, N, BLOCK_SIZE_N):
#        acc = zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=float32)
#        for k in range(0, K, BLOCK_SIZE_K):
#          a = A[m : m+BLOCK_SIZE_M, k : k+BLOCK_SIZE_K]
#          b = B[k : k+BLOCK_SIZE_K, n : n+BLOCK_SIZE_N]
#          acc += dot(a, b)
#        C[m : m+BLOCK_SIZE_M, n : n+BLOCK_SIZE_N] = acc
#
# where each iteration of the doubly-nested for-loop is performed by a dedicated Triton program instance.

# %%
# Compute Kernel
# --------------
#
# The above algorithm is, actually, fairly straightforward to implement in Triton.
# The main difficulty comes from the computation of the memory locations at which blocks
# of :code:`A` and :code:`B` must be read in the inner loop. For that, we need
# multi-dimensional pointer arithmetic.
#
# Pointer Arithmetic
# ~~~~~~~~~~~~~~~~~~~
#
# For a row-major 2D tensor :code:`X`, the memory location of :code:`X[i, j]` is given
# by :code:`&X[i, j] = X + i*stride_xi + j*stride_xj`.
# Therefore, blocks of pointers for :code:`A[m : m+BLOCK_SIZE_M, k:k+BLOCK_SIZE_K]` and
# :code:`B[k : k+BLOCK_SIZE_K, n : n+BLOCK_SIZE_N]` can be defined in pseudo-code as:
#
#  .. code-block:: python
#
#    &A[m : m+BLOCK_SIZE_M, k:k+BLOCK_SIZE_K] =  a_ptr + (m : m+BLOCK_SIZE_M)[:, None]*A.stride(0) + (k : k+BLOCK_SIZE_K)[None, :]*A.stride(1);
#    &B[k : k+BLOCK_SIZE_K, n:n+BLOCK_SIZE_N] =  b_ptr + (k : k+BLOCK_SIZE_K)[:, None]*B.stride(0) + (n : n+BLOCK_SIZE_N)[None, :]*B.stride(1);
#
# Which means that pointers for blocks of A and B can be initialized (i.e., :code:`k=0`) in Triton as the following
# code. Also note that we need an extra modulo to handle the case where :code:`M` is not a multiple of
# :code:`BLOCK_SIZE_M` or :code:`N` is not a multiple of :code:`BLOCK_SIZE_N`, in which case we can pad the data with
# some useless values, which will not contribute to the results. For the :code:`K` dimension, we will handle that later
# using masking load semantics.
#
#  .. code-block:: python
#
#    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
#    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
#    offs_k = tl.arange(0, BLOCK_SIZE_K)
#    a_ptrs = a_ptr + (offs_am[:, None]*stride_am + offs_k [None, :]*stride_ak)
#    b_ptrs = b_ptr + (offs_k [:, None]*stride_bk + offs_bn[None, :]*stride_bn)
#
# And then updated in the inner loop as follows:
#
#  .. code-block:: python
#
#    a_ptrs += BLOCK_SIZE_K * stride_ak;
#    b_ptrs += BLOCK_SIZE_K * stride_bk;
#
#
# L2 Cache Optimizations
# ~~~~~~~~~~~~~~~~~~~~~~
#
# As mentioned above, each program instance computes a :code:`[BLOCK_SIZE_M, BLOCK_SIZE_N]`
# block of :code:`C`.
# It is important to remember that the order in which these blocks are computed does
# matter, since it affects the L2 cache hit rate of our program, and unfortunately, a
# simple row-major ordering
#
#  .. code-block:: Python
#
#    pid = tl.program_id(axis=0)
#    grid_n = tl.cdiv(N, BLOCK_SIZE_N)
#    pid_m = pid // grid_n
#    pid_n = pid % grid_n
#
# is just not going to cut it.
#
# One possible solution is to launch blocks in an order that promotes data reuse.
# This can be done by 'super-grouping' blocks in groups of :code:`GROUP_M` rows before
# switching to the next column:
#
#  .. code-block:: python
#
#    # Program ID
#    pid = tl.program_id(axis=0)
#    # Number of program ids along the M axis
#    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
#    # Number of programs ids along the N axis
#    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
#    # Number of programs in group
#    num_pid_in_group = GROUP_SIZE_M * num_pid_n
#    # Id of the group this program is in
#    group_id = pid // num_pid_in_group
#    # Row-id of the first program in the group
#    first_pid_m = group_id * GROUP_SIZE_M
#    # If `num_pid_m` isn't divisible by `GROUP_SIZE_M`, the last group is smaller
#    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
#    # *Within groups*, programs are ordered in a column-major order
#    # Row-id of the program in the *launch grid*
#    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
#    # Col-id of the program in the *launch grid*
#    pid_n = (pid % num_pid_in_group) // group_size_m
#
# For example, in the following matmul where each matrix is 9 blocks by 9 blocks,
# we can see that if we compute the output in row-major ordering, we need to load 90
# blocks into SRAM to compute the first 9 output blocks, but if we do it in grouped
# ordering, we only need to load 54 blocks.
#
#   .. image:: grouped_vs_row_major_ordering.png
#
# In practice, this can improve the performance of our matrix multiplication kernel by
# more than 10% on some hardware architecture (e.g., 220 to 245 TFLOPS on A100).
#

# %%
# Final Result
# ------------

import torch

import triton
import triton.language as tl
#from kernel_utils import make_match_reference
from utils import make_match_reference

from typing import TypeVar, TypedDict
import time
import math
import numpy as np


input_t = TypeVar("input_t", bound=tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor])
output_t = TypeVar("output_t", bound=torch.Tensor)


DEVICE = triton.runtime.driver.active.get_active_torch_device()
print (f" DEVICE {DEVICE}")


def is_cuda():
    return triton.runtime.driver.active.get_current_target().backend == "cuda"


def is_hip_cdna2():
    target = triton.runtime.driver.active.get_current_target()
    return target.backend == 'hip' and target.arch == 'gfx90a'


def get_cuda_autotune_config():
    return [
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3,
                      num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5,
                      num_warps=2),
        triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5,
                      num_warps=2),
        # Good config for fp8 inputs.
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3,
                      num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3,
                      num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4,
                      num_warps=4)
    ]


def get_hip_autotune_config():
    return [
        triton.Config(
            {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 16, 'GROUP_SIZE_M': 1, 'waves_per_eu': 2},
            num_warps=4, num_stages=2),
        triton.Config(
            {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 16, 'GROUP_SIZE_M': 4, 'waves_per_eu': 2},
            num_warps=8, num_stages=2),
        triton.Config(
            {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 1, 'waves_per_eu': 2},
            num_warps=8, num_stages=2),
        triton.Config(
            {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8, 'waves_per_eu': 3},
            num_warps=4, num_stages=2),
        triton.Config(
            {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 1, 'waves_per_eu': 8},
            num_warps=4, num_stages=2),
    ]


def get_autotune_config():
    if is_cuda():
        return get_cuda_autotune_config()
    else:
        return get_hip_autotune_config()



# `triton.jit`'ed functions can be auto-tuned by using the `triton.autotune` decorator, which consumes:
#   - A list of `triton.Config` objects that define different configurations of
#       meta-parameters (e.g., `BLOCK_SIZE_M`) and compilation options (e.g., `num_warps`) to try
#   - An auto-tuning *key* whose change in values will trigger evaluation of all the
#       provided configs
BLOCK_SHAPE = (128,128)
def get_cuda_autotune_config2():
    tcs =  [
        #triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=4),
        #triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=1, num_warps=16),
        #triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=16),
        #triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, num_warps=16),
        #triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=1, num_warps=16),
        #triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=2, num_warps=16),
        triton.Config({'BLOCK_SIZE_M': 512, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=1, num_warps=16),
        #triton.Config({'BLOCK_SIZE_M': 1024, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=1, num_warps=16),
    ]
    print (f"get_cuda_autotune_config2 tcs { [tc.all_kwargs() for tc in tcs]}")
    return tcs

def get_cuda_autotune_config2test():
    print (f"get_cuda_autotune_config2test")
    return [
        triton.Config({'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 16, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=1,
                      num_warps=4),
    ]

debug_triton=False
def get_autotune_config2(debug=False):
    print (f"get_autotune_config2 debug {debug}")
    if debug:
        return get_cuda_autotune_config2test()
    else:
        return get_cuda_autotune_config2()

@triton.autotune(
    configs=get_autotune_config2(debug_triton),
    key=['M', 'N', 'K'],
)
@triton.jit
def matmul_kernel2(
        # Pointers to matrices
        a_ptr, b_ptr, c_ptr,
        # Matrix dimensions
        M, N, K,
        # The stride variables represent how much to increase the ptr by when moving by 1
        # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr`
        # by to get the element one row down (A has M rows).
        stride_am, stride_ak,  #
        stride_bk, stride_bn,  #
        stride_cm, stride_cn,
        as_ptr, bs_ptr,
        SM, SN, SK,
        stride_asm, stride_ask,  #
        stride_bsk, stride_bsn,  #
        # Meta-parameters
        BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,  #
        GROUP_SIZE_M: tl.constexpr,  #
        ACTIVATION: tl.constexpr  #
):
    """Kernel for computing the matmul C = A x B.
    A has shape (M, K), B has shape (K, N) and C has shape (M, N)
    """
    # -----------------------------------------------------------
    # Map program ids `pid` to the block of C it should compute.
    # This is done in a grouped ordering to promote L2 data reuse.
    # See above `L2 Cache Optimizations` section for details.
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m
    #print (f"########## pid {pid} pid_m {pid_m} pid_n {pid_n}")

    # ----------------------------------------------------------
    # Create pointers for the first blocks of A and B.
    # We will advance this pointer as we move in the K direction
    # and accumulate
    # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers
    # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers
    # See above `Pointer Arithmetic` section for details
    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)


    offs_asm = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bsn = pid_n
    #print (f" offs_am {offs_am} offs_bn {offs_bn } offs_asm {offs_asm}  offs_bsn {offs_bsn}")
    offs_sk = tl.arange(0, 1)
    as_ptrs = as_ptr  + (offs_asm[:, None] * stride_asm + offs_sk[None, :] * stride_ask)
    bs_ptrs = bs_ptr  + (offs_sk[:, None] * stride_bsk + offs_bsn[None, :] * stride_bsn)
    #print (f"as_ptrs {as_ptrs}  bs_ptrs {bs_ptrs}")

    # -----------------------------------------------------------
    # Iterate to compute a block of the C matrix.
    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
    # of fp32 values for higher accuracy.
    # `accumulator` will be converted back to fp16 after the loop.
    accm_type = tl.float32
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accm_type)
    #k=0
    #a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
    #b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
    #asc =  tl.load(as_ptrs,) # mask=offs_sk, other=0.0)
    #bsc =  tl.load(bs_ptrs,) #mask=offs_sk, other=0.0)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        # Load the next block of A and B, generate a mask by checking the K dimension.
        # If it is out of bounds, set it to 0.
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
        asc =  tl.load(as_ptrs,).to(accm_type) # mask=offs_sk, other=0.0)
        bsc =  tl.load(bs_ptrs,).to(accm_type) #mask=offs_sk, other=0.0)
        #print ( "pid, k, ptrs scalings: ", pid, k, as_ptrs, bs_ptrs," scaling", asc, bsc)
        # We accumulate along the K dimension.
        acck = tl.dot(a, b, out_dtype=accm_type)
        #print(f" dtype a {a.dtype} b {b.dtype} acck {acck.dtype}")
        accumulator += acck*asc*bsc
        #print ( "accumulator :", accumulator)
        # Advance the ptrs to the next K block.
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

        as_ptrs += stride_ask
        bs_ptrs += stride_bsk
    # You can fuse arbitrary activation functions here
    # while the accumulator is still in FP32!
    if ACTIVATION == "leaky_relu":
        accumulator = leaky_relu(accumulator)
    c = accumulator.to(tl.bfloat16)

    # -----------------------------------------------------------
    # Write back the block of the output matrix C with masks.
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

def matmul2(a, b, c, a_scale, b_scale, activation=""):
    # Check constraints.
    a = a.contiguous()
    assert a.shape[1] == b.shape[1], "Incompatible dimensions"
    assert a.is_contiguous(), "Matrix A must be contiguous"
    M, K = a.shape
    N, K = b.shape
    SM, SK = a_scale.shape
    SN, SK = b_scale.shape
    b = b.T
    b_scale = b_scale.T
    # Allocates output.
    co = torch.empty((M, N), device=a.device, dtype=c.dtype)
    #print (f" a {a.dtype} {a.shape}  {a}  b {b.dtype} {b.shape} {b}")
    #print (f" a_scale {a_scale}  b_scale {b_scale}")
    # 1D launch kernel where each block gets its own program.
    grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), )
    matmul_kernel2[grid](
        a, b, co,  #
        M, N, K,  #
        a.stride(0), a.stride(1),  #
        b.stride(0), b.stride(1),  #
        c.stride(0), c.stride(1),  #
        a_scale, b_scale,
        SM, SN, SK,  #
        a_scale.stride(0), a_scale.stride(1),  #
        b_scale.stride(0), b_scale.stride(1),  #
        ACTIVATION=activation  #
    )

    c[...] = co

    return c


# %%
# Unit Test
# ---------
#
# We can test our custom matrix multiplication operation against a native torch implementation (i.e., cuBLAS).

TORCH_HAS_FP8 = hasattr(torch, "float8_e5m2")
def ut():
    torch.manual_seed(0)
    a = torch.randn((512, 512), device=DEVICE, dtype=torch.float16)
    b = torch.randn((512, 512), device=DEVICE, dtype=torch.float16)
    triton_output = matmul(a, b)
    torch_output = torch.matmul(a, b)
    print(f"triton_output_with_fp16_inputs={triton_output}")
    print(f"torch_output_with_fp16_inputs={torch_output}")
    # Bigger tolerance for AMD CDNA2 devices.
    # CDNA2 devices use reduced precision fp16 and bf16 and flush input and
    # output denormal values to zero. Detailed info is at: https://pytorch.org/docs/stable/notes/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices
    rtol = 1e-2 if is_hip_cdna2() else 0
    if torch.allclose(triton_output, torch_output, atol=1e-2, rtol=rtol):
        print("✅ Triton and Torch match")
    else:
        print("❌ Triton and Torch differ")

    #TORCH_HAS_FP8 = hasattr(torch, "float8_e5m2")
    if TORCH_HAS_FP8 and is_cuda():
        torch.manual_seed(0)
        a = torch.randn((512, 512), device=DEVICE, dtype=torch.float16)
        b = torch.randn((512, 512), device=DEVICE, dtype=torch.float16)
        a = a.to(torch.float8_e5m2)
        # pre-transpose b for efficiency.
        b = b.T
        b = b.to(torch.float8_e5m2)
        triton_output = matmul(a, b)
        torch_output = torch.matmul(a.to(torch.float16), b.to(torch.float16))
        print(f"triton_output_with_fp8_inputs={triton_output}")
        print(f"torch_output_with_fp8_inputs={torch_output}")
        if torch.allclose(triton_output, torch_output, atol=0.125, rtol=0):
            print("✅ Triton and Torch match")
        else:
            print("❌ Triton and Torch differ")

# %%
# Benchmark
# ---------
#
# Square Matrix Performance
# ~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# We can now compare the performance of our kernel against that of cuBLAS or rocBLAS. Here we focus on square matrices,
# but feel free to arrange this script as you wish to benchmark any other matrix shape.

ref_lib = 'cuBLAS' if is_cuda() else 'rocBLAS'

configs = []
for fp8_inputs in [False, True]:
    if fp8_inputs and (not TORCH_HAS_FP8 or not is_cuda()):
        continue
    configs.append(
        triton.testing.Benchmark(
            x_names=["M", "N", "K"],  # Argument names to use as an x-axis for the plot
            x_vals=[128 * i for i in range(2, 33)],  # Different possible values for `x_name`
            line_arg="provider",  # Argument name whose value corresponds to a different line in the plot
            # Possible values for `line_arg`
            # Don't compare to cublas for fp8 cases as torch.matmul doesn't support fp8 at the moment.
            line_vals=["triton"] if fp8_inputs else [ref_lib.lower(), "triton"],  # Label name for the lines
            line_names=["Triton"] if fp8_inputs else [ref_lib, "Triton"],  # Line styles
            styles=[("green", "-"), ("blue", "-")],
            ylabel="TFLOPS",  # Label name for the y-axis
            plot_name="matmul-performance-" +
            ("fp16" if not fp8_inputs else "fp8"),  # Name for the plot, used also as a file name for saving the plot.
            args={"fp8_inputs": fp8_inputs},
        ))


@triton.testing.perf_report(configs)
def benchmark(M, N, K, provider, fp8_inputs):
    a = torch.randn((M, K), device=DEVICE, dtype=torch.float16)
    b = torch.randn((K, N), device=DEVICE, dtype=torch.float16)
    if TORCH_HAS_FP8 and fp8_inputs:
        a = a.to(torch.float8_e5m2)
        b = b.T
        b = b.to(torch.float8_e5m2)
    quantiles = [0.5, 0.2, 0.8]
    if provider == ref_lib.lower():
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: torch.matmul(a, b), quantiles=quantiles)
    if provider == 'triton':
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: matmul(a, b), quantiles=quantiles)
    perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
    return perf(ms), perf(max_ms), perf(min_ms)


#benchmark.run(show_plots=False, print_data=True)

def local_generate_input(m: int, n: int, k: int, seed: int, block_shape=BLOCK_SHAPE, dtype=torch.float8_e5m2):
    """
    Generate random input and weights for Blockwise W8A8 Matmul scaled to FP32.

    Returns:
        Tuple of (
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k],
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k],
            a_scale: torch.Tensor[float32] of shape [m, k // 128],
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128],
            c: torch.Tensor[bfloat16] of shape [m, n]
        )
    """
    gen = torch.Generator(device='cuda')
    gen.manual_seed(seed)
    block_shape_n, block_shape_k = block_shape
    scale_n =  (n + block_shape_n - 1) // block_shape_n
    scale_k =  (k + block_shape_k - 1) // block_shape_k

    # Generate random inputs with FP8 quantization
    a = (torch.randn((k, m), dtype=torch.bfloat16, device="cuda:0", generator=gen)).to(dtype)
    b = (torch.randn((k, n), dtype=torch.bfloat16, device="cuda:0", generator=gen)).to(dtype)

    # Generate scaling factors with FP32
    a_scale = torch.randn([scale_k, m], dtype=torch.float32, device="cuda", generator=gen)
    b_scale = torch.randn([scale_k, scale_n], dtype=torch.float32, device="cuda", generator=gen)

    c = torch.zeros((m, n), dtype=torch.bfloat16, device="cuda")
    #print (f"generate a {a.T}  b.T {b} ")
    #print (f"geberate pre-exapansion   a_scale {a_scale.T}  b_scale.T {b_scale}")
    return (a.T, b.T, a_scale.T, b_scale.T, c)

def local_ref_kernel(data: input_t, block_shape=BLOCK_SHAPE) -> output_t:
    """
    Highly inefficient torch reference implementation of FP8 GEMM.
    You can use this as a reference / starting template for your implementation.
    """
    # c: [m, n] is pre-allocated memory to help remove allocation overhead.
    a, b, a_scale, b_scale, c = data

    # a is M x K in column-major order, we convert here for simplicity.
    a = a.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()

    # constants
    m = a.shape[0]
    n = b.shape[0]
    k = a.shape[1]
    block_shape_n = block_shape[1]
    block_shape_k = block_shape[0]
    scale_n = b_scale.shape[0]
    scale_k = b_scale.shape[1]


    # Apply blockwise scaling to input 'a'
    a_scale = a_scale.unsqueeze(-1).repeat(1, 1, block_shape_k)  # Shape: [m, scale_k, block_shape_k]
    a_scale = a_scale.reshape(m, scale_k * block_shape_k)
    a_scale = a_scale[:, :k]

    #print (f"ref_kernel  a {a}  b.T {b.T} ")

    # Dequantize 'a', in your implementation you should do this at the end.
    a = a.to(a_scale.dtype) * a_scale

    # Apply blockwise scaling to input 'b'
    b_scale = (
        b_scale.view(-1, 1)
        .repeat(1, block_shape_n * block_shape_k)
        .view(scale_n, scale_k, block_shape_n, block_shape_k)
        .permute(0, 2, 1, 3)  # Reorder dimensions: [scale_n, blk_n, scale_k, blk_k]
        .reshape(scale_n * block_shape_n, scale_k * block_shape_k)
    )
    b_scale = b_scale[:n, :k]
    #print (f"ref_kernel post_expansion  a_scale {a_scale}  b_scale.T {b_scale.T}")

    # Dequantize 'b', in your implementation you should do this at the end.
    b = b.to(b_scale.dtype) * b_scale

    # Compute FP8 GEMM and write to 'c'.
    c[...] = (a @ b.T).to(torch.bfloat16)
    return c


def calculate_stats(numbers, mult=1.0e-3):
  """Calculates the min, max, average, and standard deviation of a list.

  Args:
    numbers: A list of numbers.

  Returns:
    A tuple containing the min, max, average, and standard deviation.
  """
  if not numbers:
    return None, None, None, None

  min_val = np.min(numbers)*mult
  max_val = np.max(numbers)*mult
  avg_val = np.mean(numbers)*mult
  std_dev_val = np.std(numbers)*mult

  return min_val, max_val, avg_val, std_dev_val


def localbm() :
    #ut()
    #dtype = torch.float8_e4m3fnuz
    #dtype = torch.float8_e5m2fnuz
    #dtype = torch.bfloat16
    #dtype = torch.float16
    #dtype = torch.float32

    FP8TOPS = 2.614e15
    dtype=torch.bfloat16
    if 'cu' in torch.__version__:
        dtype=torch.float8_e5m2
        FP8TOPS = 3.341e15
    elif 'rocm' in torch.__version__:
        dtype=torch.float8_e4m3fnuz
        FP8TOPS = 2.614e15

    print( f" dtype : {torch.finfo(dtype)} FP8TOPS {FP8TOPS}")
    # debug_triton = True
    #block_shape = (16,16)
    #data = generate_input(4, 4, 4, 6789, dtype=dtype)
    #data = generate_input(16, 16, 16, 6789, dtype=dtype)
    SOLMNKS = [[1024,    7168,     256,      3.17],
            [1024,    1536,    7168,      8.63],
            [1024,    4608,    7168,     25.89],
            [6144,    7168,     256,     17.27],
            [6144,    1536,    7168,     51.78],
            [6144,    4608,    7168,    155.30]]

    TESTMNKS = [[64, 64, 128, 6635],
            [64, 1536, 7168, 6635],
            [64, 3072, 1536, 1236],
            [64, 576, 7168, 542],
            [96, 7168, 256, 1234],
            [96, 7168, 2048, 4153],
            [96, 4608, 7168, 412],
            [128, 7168, 2304, 624],
            [128, 512, 7168, 2514],
            [512, 4096, 512, 543],
            [512, 1536, 7168, 12341]]

    BMMNKS = [[1024,  1536, 7168, 8135],
            [1024,  3072, 1536, 6251],
            [1024,  576, 7168, 12346],
            [1024,  7168, 256, 5364],
            [1024,  7168, 2048, 6132],
            [1024,  4608, 7168, 7531],
            [1024,  7168, 2304, 12345],
            [1024,  512, 7168, 6563],
            [1024,  4096, 512, 17512],
            [6144,  1536, 7168, 6543],
            [6144,  3072, 1536, 234],
            [6144,  576, 7168, 9863],
            [6144,  7168, 256, 764243],
            [6144,  7168, 2048, 76547],
            [6144,  4608, 7168, 65436],
            [6144,  7168, 2304, 452345],
            [6144,  512, 7168, 12341],
            [6144,  4096, 512, 45245]]


    PROFMNKS = [
            [1024,  1536, 7168, 8135],
            [6144,  4608, 7168, 65436],
            ]

    MNKS = BMMNKS
    cumprod = 1.0
    sol = 1.0
    for MNK in MNKS:
        M=MNK[0]
        N=MNK[1]
        K=MNK[2]
        seed= MNK[3]
        #SoL=MNK[3]
        runtimes = []
        for i in range(30):
            data = local_generate_input(M,N,K, seed, dtype=dtype)
            a, b, a_scale, b_scale, c = data
            time0 = time.time_ns()
            cout =  matmul2 (a,b,c, a_scale, b_scale).clone().detach()
            runtimes.append(time.time_ns() -time0)
            check_implementation = make_match_reference(local_ref_kernel, rtol=2e-02, atol=1e-03)
            result, notes = check_implementation (data, cout)
            #print ( f"completed {MNK} with {result} ", notes)
        minimum, maximum, average, standard_deviation = calculate_stats(runtimes[10:])
        solmnk = (2*M*N*K/FP8TOPS)/1.0e-6
        print ( f" MNK {M} {N} {K} avg {average} min {minimum} max {maximum} std {standard_deviation} SoL {solmnk} ")
        cumprod *= average
        sol *= solmnk
    geom = math.pow(cumprod, 1.0/len(MNKS))
    geomsol = math.pow(sol, 1.0/len(MNKS))
    print ( f" geom: {geom} sol {geomsol} ")
    print ( "localbm completed")

def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of block-scale fp8 gemm 
    Args:
        data: Tuple that expands to:
            a: torch.Tensor[float8_e4m3fnuz] of shape [m, k], 
            b: torch.Tensor[float8_e4m3fnuz] of shape [n, k], 
            a_scale: torch.Tensor[float32] of shape [m, k // 128], 
            b_scale: torch.Tensor[float32] of shape [n // 128, k // 128], 
            c: torch.Tensor[bfloat16] of shape [m, n]
    Returns:
        Tensor containing output in bf16
    """
    localbm()
    a, b, a_scale, b_scale, c = data
    print ( f"types shapes:  a {a.dtype} {a.shape}  b {b.dtype} {b.shape}  c {c.dtype} {c.shape} a_scale {a_scale.dtype} {a_scale.shape} b_scale {b_scale.dtype} {b_scale.shape} ")
    #benchmark.run(print_data=True, show_plots=False)
    #unit_test()
    matmul2(a,b,c, a_scale, b_scale)

    return c

if False: #__name__ == "__main__":
    localbm()


python eval.py test /tmp/tmp8tqi33h3
 DEVICE cuda:0
get_autotune_config2 debug False
get_cuda_autotune_config2 tcs [{'BLOCK_SIZE_M': 512, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8, 'num_warps': 16, 'num_ctas': 1, 'num_stages': 1}]
 dtype : finfo(resolution=1, min=-240, max=240, eps=0.125, smallest_normal=0.0078125, tiny=0.0078125, dtype=float8_e4m3fnuz) FP8TOPS 2614000000000000.0
 MNK 1024 1536 7168 avg 107.42755000000001 min 102.749 max 113.86 std 2.8813932823375565 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 106.81564999999999 min 100.578 max 112.64 std 3.215194057518146 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 106.35965 min 100.72800000000001 max 112.715 std 3.0190098256713243 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 104.5931 min 100.764 max 111.123 std 3.0554632038366947 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 105.01245 min 100.027 max 110.19800000000001 std 2.719429507727678 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 107.99325 min 103.762 max 113.504 std 2.503555449256117 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 105.7001 min 102.107 max 111.899 std 2.5492305290028203 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 104.86775 min 98.977 max 109.377 std 2.4854708583083407 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 105.17830000000001 min 100.052 max 114.288 std 3.339348635587485 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 107.45360000000001 min 103.238 max 113.878 std 2.8957607877723603 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 107.07835000000001 min 101.298 max 112.426 std 3.291120041490435 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 105.96175000000001 min 100.833 max 116.21900000000001 std 3.56203382739412 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 104.64930000000001 min 99.474 max 110.827 std 3.404511522964785 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 106.8477 min 102.02 max 112.36500000000001 std 2.6797140910925554 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 108.09514999999999 min 101.941 max 118.864 std 3.4367050248021 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 106.0356 min 101.712 max 111.426 std 2.5051912382091714 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 105.7265 min 99.933 max 111.179 std 2.911397576766183 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 105.71114999999999 min 100.968 max 116.06 std 3.4904106961072645 SoL 9.858379409334354 
 geom: 106.18921273777325 sol 12.695267751384295 
localbm completed
types shapes:  a torch.float8_e4m3fnuz torch.Size([64, 128])  b torch.float8_e4m3fnuz torch.Size([64, 128])  c torch.bfloat16 torch.Size([64, 64]) a_scale torch.float32 torch.Size([64, 1]) b_scale torch.float32 torch.Size([1, 1]) 
 dtype : finfo(resolution=1, min=-240, max=240, eps=0.125, smallest_normal=0.0078125, tiny=0.0078125, dtype=float8_e4m3fnuz) FP8TOPS 2614000000000000.0
 MNK 1024 1536 7168 avg 108.4441 min 102.47500000000001 max 117.822 std 3.737591964085967 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 106.3224 min 101.412 max 113.158 std 3.2406404212747826 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 105.4132 min 101.46000000000001 max 110.956 std 2.5726286284654454 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 104.59345 min 100.056 max 109.935 std 2.8360054209221817 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 104.3922 min 99.855 max 109.649 std 2.37425831787529 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 106.28805 min 101.913 max 112.913 std 2.4423411202164207 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 105.37205 min 99.705 max 111.07000000000001 std 3.3252956782066763 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 105.8749 min 101.247 max 114.307 std 3.148195513305996 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 103.9127 min 100.304 max 110.102 std 2.4572312080876717 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 106.6887 min 103.23100000000001 max 109.875 std 2.091158126493547 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 106.25965 min 103.553 max 115.936 std 2.805978711163005 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 105.94565 min 100.763 max 110.287 std 2.2387394952294026 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 106.94725 min 102.535 max 112.251 std 2.9868426619927604 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 105.94875 min 102.346 max 110.026 std 2.3886402381899208 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 106.6194 min 101.933 max 113.58500000000001 std 3.159767814254712 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 108.379 min 103.852 max 113.199 std 3.031183976600563 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 104.75685000000001 min 100.33800000000001 max 110.836 std 3.0643984609544495 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 104.20775 min 100.892 max 110.571 std 2.701623972261869 SoL 9.858379409334354 
 geom: 105.90191936470558 sol 12.695267751384295 
localbm completed
types shapes:  a torch.float8_e4m3fnuz torch.Size([64, 7168])  b torch.float8_e4m3fnuz torch.Size([1536, 7168])  c torch.bfloat16 torch.Size([64, 1536]) a_scale torch.float32 torch.Size([64, 56]) b_scale torch.float32 torch.Size([12, 56]) 
 dtype : finfo(resolution=1, min=-240, max=240, eps=0.125, smallest_normal=0.0078125, tiny=0.0078125, dtype=float8_e4m3fnuz) FP8TOPS 2614000000000000.0
 MNK 1024 1536 7168 avg 103.1517 min 98.548 max 113.31 std 3.487365081261209 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 101.69669999999999 min 97.149 max 109.197 std 2.89715707375351 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 103.81089999999999 min 99.307 max 110.19500000000001 std 2.9563985844266667 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 101.92515 min 97.392 max 110.62100000000001 std 3.2715790113491066 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 103.67525 min 99.008 max 110.884 std 3.4804829531977313 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 103.7042 min 99.305 max 109.732 std 2.5217687165955565 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 103.0052 min 97.937 max 108.926 std 2.7417860711587254 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 102.46180000000001 min 97.771 max 106.977 std 2.466618446375523 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 101.68610000000001 min 98.601 max 109.718 std 2.8388565814426063 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 106.94775 min 102.265 max 116.525 std 3.1344779928243236 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 105.63 min 100.08 max 114.455 std 3.7700699860877913 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 103.72695 min 100.921 max 107.872 std 1.825378521704471 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 102.90995 min 98.011 max 110.447 std 2.9600692639700172 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 103.85185000000001 min 99.453 max 111.32000000000001 std 2.854486088160179 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 105.25614999999999 min 99.202 max 111.298 std 3.2580850399429417 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 104.43485000000001 min 100.345 max 112.18 std 3.4878759048309043 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 104.36845 min 100.702 max 113.497 std 3.247230550407531 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 101.69080000000001 min 97.19500000000001 max 108.45 std 2.6650422810904892 SoL 9.858379409334354 
 geom: 103.54239475994218 sol 12.695267751384295 
localbm completed
types shapes:  a torch.float8_e4m3fnuz torch.Size([64, 1536])  b torch.float8_e4m3fnuz torch.Size([3072, 1536])  c torch.bfloat16 torch.Size([64, 3072]) a_scale torch.float32 torch.Size([64, 12]) b_scale torch.float32 torch.Size([24, 12]) 
 dtype : finfo(resolution=1, min=-240, max=240, eps=0.125, smallest_normal=0.0078125, tiny=0.0078125, dtype=float8_e4m3fnuz) FP8TOPS 2614000000000000.0
 MNK 1024 1536 7168 avg 103.7086 min 98.97200000000001 max 107.805 std 2.3771438618644853 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 103.92645 min 99.227 max 111.781 std 3.4326966436753477 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 101.96610000000001 min 97.44 max 105.561 std 2.270791313617348 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 102.10795 min 98.849 max 106.688 std 2.2101746871005465 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 102.97305 min 99.052 max 107.662 std 2.5678769533410284 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 106.3909 min 100.624 max 115.473 std 3.7303041283520035 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 105.71315 min 100.333 max 111.95400000000001 std 3.4369808447967816 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 103.95925 min 100.287 max 108.084 std 2.36311825508162 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 104.45815 min 100.239 max 109.977 std 2.406432385815151 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 104.61875 min 101.399 max 109.34 std 1.9343154570803596 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 105.42215 min 100.91 max 115.815 std 3.3276243519213526 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 106.12195 min 100.032 max 110.973 std 3.0035096050287575 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 104.081 min 100.85000000000001 max 109.98 std 2.396157382143335 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 105.8288 min 100.16 max 110.854 std 2.498407164575062 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 108.4285 min 103.352 max 114.312 std 3.1062601066233975 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 108.27095 min 103.297 max 114.928 std 3.253332329704422 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 104.1815 min 100.637 max 109.792 std 2.2557247504959466 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 103.1267 min 98.977 max 109.95100000000001 std 2.692840305328186 SoL 9.858379409334354 
 geom: 104.72300657688014 sol 12.695267751384295 
localbm completed
types shapes:  a torch.float8_e4m3fnuz torch.Size([64, 7168])  b torch.float8_e4m3fnuz torch.Size([576, 7168])  c torch.bfloat16 torch.Size([64, 576]) a_scale torch.float32 torch.Size([64, 56]) b_scale torch.float32 torch.Size([5, 56]) 
 dtype : finfo(resolution=1, min=-240, max=240, eps=0.125, smallest_normal=0.0078125, tiny=0.0078125, dtype=float8_e4m3fnuz) FP8TOPS 2614000000000000.0
 MNK 1024 1536 7168 avg 103.3099 min 99.22500000000001 max 106.519 std 1.669663076791243 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 103.64375 min 99.145 max 110.399 std 3.1049583229892153 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 103.3959 min 98.815 max 115.104 std 3.379083883244096 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 102.75165 min 98.86500000000001 max 108.525 std 2.5264183991374036 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 104.199 min 98.973 max 110.441 std 2.8951298934590137 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 105.11975 min 100.169 max 108.73100000000001 std 2.2412202898198115 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 102.1289 min 98.456 max 106.108 std 2.242757095630287 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 103.01050000000001 min 98.767 max 109.16 std 2.621998693744908 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 103.12145 min 98.36 max 117.691 std 4.334919543371019 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 105.05680000000001 min 100.071 max 112.24000000000001 std 3.3383321075051837 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 105.42485 min 103.08200000000001 max 110.747 std 2.3413308240186823 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 105.35485000000001 min 102.635 max 109.188 std 1.931200281560667 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 104.06055 min 99.885 max 111.34 std 2.868687739629394 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 106.16310000000001 min 101.827 max 110.539 std 2.3649575662155127 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 105.92805 min 101.077 max 113.774 std 3.114628027148667 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 105.55885 min 102.161 max 110.544 std 2.456781945452221 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 103.04710000000001 min 100.138 max 109.538 std 2.4365748274986347 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 103.82600000000001 min 98.502 max 107.768 std 2.4205316564754944 SoL 9.858379409334354 
 geom: 104.16553847983477 sol 12.695267751384295 
localbm completed
types shapes:  a torch.float8_e4m3fnuz torch.Size([96, 256])  b torch.float8_e4m3fnuz torch.Size([7168, 256])  c torch.bfloat16 torch.Size([96, 7168]) a_scale torch.float32 torch.Size([96, 2]) b_scale torch.float32 torch.Size([56, 2]) 
 dtype : finfo(resolution=1, min=-240, max=240, eps=0.125, smallest_normal=0.0078125, tiny=0.0078125, dtype=float8_e4m3fnuz) FP8TOPS 2614000000000000.0
 MNK 1024 1536 7168 avg 103.1307 min 99.379 max 113.687 std 3.3980081091721956 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 102.065 min 97.471 max 107.8 std 2.995762857770955 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 101.54975 min 97.783 max 106.521 std 2.4965796777791813 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 102.1674 min 98.02 max 109.822 std 3.5189706932567657 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 102.74365 min 96.999 max 110.397 std 3.6076361135097867 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 103.68889999999999 min 99.661 max 109.894 std 2.6773520481998627 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 101.62780000000001 min 96.71000000000001 max 106.98400000000001 std 2.7272371477376147 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 102.42165 min 99.403 max 108.61 std 2.1912207391086826 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 101.7464 min 98.019 max 109.281 std 2.664160325505956 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 106.00485 min 101.82000000000001 max 112.366 std 3.017102190430414 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 102.64995 min 99.119 max 108.928 std 2.593026021369628 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 102.65085 min 99.505 max 112.646 std 2.764577314437055 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 103.0212 min 97.098 max 110.001 std 2.983360229673916 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 105.9662 min 102.17 max 111.699 std 2.4628149463571147 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 107.15644999999999 min 102.257 max 115.519 std 3.502064640679838 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 104.1969 min 100.71000000000001 max 109.184 std 2.1249374084899535 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 101.5954 min 96.139 max 117.852 std 4.302366934607042 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 103.18145 min 99.297 max 108.693 std 2.33668235913656 SoL 9.858379409334354 
 geom: 103.18576760395018 sol 12.695267751384295 
localbm completed
types shapes:  a torch.float8_e4m3fnuz torch.Size([96, 2048])  b torch.float8_e4m3fnuz torch.Size([7168, 2048])  c torch.bfloat16 torch.Size([96, 7168]) a_scale torch.float32 torch.Size([96, 16]) b_scale torch.float32 torch.Size([56, 16]) 
 dtype : finfo(resolution=1, min=-240, max=240, eps=0.125, smallest_normal=0.0078125, tiny=0.0078125, dtype=float8_e4m3fnuz) FP8TOPS 2614000000000000.0
 MNK 1024 1536 7168 avg 101.75735 min 98.297 max 106.433 std 2.1651893976047454 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 102.1443 min 97.036 max 112.803 std 3.7492811190946993 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 102.91585 min 97.765 max 113.793 std 3.5397244281864655 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 103.24955 min 99.027 max 112.525 std 3.4952661769170024 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 102.30735000000001 min 97.891 max 107.145 std 2.686552889391906 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 106.0099 min 100.403 max 111.592 std 2.9849047539243188 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 103.5455 min 99.71900000000001 max 107.921 std 2.4136086364611806 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 102.3215 min 98.659 max 106.595 std 2.1664723284639478 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 102.57045 min 99.008 max 109.18900000000001 std 2.742283673783586 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 105.3503 min 100.96900000000001 max 112.84 std 3.161308417412006 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 105.04344999999999 min 101.422 max 117.98100000000001 std 3.7495648210825747 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 102.19680000000001 min 98.437 max 109.096 std 2.540552746943074 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 102.75619999999999 min 98.717 max 106.733 std 2.3496926735213695 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 105.60445 min 99.111 max 111.645 std 3.246106706117345 SoL 69.00865586534047 
[...] 96 lines omitted