submission 677065

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

"""
优化版本2：条件性 contiguous() + 减少重复导入

优化点：
1. 只在必要时调用 contiguous()
2. 将导入移到函数外部（避免每次调用都导入）
3. 内联量化函数（减少函数调用开销）

预期提升：5-12%（比 v1 更好）
"""
import aiter
from aiter import QuantType, dtypes
from aiter.ops.triton.quant import dynamic_mxfp4_quant 
from aiter.utility.fp4_utils import e8m0_shuffle
from task import input_t, output_t


def custom_kernel(data: input_t) -> output_t:
    """
    Optimized: Conditional contiguous() + reduced overhead.
    """
    A, B, B_q, B_shuffle, B_scale_sh = data
    
    # 🔥 优化1: 条件性 contiguous()
    # 只在必要时才调用，避免不必要的内存拷贝
    if not A.is_contiguous():
        A = A.contiguous()
    if not B.is_contiguous():
        B = B.contiguous()
    
    # 🔥 优化2: 内联量化逻辑，减少函数调用开销
    A_fp4, A_scale_raw = dynamic_mxfp4_quant(A)
    A_scale_sh = e8m0_shuffle(A_scale_raw)
    A_q = A_fp4.view(dtypes.fp4x2)
    A_scale_sh = A_scale_sh.view(dtypes.fp8_e8m0)
    
    out_gemm = aiter.gemm_a4w4(
        A_q,
        B_shuffle,
        A_scale_sh,
        B_scale_sh,
        dtype=dtypes.bf16,
        bpreshuffle=True,
    )
    return out_gemm