submission 533626

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

"""Best hybrid path, with block_m forced to 32 only for 33-expert bs128."""

import csv
import importlib
import importlib.util
import os
from dataclasses import dataclass
from pathlib import Path

os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"
os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx950")
os.environ.setdefault("CXX", "clang++")

import aiter
import torch
from torch.utils.cpp_extension import load_inline

from task import input_t, output_t
from aiter import ActivationType, QuantType, dtypes, get_hip_quant
from aiter.fused_moe import fused_moe


_FIELDNAMES = [
    "cu_num", "token", "model_dim", "inter_dim", "expert", "topk", "act_type", "dtype",
    "q_dtype_a", "q_dtype_w", "q_type", "use_g1u1", "doweight_stage1", "block_m", "ksplit",
    "us1", "kernelName1", "err1", "us2", "kernelName2", "err2", "us", "run_1stage",
    "tflops", "bw", "_tag",
]


def _write_row(writer: csv.DictWriter, *, cu_num: str, token: str, inter_dim: str, expert: str, ksplit: str, block_m: str,
               kernel1: str = "Null", kernel2: str = "Null", us1: str = "0.0", us2: str = "0.0", us: str = "0.0", err2: str = "0.0%") -> None:
    writer.writerow({
        "cu_num": cu_num,
        "token": token,
        "model_dim": "7168",
        "inter_dim": inter_dim,
        "expert": expert,
        "topk": "9",
        "act_type": "ActivationType.Silu",
        "dtype": "torch.bfloat16",
        "q_dtype_a": "torch.float4_e2m1fn_x2",
        "q_dtype_w": "torch.float4_e2m1fn_x2",
        "q_type": "QuantType.per_1x32",
        "use_g1u1": "1",
        "doweight_stage1": "0",
        "block_m": block_m,
        "ksplit": ksplit,
        "us1": us1,
        "kernelName1": kernel1,
        "err1": "0.0%",
        "us2": us2,
        "kernelName2": kernel2,
        "err2": err2,
        "us": us,
        "run_1stage": "0",
        "tflops": "0.0",
        "bw": "0.0",
        "_tag": "",
    })


def _prepare_env() -> None:
    if importlib.util.find_spec("aiter") is None:
        return
    cfg_path = Path("/tmp/gpumode_amd_moe_mxfp4_cfg_all512_hipsort_v3.csv")
    with cfg_path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=_FIELDNAMES)
        writer.writeheader()
        for cu_num in ("256", "288", "304"):
            for token in ("16", "128"):
                _write_row(writer, cu_num=cu_num, token=token, inter_dim="256", expert="257", ksplit="2", block_m="32")
            _write_row(
                writer, cu_num=cu_num, token="512", inter_dim="256", expert="257", ksplit="0", block_m="32",
                kernel1="moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
                kernel2="moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16",
                us1="97.693", us2="70.2081", us="167.9011", err2="1.3%",
            )
            for token in ("16", "128"):
                _write_row(writer, cu_num=cu_num, token=token, inter_dim="512", expert="33", ksplit="2", block_m="32")
    os.environ["AITER_CONFIG_FMOE"] = str(cfg_path)
    os.environ.setdefault("AITER_LOG_LEVEL", "ERROR")


def _block_size(config: dict) -> int | None:
    if config["n_routed_experts"] != 32 or config["n_shared_experts"] != 1:
        return None
    if config["bs"] == 128 and config["d_expert"] == 512:
        return 32
    if config["bs"] == 512 and config["d_expert"] == 512:
        return 128
    if config["bs"] == 512 and config["d_expert"] == 2048:
        return 64
    return None


def _use_direct_33_512(config: dict) -> bool:
    return (
        config["bs"] == 512
        and config["n_routed_experts"] == 32
        and config["n_shared_experts"] == 1
        and config["d_expert"] == 512
    )


_prepare_env()

_CPP_WRAPPER = """
void moe_scale_sort_vec8(torch::Tensor scale,
                         torch::Tensor sorted_ids,
                         torch::Tensor num_valid_ids,
                         torch::Tensor out,
                         int token_num,
                         int topk);
"""

_CUDA_SRC = r"""
#include <torch/extension.h>
#include <hip/hip_runtime.h>
#include <cstdint>

__global__ void moe_scale_sort_vec8_kernel(
    const uint8_t* __restrict__ scale,
    const int32_t* __restrict__ sorted_ids,
    const int32_t* __restrict__ num_valid_ids,
    uint32_t* __restrict__ out,
    int scale_cols,
    int out_blocks_n,
    int token_num,
    int topk) {
  __shared__ uint64_t packed_rows[32];

  const int tid = threadIdx.x;
  const int tile_m = blockIdx.x;
  const int tile_n = blockIdx.y;
  const int row_base = tile_m * 32;
  const int col_base = tile_n * 8;
  const int valid = num_valid_ids[0];

  if (tid < 32) {
    const int out_row = row_base + tid;
    uint64_t bytes = 0;
    if (out_row < valid && col_base + 8 <= scale_cols) {
      const uint32_t sid = static_cast<uint32_t>(sorted_ids[out_row]);
      const int topk_id = static_cast<int>(sid >> 24);
      const int tok = static_cast<int>(sid & 0xFFFFFFu);
      const int src_row = (topk == 1) ? tok : tok * topk + topk_id;
      if (tok < token_num) {
        const uint64_t* src = reinterpret_cast<const uint64_t*>(scale + static_cast<size_t>(src_row) * scale_cols + col_base);
        bytes = *src;
      }
    }
    packed_rows[tid] = bytes;
  }
  __syncthreads();

  if (tid < 64) {
    const int r = tid & 15;
    const int c = tid >> 4;
    const uint64_t lo = packed_rows[r];
    const uint64_t hi = packed_rows[r + 16];
    const uint32_t word =
        ((lo >> (8 * c)) & 0xFFu) |
        (((hi >> (8 * c)) & 0xFFu) << 8) |
        (((lo >> (8 * (c + 4))) & 0xFFu) << 16) |
        (((hi >> (8 * (c + 4))) & 0xFFu) << 24);
    const size_t out_idx = (((static_cast<size_t>(tile_m) * out_blocks_n + tile_n) * 4 + c) * 16 + r);
    out[out_idx] = word;
  }
}

void moe_scale_sort_vec8(torch::Tensor scale,
                         torch::Tensor sorted_ids,
                         torch::Tensor num_valid_ids,
                         torch::Tensor out,
                         int token_num,
                         int topk) {
  const int scale_cols = static_cast<int>(scale.size(1));
  const int out_blocks_m = static_cast<int>(out.size(0));
  const int out_blocks_n = static_cast<int>(out.size(1));
  dim3 block(64);
  dim3 grid(out_blocks_m, out_blocks_n);
  hipLaunchKernelGGL(
      moe_scale_sort_vec8_kernel,
      grid,
      block,
      0,
      0,
      reinterpret_cast<const uint8_t*>(scale.data_ptr()),
      sorted_ids.data_ptr<int32_t>(),
      num_valid_ids.data_ptr<int32_t>(),
      reinterpret_cast<uint32_t*>(out.data_ptr()),
      scale_cols,
      out_blocks_n,
      token_num,
      topk);
}
"""

_sort_mod = load_inline(
    name="moe_scale_sort_vec8_gfx950_v3",
    cpp_sources=[_CPP_WRAPPER],
    cuda_sources=[_CUDA_SRC],
    functions=["moe_scale_sort_vec8"],
    verbose=False,
    extra_cuda_cflags=["--offload-arch=gfx950", "-std=c++20", "-O3"],
)

_fused_moe_mod = importlib.import_module("aiter.fused_moe")
_orig_fused_qsort = _fused_moe_mod.fused_dynamic_mxfp4_quant_moe_sort
_hip_quant_fp4 = get_hip_quant(QuantType.per_1x32)


def _hip_vec8_sort(scale_e8m0, sorted_ids, num_valid_ids, token_num, topk):
    if topk > 1:
        scale_e8m0 = scale_e8m0[: token_num * topk, :].view(token_num * topk, -1)
    _, n_i = scale_e8m0.shape
    out = torch.empty(((sorted_ids.shape[0] + 31) // 32, (n_i + 7) // 8, 4, 16), dtype=torch.uint32, device=scale_e8m0.device)
    _sort_mod.moe_scale_sort_vec8(scale_e8m0.view(torch.uint8), sorted_ids, num_valid_ids, out, token_num, topk)
    return out.view(dtypes.fp8_e8m0).view(-1, n_i)


def _patched_fused_dynamic_mxfp4_quant_moe_sort(x, sorted_ids, num_valid_ids, token_num, topk, block_size=32, scaling_mode="even"):
    if token_num == 512:
        x_q, x_scale = _hip_quant_fp4(x, scale=None, quant_dtype=dtypes.fp4x2, num_rows=None, num_rows_factor=topk)
        return x_q, _hip_vec8_sort(x_scale, sorted_ids, num_valid_ids, token_num, topk)
    return _orig_fused_qsort(
        x,
        sorted_ids=sorted_ids,
        num_valid_ids=num_valid_ids,
        token_num=token_num,
        topk=topk,
        block_size=block_size,
        scaling_mode=scaling_mode,
    )


_fused_moe_mod.fused_dynamic_mxfp4_quant_moe_sort = _patched_fused_dynamic_mxfp4_quant_moe_sort


@dataclass
class _DirectWorkspace:
    sorted_ids: torch.Tensor
    sorted_weights: torch.Tensor
    sorted_expert_ids: torch.Tensor
    num_valid_ids: torch.Tensor
    moe_buf: torch.Tensor
    stage1_out: torch.Tensor


_DIRECT_WORKSPACES: dict[tuple[int, int, int, int, int], _DirectWorkspace] = {}


def _get_direct_workspace(
    hidden_states: torch.Tensor, topk_ids: torch.Tensor, config: dict
) -> _DirectWorkspace:
    key = (
        hidden_states.device.index if hidden_states.device.index is not None else -1,
        config["bs"],
        config["d_hidden"],
        config["d_expert_pad"],
        config["n_routed_experts"] + config["n_shared_experts"],
    )
    ws = _DIRECT_WORKSPACES.get(key)
    if ws is not None:
        return ws

    block_m = 128
    num_experts = config["n_routed_experts"] + config["n_shared_experts"]
    token_num, topk = topk_ids.shape
    max_num_tokens_padded = int(topk_ids.numel() + num_experts * block_m - topk)
    max_num_m_blocks = (max_num_tokens_padded + block_m - 1) // block_m
    ws = _DirectWorkspace(
        sorted_ids=torch.empty(max_num_tokens_padded, dtype=dtypes.i32, device=hidden_states.device),
        sorted_weights=torch.empty(max_num_tokens_padded, dtype=dtypes.fp32, device=hidden_states.device),
        sorted_expert_ids=torch.empty(max_num_m_blocks, dtype=dtypes.i32, device=hidden_states.device),
        num_valid_ids=torch.empty(2, dtype=dtypes.i32, device=hidden_states.device),
        moe_buf=torch.empty((token_num, config["d_hidden"]), dtype=hidden_states.dtype, device=hidden_states.device),
        stage1_out=torch.empty((token_num, topk, config["d_expert_pad"]), dtype=hidden_states.dtype, device=hidden_states.device),
    )
    _DIRECT_WORKSPACES[key] = ws
    return ws


def _run_direct_33_512(
    hidden_states: torch.Tensor,
    gate_up_weight_shuffled: torch.Tensor,
    down_weight_shuffled: torch.Tensor,
    gate_up_weight_scale_shuffled: torch.Tensor,
    down_weight_scale_shuffled: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    config: dict,
) -> torch.Tensor:
    ws = _get_direct_workspace(hidden_states, topk_ids, config)
    num_experts = config["n_routed_experts"] + config["n_shared_experts"]
    token_num, topk = topk_ids.shape
    inter_dim = config["d_expert_pad"]
    block_m = 128

    aiter.moe_sorting_opus_fwd(
        topk_ids,
        topk_weights,
        ws.sorted_ids,
        ws.sorted_weights,
        ws.sorted_expert_ids,
        ws.num_valid_ids,
        ws.moe_buf,
        num_experts,
        block_m,
        None,
        None,
        0,
    )

    a1_q, a1_scale = _patched_fused_dynamic_mxfp4_quant_moe_sort(
        hidden_states,
        sorted_ids=ws.sorted_ids,
        num_valid_ids=ws.num_valid_ids,
        token_num=token_num,
        topk=1,
        block_size=block_m,
    )
    aiter.ck_moe_stage1_fwd(
        a1_q,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        ws.sorted_ids,
        ws.sorted_expert_ids,
        ws.num_valid_ids,
        ws.stage1_out,
        topk,
        "",
        gate_up_weight_scale_shuffled.view(dtypes.fp8_e8m0),
        a1_scale,
        block_m,
        None,
        QuantType.per_1x32,
        ActivationType.Silu,
        0,
        True,
        ws.stage1_out.dtype,
    )

    a2_q, a2_scale = _patched_fused_dynamic_mxfp4_quant_moe_sort(
        ws.stage1_out.view(-1, inter_dim),
        sorted_ids=ws.sorted_ids,
        num_valid_ids=ws.num_valid_ids,
        token_num=token_num,
        topk=topk,
        block_size=block_m,
    )
    a2_q = a2_q.view(token_num, topk, -1)
    aiter.ck_moe_stage2_fwd(
        a2_q,
        gate_up_weight_shuffled,
        down_weight_shuffled,
        ws.sorted_ids,
        ws.sorted_expert_ids,
        ws.num_valid_ids,
        ws.moe_buf,
        topk,
        "",
        down_weight_scale_shuffled.view(dtypes.fp8_e8m0),
        a2_scale,
        block_m,
        ws.sorted_weights,
        QuantType.per_1x32,
        ActivationType.Silu,
        True,
    )
    return ws.moe_buf


def custom_kernel(data: input_t) -> output_t:
    (
        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

    if _use_direct_33_512(config):
        return _run_direct_33_512(
            hidden_states,
            gate_up_weight_shuffled,
            down_weight_shuffled,
            gate_up_weight_scale_shuffled,
            down_weight_scale_shuffled,
            topk_weights,
            topk_ids,
            config,
        )

    return 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,
        block_size_M=_block_size(config),
        hidden_pad=config["d_hidden_pad"] - config["d_hidden"],
        intermediate_pad=config["d_expert_pad"] - config["d_expert"],
    )