submission 529158

import torch
import math
import aiter
from aiter import QuantType, dtypes, ActivationType
from aiter.fused_moe import fused_moe

# --- [1] LOGICA MXFP4 GEMM (aiter NATIVE) ---
def gemm_logic(data):
    # Firma: (A, B, B_q, B_shuffle, B_scale_sh)
    A, B, B_q, B_shuffle, B_scale_sh = data[:5]
    A = A.contiguous()
    quant_func = aiter.get_triton_quant(QuantType.per_1x32)
    A_q, A_scale_sh = quant_func(A, shuffle=True)
    return aiter.gemm_a4w4(
        A_q, B_shuffle, A_scale_sh, B_scale_sh,
        dtype=dtypes.bf16, bpreshuffle=True
    )

# --- [2] LOGICA MXFP4 MoE ---
def moe_logic(data):
    # Firma: 12 elementos
    (h_states, _, _, _, _, gu_shuf, d_shuf, gus_shuf, ds_shuf, 
     tk_weights, tk_ids, config) = data[:12]
    h_pad = config.get("d_hidden_pad", config.get("d_hidden", 0)) - config.get("d_hidden", 0)
    i_pad = config.get("d_expert_pad", config.get("d_expert", 0)) - config.get("d_expert", 0)
    return fused_moe(
        h_states, gu_shuf, d_shuf, tk_weights, tk_ids, activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32, w1_scale=gus_shuf, w2_scale=ds_shuf,
        hidden_pad=h_pad, intermediate_pad=i_pad
    )

# --- [3] LOGICA MLA ---
def mla_logic(data):
    q, kv_data, qo_indptr, kv_indptr, config = data[:5]
    sm_scale = config.get("sm_scale", 1.0 / math.sqrt(576))
    kv_buffer = kv_data['bf16'] if 'bf16' in kv_data else kv_data['kv']
    from aiter.mla import mla_decode_fwd
    return mla_decode_fwd(q, kv_buffer, None, qo_indptr, kv_indptr, None, None, 
                          q.shape[0], sm_scale=sm_scale, nhead_kv=1)

# --- SOBERANO DISPATCHER V5 ---
def submission(*args, **kwargs):
    # Consolidamos argumentos
    if len(args) == 0: return None
    
    # Caso 1: Los argumentos vienen expandidos (*args)
    # Caso 2: Los argumentos vienen en una tupla (args[0])
    if len(args) == 1 and isinstance(args[0], (tuple, list)):
        payload = args[0]
    else:
        payload = args
        
    n = len(payload)
    
    if n == 12: # MoE
        return moe_logic(payload)
    if n == 5:
        # MLA vs GEMM
        # MLA payload[1] es dict (kv_data)
        # GEMM payload[1] es tensor (B)
        if isinstance(payload[1], dict):
            return mla_logic(payload)
        return gemm_logic(payload)
        
    return None

check_implementation = submission
custom_kernel = submission