#!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.
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    #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,) # mask=offs_sk, other=0.0)
        bsc =  tl.load(bs_ptrs,) #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, )
        #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/tmpvnz6qapk
 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 132.5203 min 125.643 max 140.828 std 4.784223104120459 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 130.34845 min 125.268 max 137.725 std 3.551300585348979 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 127.15610000000001 min 120.36800000000001 max 134.287 std 3.028301139252832 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 123.86030000000001 min 118.75 max 129.962 std 2.8005930818310607 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 131.8028 min 123.258 max 144.481 std 4.7922821765835115 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 131.08235000000002 min 125.72200000000001 max 139.054 std 3.5657779834841095 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 130.6317 min 125.854 max 141.092 std 3.792045254213088 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 128.90664999999998 min 123.645 max 140.037 std 3.68309613335031 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 127.81689999999999 min 123.748 max 133.885 std 2.807178617402178 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 130.89095 min 125.18 max 141.559 std 4.269509251366015 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 132.478 min 123.58200000000001 max 139.922 std 3.7740752244755273 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 130.349 min 125.342 max 136.16400000000002 std 3.0631586312171297 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 130.1851 min 124.768 max 138.23 std 3.677306363358919 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 130.4935 min 122.655 max 138.785 std 4.072135502411481 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 133.9091 min 127.376 max 175.65 std 10.132952639285353 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 131.26270000000002 min 128.13 max 136.171 std 2.2332938476608937 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 129.44575 min 125.08500000000001 max 139.346 std 3.8214967470220356 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 129.85855 min 124.581 max 137.822 std 3.9362778925655135 SoL 9.858379409334354 
 geom: 130.1477450410998 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 133.88510000000002 min 127.316 max 143.409 std 3.891020558927953 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 130.2151 min 123.824 max 146.26500000000001 std 4.995022111062172 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 129.1073 min 124.807 max 134.797 std 3.0573394332327575 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 129.13675 min 121.45100000000001 max 136.912 std 4.00716404549402 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 128.57195 min 124.438 max 135.343 std 2.6621518640941577 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 132.14675 min 127.161 max 140.769 std 4.268609186550111 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 132.23035000000002 min 126.656 max 138.77100000000002 std 3.294660290151323 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 129.82955 min 123.57900000000001 max 140.066 std 3.809603187669288 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 128.5918 min 122.914 max 140.525 std 3.9177331404780493 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 131.03215 min 125.757 max 137.235 std 2.6455833057191755 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 130.55275 min 125.697 max 140.856 std 3.6566579259619023 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 128.50705 min 123.026 max 134.614 std 2.763598857920592 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 129.32825 min 125.139 max 138.228 std 3.2934783721014473 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 133.14125 min 128.379 max 140.165 std 3.3142926677497866 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 130.96785 min 125.239 max 137.823 std 3.791253015495009 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 133.65005 min 127.941 max 140.005 std 3.2021184468254766 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 128.3831 min 124.10000000000001 max 135.247 std 2.7432234487915856 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 128.5507 min 121.586 max 136.481 std 3.7673603769748394 SoL 9.858379409334354 
 geom: 130.4222092131531 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 127.37460000000002 min 123.527 max 138.404 std 3.861943233658413 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 128.08075 min 122.913 max 136.136 std 3.6996718486238747 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 128.9238 min 122.166 max 141.482 std 4.484431709815637 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 128.0843 min 122.284 max 135.18800000000002 std 3.286433083146529 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 128.34810000000002 min 123.325 max 138.413 std 3.7603876515593444 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 127.6497 min 123.456 max 133.251 std 2.665316568439854 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 127.94064999999999 min 121.308 max 133.56900000000002 std 3.225196726945505 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 128.19045 min 122.12700000000001 max 135.84 std 3.435615453961634 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 126.15195 min 119.342 max 132.675 std 3.910378555523749 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 130.09810000000002 min 122.896 max 134.607 std 3.51254996975132 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 129.8804 min 123.693 max 140.21 std 4.119365125841602 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 129.64985000000001 min 123.593 max 140.303 std 4.329185480838169 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 127.13289999999999 min 122.911 max 133.096 std 2.997254091664569 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 131.43645 min 127.199 max 137.339 std 2.9993739259218746 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 130.6659 min 126.258 max 137.872 std 3.107572217986253 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 130.24935000000002 min 124.735 max 139.732 std 3.258934446640497 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 128.32345 min 123.143 max 137.439 std 3.7125330230854514 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 127.66915 min 122.32600000000001 max 141.404 std 4.127156930321406 SoL 9.858379409334354 
 geom: 128.6512686008295 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 128.891 min 123.305 max 135.05 std 3.148789799272095 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 128.42525 min 122.87100000000001 max 134.701 std 2.9555868770009113 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 126.2716 min 121.159 max 131.821 std 2.7480845219898167 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 126.9212 min 121.467 max 141.465 std 4.408979559943549 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 128.9017 min 121.869 max 177.906 std 11.594371893724992 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 131.0975 min 125.031 max 140.014 std 4.297090777491208 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 129.58335 min 124.215 max 140.548 std 3.9606044522900796 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 130.41505 min 126.34 max 137.383 std 3.4425575881167187 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 128.6139 min 122.894 max 141.768 std 4.280028491727597 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 130.06130000000002 min 124.46300000000001 max 138.823 std 3.5550289605008847 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 130.61939999999998 min 125.83200000000001 max 138.105 std 3.224952284918337 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 129.58505 min 124.655 max 136.513 std 3.5948058984457005 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 127.44515 min 121.99300000000001 max 135.756 std 3.3633936028214126 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 132.64170000000001 min 127.977 max 138.854 std 3.7682962078371705 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 130.77225 min 125.01 max 138.96 std 3.826841058562532 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 132.8631 min 127.848 max 139.674 std 3.885645595007347 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 129.54205000000002 min 125.84100000000001 max 137.233 std 2.585972553508641 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 129.3332 min 121.97800000000001 max 135.653 std 3.695705529394895 SoL 9.858379409334354 
 geom: 129.54357123434107 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 128.1489 min 123.511 max 134.306 std 3.0243423566124257 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 129.37064999999998 min 124.827 max 137.937 std 3.0839898228593423 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 129.72135 min 121.748 max 137.993 std 3.9487458930019796 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 128.4343 min 121.884 max 140.216 std 4.882862337809658 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 127.9897 min 122.732 max 136.298 std 3.83092485047671 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 130.6417 min 124.97500000000001 max 135.582 std 2.894073584067966 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 128.16595 min 122.45 max 140.84 std 3.9154308125032675 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 128.79055 min 123.07900000000001 max 137.573 std 4.100042188502455 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 125.94239999999999 min 122.026 max 130.404 std 2.0444860332122596 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 127.83655 min 123.82300000000001 max 134.808 std 2.8800132026607104 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 129.4089 min 124.78 max 137.496 std 3.445619754122616 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 129.4108 min 125.113 max 134.667 std 2.7162909012106935 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 129.08405000000002 min 121.803 max 141.117 std 4.38905582642782 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 130.07925 min 125.068 max 142.89000000000001 std 4.300611489951168 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 131.3101 min 123.31400000000001 max 140.114 std 4.145465931593215 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 133.2346 min 124.745 max 145.67600000000002 std 4.551583827636266 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 128.13405 min 123.345 max 135.216 std 3.0012938289177886 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 129.54315 min 123.229 max 138.924 std 3.746739492879109 SoL 9.858379409334354 
 geom: 129.1713747500342 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 127.0467 min 122.35300000000001 max 136.64000000000001 std 3.238653394545332 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 126.668 min 120.248 max 134.222 std 4.051641210670066 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 129.989 min 124.32600000000001 max 138.581 std 3.7852144060805855 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 126.99860000000001 min 122.083 max 132.805 std 2.837479804333416 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 127.269 min 122.769 max 136.44 std 3.654592945869622 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 132.2392 min 128.02100000000002 max 139.20000000000002 std 3.1822371784642325 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 128.44825 min 123.46600000000001 max 134.477 std 3.2255788763414235 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 127.75585000000001 min 122.009 max 131.98 std 2.480457725400697 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 124.38075 min 120.709 max 128.80700000000002 std 2.309443328488491 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 131.78135 min 123.367 max 139.733 std 4.27839041316942 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 130.4979 min 124.827 max 135.591 std 3.108373560240146 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 130.40800000000002 min 125.727 max 137.725 std 3.912059917741547 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 129.89315 min 124.422 max 136.389 std 3.485615286789407 SoL 8.626081983167559 
 MNK 6144 7168 2048 avg 131.16085 min 122.88 max 144.54500000000002 std 4.739491663406531 SoL 69.00865586534047 
 MNK 6144 4608 7168 avg 132.75820000000002 min 127.809 max 138.975 std 3.4242629659534036 SoL 155.26947569701608 
 MNK 6144 7168 2304 avg 130.40460000000002 min 124.123 max 139.381 std 4.370876335473243 SoL 77.63473784850804 
 MNK 6144 512 7168 avg 128.19050000000001 min 122.23100000000001 max 133.451 std 2.850810156779999 SoL 17.252163966335118 
 MNK 6144 4096 512 avg 131.208 min 125.988 max 137.52100000000002 std 3.375392317938761 SoL 9.858379409334354 
 geom: 129.26411265446822 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 128.04515 min 122.158 max 136.224 std 4.14799563976386 SoL 8.626081983167559 
 MNK 1024 3072 1536 avg 127.81385 min 123.038 max 136.377 std 4.1305695645394955 SoL 3.696892278500383 
 MNK 1024 576 7168 avg 128.9174 min 123.53 max 136.873 std 3.871113023408126 SoL 3.234780743687835 
 MNK 1024 7168 256 avg 126.8797 min 121.325 max 131.78300000000002 std 2.9718416529149057 SoL 1.4376803305279267 
 MNK 1024 7168 2048 avg 126.59230000000001 min 120.26 max 131.80100000000002 std 3.183375097282757 SoL 11.501442644223413 
 MNK 1024 4608 7168 avg 133.44485 min 127.04 max 140.56900000000002 std 4.0832677756302 SoL 25.87824594950268 
 MNK 1024 7168 2304 avg 127.4178 min 121.45100000000001 max 133.475 std 3.3183743248765656 SoL 12.93912297475134 
 MNK 1024 512 7168 avg 128.4973 min 121.19 max 134.857 std 2.990037743908929 SoL 2.8753606610558533 
 MNK 1024 4096 512 avg 126.98935 min 120.884 max 135.14000000000001 std 3.605665254498814 SoL 1.643063234889059 
 MNK 6144 1536 7168 avg 134.10595 min 126.938 max 171.223 std 9.439774999834476 SoL 51.75649189900536 
 MNK 6144 3072 1536 avg 131.17555 min 124.422 max 144.653 std 4.2633404095263145 SoL 22.181353671002295 
 MNK 6144 576 7168 avg 128.88845 min 124.15 max 137.949 std 3.5592589042524008 SoL 19.40868446212701 
 MNK 6144 7168 256 avg 130.4817 min 123.467 max 139.053 std 3.8962819854317536 SoL 8.626081983167559 
[...] 97 lines omitted