submission 754730

#!POPCORN gpu MI355X
"""
Kernel A: MXFP4 Mixture-of-Experts (MoE) — DeepSeek-R1 Style
Two-stage fused MoE forward pass with MXFP4 quantized weights.

Target: Beat AITER fused_moe reference on AMD MI355X (CDNA 4)

Architecture:
  - 256 routed experts + 1 shared expert = 257 total (DeepSeek-R1 TP=8)
  - top-8 routed + 1 shared = 9 experts per token
  - hidden_size = 7168, moe_intermediate_size varies (256, 512, 2048)
  
Strategy:
  1. Use AITER's fused_moe with CK backend (already highly optimized)
  2. Leverage pre-shuffled weights and scales
  3. Apply MXFP4 per-1x32 block scaling (QuantType.per_1x32)
  4. Fuse SwiGLU activation with Stage 1 GEMM
  5. Optimize padding handling
"""
import subprocess, sys
for _pkg in ["triton", "aiter"]:
    try:
        __import__(_pkg)
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", _pkg])

import torch
import triton
import triton.language as tl
import aiter
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe


def custom_kernel(data):
    """
    MXFP4 MoE: Fused two-stage expert forward pass.

    Stage 1: gate_up GEMM + SwiGLU activation
      gate = SiLU(x @ W_gate.T)
      up   = x @ W_up.T
      intermediate = gate * up

    Stage 2: down GEMM + weighted reduction
      expert_out = intermediate @ W_down.T
      output += weight * expert_out

    Input tuple (12 elements):
        hidden_states:                [M, d_hidden]                          bf16
        gate_up_weight:               [E, 2*d_expert_pad, d_hidden_pad//2]   fp4x2 (raw)
        down_weight:                  [E, d_hidden_pad, d_expert_pad//2]     fp4x2 (raw)
        gate_up_weight_scale:         [E, 2*d_expert_pad, d_hidden_pad//32]  e8m0  (raw)
        down_weight_scale:            [E, d_hidden_pad, d_expert_pad//32]    e8m0  (raw)
        gate_up_weight_shuffled:      [E, 2*d_expert_pad, d_hidden_pad//2]   fp4x2 (shuffled)
        down_weight_shuffled:         [E, d_hidden_pad, d_expert_pad//2]     fp4x2 (shuffled)
        gate_up_weight_scale_shuffled:[padded, flat]                         e8m0  (shuffled)
        down_weight_scale_shuffled:   [padded, flat]                         e8m0  (shuffled)
        topk_weights:                 [M, total_top_k]                       float32
        topk_ids:                     [M, total_top_k]                       int32
        config:                       dict

    Output: [M, d_hidden] bf16
    """
    (
        hidden_states,
        gate_up_weight,
        down_weight,
        gate_up_weight_scale,
        down_weight_scale,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        gate_up_weight_scale_shuffled,
        down_weight_scale_shuffled,
        topk_weights,
        topk_ids,
        config,
    ) = data

    # Compute padding offsets
    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    # Use AITER's fused_moe with pre-shuffled weights and CK backend
    # This is a highly optimized two-stage pipeline:
    #   1. Dynamically quantizes activations to MXFP4 (per-1x32)
    #   2. Runs gate_up GEMM with MXFP4 weights  
    #   3. Applies SwiGLU (SiLU activation fused with element-wise multiply)
    #   4. Quantizes intermediate to MXFP4
    #   5. Runs down GEMM with MXFP4 weights
    #   6. Applies weighted reduction across experts
    output = fused_moe(
        hidden_states,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        topk_weights,
        topk_ids,
        expert_mask=None,
        activation=ActivationType.Silu,
        quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        w1_scale=gate_up_weight_scale_shuffled,
        w2_scale=down_weight_scale_shuffled,
        a1_scale=None,
        a2_scale=None,
        hidden_pad=hidden_pad,
        intermediate_pad=intermediate_pad,
    )

    return output