#!POPCORN leaderboard amd-fp8-mm
import triton
import triton.language as tl
from task import input_t, output_t
import sys

_printed_config_summaries = {}

@triton.autotune(
    configs=[
        # M=128, K=64
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=8),

        # M=128, K=128
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=8),

        # M=256, K=64
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=8),

        # M=256, K=128
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=8),

        # --- Add K=256 Probes ---
        # Pair K=256 with promising M=128 and common winner W=8, S=1/2 and W=4, S=3
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4}, num_stages=2, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4}, num_stages=3, num_warps=4),
        # Pair K=256 with promising M=256 and common winner W=8, S=1
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 256, 'GROUP_SIZE_M': 4}, num_stages=1, num_warps=8),

    ],
    key=['M', 'N', 'K'], # Autotuning key based on matrix dimensions
)
@triton.jit
def matmul_fp8_kernel(
    a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, c_ptr, # Matrix dimensions
    M, N, K,
    stride_am, stride_ak,   # A (M x K): stride_am=K, stride_ak=1
    stride_bn, stride_bk,   # B (N x K): stride_bn=K, stride_bk=1
    stride_a_scale_m, stride_a_scale_k, # A scales (M x K/BLK_K_S): stride_a_scale_m=K/BLK_K_S, stride_a_scale_k=1
    stride_b_scale_n, stride_b_scale_k, # B scales (N/BLK_N_S x K/BLK_K_S): stride_b_scale_n=K/BLK_K_S, stride_b_scale_k=1
    stride_cm, stride_cn,   # C (M x N): stride_cm=N, stride_cn=1
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
    BLOCK_SIZE_K_SCALE: tl.constexpr, BLOCK_SIZE_N_SCALE: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    """
    Kernel for FP8 matrix multiplication (C = A * B.T) with per-tensor/per-group scaling.
    A is M x K (row-major)
    B is N x K (row-major) -> We compute A @ B_transposed_conceptually
    C is M x N (row-major)
    a_scale is M x (K / BLOCK_K_SCALE) (row-major) - Per-row scaling for A, grouped along K
    b_scale is (N / BLOCK_N_SCALE) x (K / BLOCK_K_SCALE) (row-major) - Per-group scaling for B
    """
    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 % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    offs_k = tl.arange(0, BLOCK_SIZE_K)

    a_ptrs = a_ptr + offs_m[:, None] * stride_am
    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    b_ptrs = b_ptr + offs_bn[None, :] * stride_bn
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    a_scale_vec = tl.zeros((BLOCK_SIZE_M, 1), dtype=tl.float32) + 1.0 # Default to 1.0
    b_scale = 1.0 # Default to 1.0
    last_k_block_idx = -1 # Track the last scale block index processed
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        k_offset = k * BLOCK_SIZE_K # Base offset for K in this iteration
        current_k_block_idx = k_offset // BLOCK_SIZE_K_SCALE

        if current_k_block_idx != last_k_block_idx:
            a_scale_ptrs = a_scale_ptr + offs_m[:, None] * stride_a_scale_m + current_k_block_idx * stride_a_scale_k
            a_scale_mask = (offs_m[:, None] < M) & (current_k_block_idx < tl.cdiv(K, BLOCK_SIZE_K_SCALE))
            a_scale_vec = tl.load(a_scale_ptrs, mask=a_scale_mask, other=1.0) # Default scale to 1.0 if out of bounds

            b_scale_n_idx = pid_n # N scale block index is simply the N block index
            b_scale_ptr_scalar = b_scale_ptr + (b_scale_n_idx * stride_b_scale_n + current_k_block_idx * stride_b_scale_k)
            b_scale_mask = (b_scale_n_idx < tl.cdiv(N, BLOCK_SIZE_N_SCALE)) & (current_k_block_idx < tl.cdiv(K, BLOCK_SIZE_K_SCALE))
            b_scale = tl.load(b_scale_ptr_scalar, mask=b_scale_mask, other=1.0) # Default scale to 1.0 if out of bounds

            last_k_block_idx = current_k_block_idx # Update tracker

        offs_k_a = k_offset + offs_k # Shape: [BLOCK_SIZE_K]
        a_load_ptrs = a_ptrs + offs_k_a[None, :] * stride_ak # Shape: [BLOCK_SIZE_M, BLOCK_SIZE_K]
        a_mask = (offs_m[:, None] < M) & (offs_k_a[None, :] < K) # Shape: [BLOCK_SIZE_M, BLOCK_SIZE_K]
        a = tl.load(a_load_ptrs, mask=a_mask, other=0.0)

        offs_k_b = k_offset + offs_k # Shape: [BLOCK_SIZE_K]
        b_load_ptrs = b_ptrs + offs_k_b[:, None] * stride_bk # Shape: [BLOCK_SIZE_K, BLOCK_SIZE_N]
        b_mask = (offs_k_b[:, None] < K) & (offs_bn[None, :] < N) # Shape: [BLOCK_SIZE_K, BLOCK_SIZE_N]
        b = tl.load(b_load_ptrs, mask=b_mask, other=0.0)

        accumulator += tl.dot(a, b, out_dtype=tl.float32) * a_scale_vec * b_scale

    accumulator = accumulator.to(c_ptr.dtype.element_ty)
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)


def custom_kernel(data: input_t) -> output_t:
    global _printed_config_summaries
    a, b, a_scale, b_scale, c = data

    M, K = a.shape
    N = b.shape[0]

    BLOCK_SIZE_K_SCALE = 128
    BLOCK_SIZE_N_SCALE = 128

    def grid(META):
        # Use the META dict provided by the autotuner in the launch call
        num_m_blocks = triton.cdiv(M, META['BLOCK_SIZE_M'])
        num_n_blocks = triton.cdiv(N, META['BLOCK_SIZE_N'])
        num_m_groups = triton.cdiv(num_m_blocks, META['GROUP_SIZE_M'])
        return (num_m_groups * num_n_blocks * META['GROUP_SIZE_M'],) # Total number of programs launched

    # Launch the kernel
    matmul_fp8_kernel[grid]( # Call the correct kernel
        a, b, a_scale, b_scale, c,   # Data pointers
        M, N, K,                     # Dimensions
        # Pass appropriate strides based on memory layout
        a.stride(0), a.stride(1),    # A col-major strides (1, M)
        b.stride(0), b.stride(1),    # B col-major strides (1, N)
        a_scale.stride(0), a_scale.stride(1), # A_scale row-major strides
        b_scale.stride(0), b_scale.stride(1), # B_scale row-major strides
        c.stride(0), c.stride(1),    # C row-major strides (N, 1)
        BLOCK_SIZE_K_SCALE=BLOCK_SIZE_K_SCALE, BLOCK_SIZE_N_SCALE=BLOCK_SIZE_N_SCALE,
    )

    config_key = (M, N, K)
    # Check if autotuning has determined a best_config AND we haven't printed for this key yet
    if config_key not in _printed_config_summaries and hasattr(matmul_fp8_kernel, 'best_config'):
         best_config = matmul_fp8_kernel.best_config
         if best_config: # Check if not None (tuning might not be instant)
            try:
                # Access config values from the triton.Config object
                # The parameters passed in the constructor are stored in .kwargs
                cfg = best_config.kwargs
                m = cfg.get('BLOCK_SIZE_M', 'NA') # Use .get for safety
                n = cfg.get('BLOCK_SIZE_N', 'NA')
                k = cfg.get('BLOCK_SIZE_K', 'NA')
                gm = cfg.get('GROUP_SIZE_M', 'NA')
                # num_warps and num_stages are attributes of the Config object itself
                warps = best_config.num_warps if hasattr(best_config, 'num_warps') else 'NA'
                stages = best_config.num_stages if hasattr(best_config, 'num_stages') else 'NA'

                # Format: Mblock-Nblock-Kblock-GroupM-Warps-Stages
                concise_summary = f"{m}-{n}-{k}-{gm}-{warps}-{stages}"

                # Print to stderr, prefixed with the matrix dimensions for context
                print(f"Config [{M}x{N}x{K}]: {concise_summary}", file=sys.stderr)

                # Mark this key as printed
                _printed_config_summaries[config_key] = True

            except Exception as e:
                # Fallback in case of unexpected structure or error
                print(f"Error summarizing config for {config_key}: {e}", file=sys.stderr)
                # Optionally print the full config string on error for debugging
                # print(f"Full config was: {str(best_config)}", file=sys.stderr)
                _printed_config_summaries[config_key] = True # Still mark to avoid repeat errors
    return c

python eval.py benchmark /tmp/tmplbey8ugm
Config [1024x1536x7168]: 128-128-64-4-8-2
Config [1024x3072x1536]: 128-128-64-4-8-2
Config [1024x576x7168]: 128-128-64-4-8-2
Config [1024x7168x256]: 256-128-64-4-8-1
Config [1024x7168x2048]: 256-128-64-4-8-1
Config [1024x4608x7168]: 128-128-64-4-8-2
Config [1024x7168x2304]: 256-128-64-4-8-1
Config [1024x512x7168]: 128-128-64-4-8-2
Config [1024x4096x512]: 128-128-64-4-8-2
Config [6144x1536x7168]: 256-128-64-4-8-1
Config [6144x3072x1536]: 256-128-64-4-8-1
Config [6144x576x7168]: 128-128-128-4-4-1
Config [6144x7168x256]: 256-128-64-4-8-1
Config [6144x7168x2048]: 256-128-64-4-8-1
Config [6144x4608x7168]: 256-128-64-4-8-1
Config [6144x7168x2304]: 256-128-64-4-8-1
Config [6144x512x7168]: 128-128-64-4-8-2
Config [6144x4096x512]: 256-128-64-4-8-1
