kvcache-ai--ktransformers/kt-kernel/examples/test_mxfp8_moe_avx2.py

"""AVX2 MXFP8 MoE validation for MiniMax-M3-Preview.

Forces the AVX2 backend (`kt_kernel_ext.moe.AVX2MXFP8_MOE`) for layer-`LAYER`
experts, compares output against a torch dequant+matmul reference.

Optional `--compare-amx` flag runs both AMX and AVX2 backends on identical
inputs and asserts numerical equivalence. The two paths share buffer layout
and do the same FMA arithmetic; observed differences should be at BF16-noise
level (typical mean abs ~ 1e-4).

Usage:
    python test_mxfp8_moe_avx2.py --weight-path /mnt/data/models/Minimax-M3-preview
    python test_mxfp8_moe_avx2.py --weight-path ... --compare-amx
"""

from __future__ import annotations

import argparse
import os
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/build")
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/python")

from kt_kernel import kt_kernel_ext  # noqa: E402
from kt_kernel.utils.loader import MXFP8SafeTensorLoader  # noqa: E402


# ---- Reference implementation (shared with test_mxfp8_moe_m3.py) ----

def dequantize_mxfp8(weight_u8: torch.Tensor, scale_u8: torch.Tensor, group_size: int = 32) -> torch.Tensor:
    """Dequantize [N, K] FP8 E4M3fn (as uint8) with [N, K/gs] ue8m0 scales -> [N, K] BF16."""
    n, k = weight_u8.shape
    assert k % group_size == 0
    assert scale_u8.shape == (n, k // group_size)
    w_fp32 = weight_u8.view(torch.float8_e4m3fn).float()
    s_fp32 = (2.0 ** (scale_u8.to(torch.int32) - 127)).float()
    s_full = s_fp32.repeat_interleave(group_size, dim=-1)
    return (w_fp32 * s_full).to(torch.bfloat16).contiguous()


def reference_mlp_m3(x, gate, up, down, alpha=1.702, limit=7.0):
    g = torch.mm(x.float(), gate.float().t()).clamp(-limit, limit)
    u = torch.mm(x.float(), up.float().t()).clamp(-limit, limit)
    act = g * torch.sigmoid(g * alpha) * (u + 1.0)
    return torch.mm(act, down.float().t())


def reference_moe_m3(hidden, expert_ids, weights, gate_w, up_w, down_w, alpha=1.702, limit=7.0):
    out = torch.zeros(hidden.shape[0], down_w.shape[1], dtype=torch.float32)
    for tok in range(hidden.shape[0]):
        for slot in range(expert_ids.shape[1]):
            eid = int(expert_ids[tok, slot])
            if eid < 0:
                continue
            w = float(weights[tok, slot])
            y = reference_mlp_m3(hidden[tok:tok+1], gate_w[eid], up_w[eid], down_w[eid], alpha, limit)
            out[tok] += w * y[0]
    return out.to(hidden.dtype)


# ---- Backend dispatch ----

def _backend_cls(backend_name: str):
    if backend_name == "amx":
        cls = getattr(kt_kernel_ext.moe, "AMXMXFP8_KGroup_MOE", None)
        if cls is None:
            raise RuntimeError("AMXMXFP8_KGroup_MOE not in .so — rebuild with AVX-512 + VBMI + AMX")
        return cls
    if backend_name == "avx2":
        cls = getattr(kt_kernel_ext.moe, "AVX2MXFP8_MOE", None)
        if cls is None:
            raise RuntimeError("AVX2MXFP8_MOE not in .so — rebuild with the AVX2 MXFP8 wiring (PR #2041 + this fix)")
        return cls
    raise ValueError(f"unknown backend {backend_name}")


def run_backend(backend_name, weights, expert_num, top_k, hidden_size, intermediate_size,
                layer_idx, qlen, cpu_threads, x, expert_ids, routing, physical_to_logical):
    cpu_infer = kt_kernel_ext.CPUInfer(cpu_threads)
    cfg = kt_kernel_ext.moe.MOEConfig(expert_num, top_k, hidden_size, intermediate_size, 0)
    cfg.layer_idx = layer_idx
    cfg.max_len = max(qlen, 1)
    cfg.pool = cpu_infer.backend_
    cfg.quant_config.bits = 8
    cfg.quant_config.group_size = 32
    cfg.quant_config.zero_point = False
    cfg.swiglu_alpha = 1.702
    cfg.swiglu_limit = 7.0

    cfg.gate_projs = [[t.data_ptr() for t in weights["gate"]]]
    cfg.up_projs = [[t.data_ptr() for t in weights["up"]]]
    cfg.down_projs = [[t.data_ptr() for t in weights["down"]]]
    cfg.gate_scales = [[t.data_ptr() for t in weights["gate_scale"]]]
    cfg.up_scales = [[t.data_ptr() for t in weights["up_scale"]]]
    cfg.down_scales = [[t.data_ptr() for t in weights["down_scale"]]]

    moe = _backend_cls(backend_name)(cfg)
    cpu_infer.submit(moe.load_weights_task(physical_to_logical.data_ptr()))
    cpu_infer.sync()

    bsz = torch.tensor([qlen], dtype=torch.int32)
    y = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
    cpu_infer.submit(
        moe.forward_task(bsz.data_ptr(), top_k, expert_ids.data_ptr(), routing.data_ptr(),
                         x.data_ptr(), y.data_ptr(), False)
    )
    cpu_infer.sync()
    return y


def parse_args():
    p = argparse.ArgumentParser(description="MXFP8 MoE AVX2 backend test for MiniMax M3")
    p.add_argument("--weight-path", required=True)
    p.add_argument("--layer", type=int, default=3, help="Layer index (default 3 = first MoE layer)")
    p.add_argument("--qlen", type=int, default=1)
    p.add_argument("--top-k", type=int, default=4)
    p.add_argument("--cpu-threads", type=int, default=32)
    p.add_argument("--max-experts", type=int, default=0, help="Cap experts loaded (0=all)")
    p.add_argument("--compare-amx", action="store_true",
                   help="Also run AMX backend and assert numerical equivalence with AVX2.")
    p.add_argument("--ref-threshold", type=float, default=0.10,
                   help="Max relative error vs torch reference (default 10%).")
    p.add_argument("--equiv-threshold", type=float, default=0.01,
                   help="Max relative error AMX vs AVX2 (default 1%).")
    return p.parse_args()


def main():
    args = parse_args()
    torch.manual_seed(42)

    print(f"[AVX2-MXFP8] Loading layer {args.layer} from {args.weight_path}")
    loader = MXFP8SafeTensorLoader(args.weight_path)
    weights = loader.load_experts(f"language_model.model.layers.{args.layer}")

    expert_num = len(weights["gate"])
    if args.max_experts and args.max_experts < expert_num:
        for k in ("gate", "up", "down", "gate_scale", "up_scale", "down_scale"):
            weights[k] = weights[k][:args.max_experts]
        expert_num = args.max_experts

    gate0 = weights["gate"][0]
    intermediate_size = gate0.shape[0]
    hidden_size = gate0.shape[1]
    group_size = hidden_size // weights["gate_scale"][0].shape[1]
    print(f"[AVX2-MXFP8] expert_num={expert_num} hidden={hidden_size} inter={intermediate_size} gs={group_size}")
    assert group_size == 32, f"Expected group_size=32, got {group_size}"

    physical_to_logical = torch.arange(expert_num, dtype=torch.int64).contiguous()

    qlen = args.qlen
    top_k = args.top_k
    expert_ids = torch.stack([torch.randperm(expert_num)[:top_k] for _ in range(qlen)]).to(torch.int64).contiguous()
    routing = torch.randn((qlen, top_k), dtype=torch.float32).abs().contiguous()
    routing = routing / routing.sum(dim=-1, keepdim=True)
    x = (torch.randn((qlen, hidden_size), dtype=torch.bfloat16) * 0.01).contiguous()

    # ---- AVX2 forward ----
    print("[AVX2-MXFP8] Running AVX2 backend...")
    y_avx2 = run_backend("avx2", weights, expert_num, top_k, hidden_size, intermediate_size,
                         args.layer, qlen, args.cpu_threads, x, expert_ids, routing,
                         physical_to_logical)

    # ---- Torch reference ----
    print("[AVX2-MXFP8] Building torch reference (dequant+matmul)...")
    gate_bf16 = torch.stack([dequantize_mxfp8(weights["gate"][i], weights["gate_scale"][i], group_size)
                             for i in range(expert_num)])
    up_bf16 = torch.stack([dequantize_mxfp8(weights["up"][i], weights["up_scale"][i], group_size)
                           for i in range(expert_num)])
    down_bf16 = torch.stack([dequantize_mxfp8(weights["down"][i], weights["down_scale"][i], group_size)
                             for i in range(expert_num)])
    y_ref = reference_moe_m3(x, expert_ids, routing, gate_bf16, up_bf16, down_bf16,
                             alpha=1.702, limit=7.0)

    diff_ref = (y_avx2.float() - y_ref.float()).abs()
    ref_mag = y_ref.float().abs().mean() + 1e-12
    rel_ref = diff_ref.mean() / ref_mag
    print(f"[AVX2-MXFP8] vs ref:   mean abs={diff_ref.mean().item():.4e}  max abs={diff_ref.max().item():.4e}  rel={rel_ref.item()*100:.3f}%")
    pass_ref = rel_ref.item() < args.ref_threshold

    # ---- (optional) AMX vs AVX2 equivalence ----
    pass_eq = True
    if args.compare_amx:
        print("[AVX2-MXFP8] Running AMX backend for equivalence check...")
        y_amx = run_backend("amx", weights, expert_num, top_k, hidden_size, intermediate_size,
                            args.layer, qlen, args.cpu_threads, x, expert_ids, routing,
                            physical_to_logical)
        diff_eq = (y_amx.float() - y_avx2.float()).abs()
        amx_mag = y_amx.float().abs().mean() + 1e-12
        rel_eq = diff_eq.mean() / amx_mag
        print(f"[AVX2-MXFP8] vs AMX:  mean abs={diff_eq.mean().item():.4e}  max abs={diff_eq.max().item():.4e}  rel={rel_eq.item()*100:.3f}%")
        pass_eq = rel_eq.item() < args.equiv_threshold

    print(f"[AVX2-MXFP8] vs ref:  {'PASS' if pass_ref else 'FAIL'} (rel < {args.ref_threshold*100:.1f}%)")
    if args.compare_amx:
        print(f"[AVX2-MXFP8] vs AMX:  {'PASS' if pass_eq else 'FAIL'} (rel < {args.equiv_threshold*100:.1f}%)")
    return 0 if (pass_ref and pass_eq) else 1


if __name__ == "__main__":
    sys.exit(main())