submission 583982

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

import torch
import triton
from task import input_t, output_t
from aiter import dtypes
import aiter
from aiter.ops.triton._triton_kernels.quant.quant import _dynamic_mxfp4_quant_kernel

# Global caches for Zero-Allocation architecture
_cache_fp4 = {}
_cache_e8m0 = {}
_cache_padded = {}
_cache_shuffled = {}

def custom_kernel(data: input_t) -> output_t:
    """
    Ultra-optimized MXFP4 Quantization + GEMM.
    Zero-Allocation pipeline: Bypasses all dynamic memory allocations in
    dynamic_mxfp4_quant and e8m0_shuffle, saving ~3us overhead per invocation.
    """
    A, B, B_q, B_shuffle, B_scale_sh = data
    M, K = A.shape
    key = (M, K)
    
    if key not in _cache_fp4:
        MXFP4_QUANT_BLOCK_SIZE = 32
        x_fp4 = torch.empty((M, K // 2), dtype=torch.uint8, device=A.device)
        blockscale_e8m0 = torch.empty(
            ((K + MXFP4_QUANT_BLOCK_SIZE - 1) // MXFP4_QUANT_BLOCK_SIZE, M),
            dtype=torch.uint8,
            device=A.device,
        ).T
        
        m, n = blockscale_e8m0.shape
        scale_padded = torch.empty(
            (m + 255) // 256 * 256,
            (n + 7) // 8 * 8,
            dtype=torch.uint8,
            device=A.device,
        )
        
        sm, sn = scale_padded.shape
        scale_shuffled = torch.empty((sm, sn), dtype=torch.uint8, device=A.device)
        
        _cache_fp4[key] = x_fp4
        _cache_e8m0[key] = blockscale_e8m0
        _cache_padded[key] = scale_padded
        _cache_shuffled[key] = scale_shuffled
        
    x_fp4 = _cache_fp4[key]
    blockscale_e8m0 = _cache_e8m0[key]
    scale_padded = _cache_padded[key]
    scale_shuffled = _cache_shuffled[key]
    
    # 1. Inline Triton MXFP4 Quantization Kernel Invocation
    if M <= 32:
        NUM_ITER = 1
        BLOCK_SIZE_M = triton.next_power_of_2(M)
        BLOCK_SIZE_N = 32
        NUM_WARPS = 1
        NUM_STAGES = 1
    else:
        NUM_ITER = 4
        BLOCK_SIZE_M = 64
        BLOCK_SIZE_N = 64
        NUM_WARPS = 4
        NUM_STAGES = 2
        if K <= 16384:
            BLOCK_SIZE_M = 32
            BLOCK_SIZE_N = 128
            
    # For small K values
    if K <= 1024:
        NUM_ITER = 1
        NUM_STAGES = 1
        NUM_WARPS = 4
        BLOCK_SIZE_N = min(256, triton.next_power_of_2(K))
        BLOCK_SIZE_N = max(32, BLOCK_SIZE_N)
        BLOCK_SIZE_M = min(8, triton.next_power_of_2(M))

    grid = (
        triton.cdiv(M, BLOCK_SIZE_M),
        triton.cdiv(K, BLOCK_SIZE_N * NUM_ITER),
    )

    _dynamic_mxfp4_quant_kernel[grid](
        A,
        x_fp4,
        blockscale_e8m0,
        *A.stride(),
        *x_fp4.stride(),
        *blockscale_e8m0.stride(),
        M=M,
        N=K,
        MXFP4_QUANT_BLOCK_SIZE=32,
        SCALING_MODE=0,
        NUM_ITER=NUM_ITER,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        NUM_STAGES=NUM_STAGES,
        num_warps=NUM_WARPS,
        waves_per_eu=0,
        num_stages=1,
    )

    # 2. In-place Padding and Memory Shuffle (Replaces e8m0_shuffle)
    m, n = blockscale_e8m0.shape
    scale_padded[:m, :n].copy_(blockscale_e8m0)
    
    sm, sn = scale_padded.shape
    viewed = scale_padded.view(sm // 32, 2, 16, sn // 8, 2, 4)
    # Permute and copy to guarantee zero new memory allocations instead of .contiguous()
    scale_shuffled.view(sm // 32, sn // 8, 4, 16, 2, 2).copy_(viewed.permute(0, 3, 5, 2, 4, 1))
    
    A_q = x_fp4.view(dtypes.fp4x2)
    A_scale_sh = scale_shuffled.view(dtypes.fp8_e8m0)
    
    # 3. Call C++ GEMM with pre-shuffled quantized layout
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=torch.bfloat16,
        bpreshuffle=True,
    )
    
    return out_gemm