submission 629606

#!POPCORN leaderboard amd-mxfp4-mm
#!POPCORN gpu MI355X

"""
Advanced FP4 quant + FP4 GEMM with custom HIP kernels for MI355X.
This version includes optional custom quantization kernel for better performance.
"""
from task import input_t, output_t

import os
# Set ROCm architecture before importing torch
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'
os.environ['CXX'] = 'clang++'

import torch
import aiter
from aiter import dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle

# Try to use custom HIP kernel for quantization (fallback to aiter if compilation fails)
USE_CUSTOM_QUANT = False  # Set to True to enable custom HIP quant kernel

if USE_CUSTOM_QUANT:
    try:
        from torch.utils.cpp_extension import load_inline

        HIP_QUANT_SRC = """
#include <hip/hip_runtime.h>
#include <hip/amd_detail/amd_hip_bf16.h>

// MXFP4 quantization kernel: per-1x32 grouping with e8m0 scale
// Each block processes multiple groups of 32 elements
__global__ void mxfp4_quant_kernel(
    const __hip_bfloat16* __restrict__ input,
    uint8_t* __restrict__ output_fp4,
    uint8_t* __restrict__ output_scale,
    int M, int K
) {
    // Group size for per-1x32 quantization
    const int GROUP_SIZE = 32;

    int group_idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total_groups = M * ((K + GROUP_SIZE - 1) / GROUP_SIZE);

    if (group_idx >= total_groups) return;

    int row = group_idx / ((K + GROUP_SIZE - 1) / GROUP_SIZE);
    int col_group = group_idx % ((K + GROUP_SIZE - 1) / GROUP_SIZE);
    int col_start = col_group * GROUP_SIZE;

    // Find max absolute value in group for scale calculation
    float max_abs = 0.0f;
    for (int i = 0; i < GROUP_SIZE && (col_start + i) < K; ++i) {
        float val = __bfloat162float(input[row * K + col_start + i]);
        max_abs = fmaxf(max_abs, fabsf(val));
    }

    // Calculate e8m0 scale (exponent only, 8-bit)
    // Scale = 2^exponent where exponent = floor(log2(max_abs)) - 2 (for FP4 range)
    uint8_t scale_e8m0 = 0;
    if (max_abs > 0.0f) {
        int exp = (int)floorf(log2f(max_abs)) - 2;
        scale_e8m0 = (uint8_t)(exp + 127);  // e8m0 bias
    }

    output_scale[group_idx] = scale_e8m0;

    // Quantize each element to FP4
    float scale_val = exp2f((float)(scale_e8m0) - 127.0f);
    float inv_scale = scale_val > 0.0f ? 1.0f / scale_val : 0.0f;

    // Pack 2 FP4 values per byte
    for (int i = 0; i < GROUP_SIZE && (col_start + i) < K; i += 2) {
        uint8_t packed = 0;
        for (int j = 0; j < 2 && (col_start + i + j) < K; ++j) {
            float val = __bfloat162float(input[row * K + col_start + i + j]);
            float scaled = val * inv_scale;
            // Simple FP4 quantization (simplified - actual MXFP4 has specific encoding)
            int quantized = (int)roundf(scaled * 6.0f);  // Scale to FP4 range
            quantized = max(-6, min(6, quantized));      // Clamp
            uint8_t fp4_bits = (uint8_t)(quantized + 8); // Offset to unsigned
            if (j == 0) {
                packed = fp4_bits & 0x0F;
            } else {
                packed |= (fp4_bits & 0x0F) << 4;
            }
        }
        output_fp4[row * ((K + 1) / 2) + (col_start + i) / 2] = packed;
    }
}

torch::Tensor mxfp4_quant(torch::Tensor input) {
    int M = input.size(0);
    int K = input.size(1);

    auto output_fp4 = torch::empty({M, (K + 1) / 2}, torch::dtype(torch::kUInt8).device(input.device()));
    auto output_scale = torch::empty({M, (K + 31) / 32}, torch::dtype(torch::kUInt8).device(input.device()));

    const int threads = 256;
    int total_groups = M * ((K + 31) / 32);
    const int blocks = (total_groups + threads - 1) / threads;

    mxfp4_quant_kernel<<<blocks, threads>>>(
        reinterpret_cast<const __hip_bfloat16*>(input.data_ptr<at::BFloat16>()),
        output_fp4.data_ptr<uint8_t>(),
        output_scale.data_ptr<uint8_t>(),
        M, K
    );

    return torch::stack({output_fp4, output_scale});
}
"""

        CPP_SRC = """
torch::Tensor mxfp4_quant(torch::Tensor input);
"""

        _hip_module = load_inline(
            name='mxfp4_quant_module',
            cpp_sources=[CPP_SRC],
            cuda_sources=[HIP_QUANT_SRC],
            functions=['mxfp4_quant'],
            verbose=True,
            extra_cuda_cflags=["--offload-arch=gfx942", "-std=c++20", "-O3"],
        )

        def _custom_quant_mxfp4(x):
            result = _hip_module.mxfp4_quant(x)
            return result[0].view(dtypes.fp4x2), result[1].view(dtypes.fp8_e8m0)

        _quant_func = _custom_quant_mxfp4
        print("Custom HIP quant kernel loaded successfully")

    except Exception as e:
        print(f"Failed to load custom HIP kernel, falling back to aiter: {e}")
        USE_CUSTOM_QUANT = False


def _aiter_quant_mxfp4(x: torch.Tensor) -> tuple:
    """MXFP4 quantization using aiter library."""
    x_fp4, bs_e8m0 = dynamic_mxfp4_quant(x)
    bs_e8m0 = e8m0_shuffle(bs_e8m0)
    return x_fp4.view(dtypes.fp4x2), bs_e8m0.view(dtypes.fp8_e8m0)


# Select quantization function
_quant_mxfp4 = _custom_quant_mxfp4 if USE_CUSTOM_QUANT else _aiter_quant_mxfp4


def custom_kernel(data: input_t) -> output_t:
    """
    MXFP4 GEMM with optimized quantization.
    Uses custom HIP kernel if available, otherwise falls back to aiter.
    """
    A, _, _, B_shuffle, B_scale_sh = data

    # Quantize A to MXFP4
    A_q, A_scale_sh = _quant_mxfp4(A)

    # 4-bit GEMM
    return aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )