From 943cc4daeb1501d2eb3fc4d711c7c87767905359 Mon Sep 17 00:00:00 2001 From: Benjamin <159887351+yyj6666667@users.noreply.github.com> Date: Thu, 18 Jun 2026 16:18:49 +0800 Subject: [PATCH] [feat] MXFP8 MoE support (#2041) Add MXFP8 kernel for Minimax-M3-MXFP8 Day 0 --- kt-kernel/examples/test_mxfp8_moe_avx2.py | 204 ++++ kt-kernel/examples/test_mxfp8_moe_m3.py | 176 ++++ kt-kernel/ext_bindings.cpp | 5 + kt-kernel/operators/amx/la/amx.hpp | 28 + kt-kernel/operators/amx/moe_base.hpp | 4 +- kt-kernel/operators/amx/mxfp8-moe.hpp | 952 +++++++++++++++++++ kt-kernel/operators/avx2/avx2_bf16_utils.hpp | 41 + kt-kernel/operators/avx2/moe_base.hpp | 23 +- kt-kernel/operators/avx2/mxfp8-moe.hpp | 740 ++++++++++++++ kt-kernel/operators/common.hpp | 5 + kt-kernel/python/experts.py | 15 +- kt-kernel/python/utils/amx.py | 105 +- kt-kernel/python/utils/loader.py | 79 ++ 13 files changed, 2355 insertions(+), 22 deletions(-) create mode 100644 kt-kernel/examples/test_mxfp8_moe_avx2.py create mode 100644 kt-kernel/examples/test_mxfp8_moe_m3.py create mode 100644 kt-kernel/operators/amx/mxfp8-moe.hpp create mode 100644 kt-kernel/operators/avx2/mxfp8-moe.hpp diff --git a/kt-kernel/examples/test_mxfp8_moe_avx2.py b/kt-kernel/examples/test_mxfp8_moe_avx2.py new file mode 100644 index 00000000..4dad97ec --- /dev/null +++ b/kt-kernel/examples/test_mxfp8_moe_avx2.py @@ -0,0 +1,204 @@ +"""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()) diff --git a/kt-kernel/examples/test_mxfp8_moe_m3.py b/kt-kernel/examples/test_mxfp8_moe_m3.py new file mode 100644 index 00000000..2bd59e91 --- /dev/null +++ b/kt-kernel/examples/test_mxfp8_moe_m3.py @@ -0,0 +1,176 @@ +"""End-to-end MXFP8 MoE validation against the native MiniMax-M3-Preview ckpt. + +Loads layer-`LAYER_ID` experts via :class:`MXFP8SafeTensorLoader`, runs the AMX +MXFP8 backend with swigluoai activation, and compares against a torch reference +that dequantizes FP8 E4M3fn weights with ue8m0 group scales. + +Usage: + python test_mxfp8_moe_m3.py --weight-path /mnt/data/models/Minimax-M3-preview [--layer 3] +""" + +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 ---- + +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: torch.Tensor, gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor, + alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: + """Single-expert MLP with swigluoai activation.""" + 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: torch.Tensor, expert_ids: torch.Tensor, weights: torch.Tensor, + gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor, + alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: + """Full MoE forward: route tokens to experts, compute, weighted sum.""" + 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) + + +# ---- Main ---- + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="MXFP8 MoE E2E test for MiniMax M3") + p.add_argument("--weight-path", required=True, help="Path to Minimax-M3-preview safetensors directory.") + p.add_argument("--layer", type=int, default=3, help="Layer index to validate (default: 3, first MoE layer).") + p.add_argument("--qlen", type=int, default=1, help="Number of tokens to test.") + p.add_argument("--top-k", type=int, default=4, help="num_experts_per_tok (M3 default 4).") + p.add_argument("--cpu-threads", type=int, default=32) + p.add_argument("--max-experts", type=int, default=0, help="Cap number of experts loaded (0 = all).") + return p.parse_args() + + +def main() -> int: + args = parse_args() + torch.manual_seed(42) + + print(f"[M3-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 + print(f"[M3-MXFP8] expert_num={expert_num}") + + gate0 = weights["gate"][0] + down0 = weights["down"][0] + intermediate_size = gate0.shape[0] # w1: [I, K] + hidden_size = gate0.shape[1] # FP8: 1 byte/element, no /2 + assert down0.shape == (hidden_size, intermediate_size), f"unexpected down shape {down0.shape}" + + group_size = hidden_size // weights["gate_scale"][0].shape[1] + print(f"[M3-MXFP8] 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() + + # ---- C++ MXFP8 forward ---- + cpu_infer = kt_kernel_ext.CPUInfer(args.cpu_threads) + cfg = kt_kernel_ext.moe.MOEConfig(expert_num, args.top_k, hidden_size, intermediate_size, 0) + cfg.layer_idx = args.layer + cfg.max_len = max(args.qlen, 1) + cfg.pool = cpu_infer.backend_ + cfg.quant_config.bits = 8 + cfg.quant_config.group_size = group_size + 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 = kt_kernel_ext.moe.AMXMXFP8_KGroup_MOE(cfg) + cpu_infer.submit(moe.load_weights_task(physical_to_logical.data_ptr())) + cpu_infer.sync() + print("[M3-MXFP8] C++ weights loaded") + + qlen = args.qlen + top_k = args.top_k + bsz = torch.tensor([qlen], dtype=torch.int32) + 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) # normalize weights + x = (torch.randn((qlen, hidden_size), dtype=torch.bfloat16) * 0.01).contiguous() + y_cpp = 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_cpp.data_ptr(), False, + ) + ) + cpu_infer.sync() + print("[M3-MXFP8] C++ forward done") + + # ---- Torch reference ---- + print("[M3-MXFP8] Building torch reference (dequantizing all loaded experts)...") + 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) + + # ---- Compare ---- + diff = (y_cpp.float() - y_ref.float()).abs() + ref_mag = y_ref.float().abs().mean() + 1e-12 + rel = diff.mean() / ref_mag + print(f"[M3-MXFP8] mean abs diff = {diff.mean().item():.4e}") + print(f"[M3-MXFP8] max abs diff = {diff.max().item():.4e}") + print(f"[M3-MXFP8] ref mean mag = {ref_mag.item():.4e}") + print(f"[M3-MXFP8] rel mean diff = {rel.item()*100:.3f}%") + print(f"[M3-MXFP8] cpp[:8] = {y_cpp.flatten()[:8].tolist()}") + print(f"[M3-MXFP8] ref[:8] = {y_ref.flatten()[:8].tolist()}") + + passed = rel.item() < 0.10 + print(f"[M3-MXFP8] {'PASS' if passed else 'FAIL'} (threshold: 10%)") + return 0 if passed else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 81ca2577..87074e41 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -45,6 +45,7 @@ static const bool _is_plain_ = false; #include "operators/amx/awq-moe.hpp" #include "operators/amx/bf16-moe.hpp" // Native BF16 MoE using CRTP pattern, with fallback for AVX512F #include "operators/amx/fp4-moe.hpp" // MXFP4 MoE: FP4 E2M1 weights × BF16 activations +#include "operators/amx/mxfp8-moe.hpp" // MXFP8 MoE: FP8 E4M3fn weights × BF16 activations (MiniMax M3) #include "operators/amx/fp8-moe.hpp" // FP8 MoE requires AVX512 BF16 support, with fallback for AVX512F+BW #include "operators/amx/fp8-perchannel-moe.hpp" // FP8 Per-Channel MoE for GLM-4.7-FP8 #include "operators/amx/k2-moe.hpp" @@ -60,6 +61,7 @@ static const bool _is_plain_ = false; #include "operators/avx2/gptq_int4-moe.hpp" #include "operators/avx2/gptq_int4_avxvnni-moe.hpp" #include "operators/avx2/mxfp4-moe.hpp" +#include "operators/avx2/mxfp8-moe.hpp" #include "operators/avx2/rawint4-moe.hpp" #include "operators/avx2/rawint4_avxvnni-moe.hpp" #endif @@ -760,6 +762,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .def_readwrite("max_cache_depth", &GeneralMOEConfig::max_cache_depth) // V4-Flash 2604B SwiGLU clamp limit (0.0 = disabled). See common.hpp. .def_readwrite("swiglu_limit", &GeneralMOEConfig::swiglu_limit) + .def_readwrite("swiglu_alpha", &GeneralMOEConfig::swiglu_alpha) ; @@ -793,6 +796,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { bind_moe_module>(moe_module, "AMXFP8_MOE"); bind_moe_module>(moe_module, "AMXFP8PerChannel_MOE"); bind_moe_module>(moe_module, "AMXFP4_KGroup_MOE"); + bind_moe_module>(moe_module, "AMXMXFP8_KGroup_MOE"); #endif #if defined(__AVX512BF16__) // SFT MoE with LoRA support (BF16, INT8, INT4, AWQ, K2) @@ -823,6 +827,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { bind_moe_module>(moe_module, "AVX2GPTQInt4_MOE"); bind_moe_module>(moe_module, "AVX2RawInt4_MOE"); bind_moe_module>(moe_module, "AVX2MXFP4_MOE"); + bind_moe_module>(moe_module, "AVX2MXFP8_MOE"); bind_moe_module>(moe_module, "AVXVNNI256GPTQInt4_MOE"); bind_moe_module>(moe_module, diff --git a/kt-kernel/operators/amx/la/amx.hpp b/kt-kernel/operators/amx/la/amx.hpp index 59b16e4c..3fe82edf 100644 --- a/kt-kernel/operators/amx/la/amx.hpp +++ b/kt-kernel/operators/amx/la/amx.hpp @@ -72,6 +72,34 @@ static inline __m512 act_fn(__m512 gate_val, __m512 up_val, float swiglu_limit = return _mm512_mul_ps(act_val, up_val); } +// MiniMax M3 "swigluoai" activation + DeepSeek V4 "silu" unified entry point. +// alpha > 0 → swigluoai: gate * sigmoid(gate * alpha) * (up + 1), symmetric clamp +// alpha == 0 → falls back to standard silu path above +static inline __m512 act_fn(__m512 gate_val, __m512 up_val, float swiglu_limit, float swiglu_alpha) { + if (swiglu_alpha > 0.0f) { + // --- MiniMax M3 swigluoai path --- + // Symmetric clamp on both gate and up (differs from V4 which clamps gate one-sided) + if (swiglu_limit > 0.0f) { + const __m512 pos_lim = _mm512_set1_ps(swiglu_limit); + const __m512 neg_lim = _mm512_set1_ps(-swiglu_limit); + gate_val = _mm512_min_ps(gate_val, pos_lim); + gate_val = _mm512_max_ps(gate_val, neg_lim); + up_val = _mm512_min_ps(up_val, pos_lim); + up_val = _mm512_max_ps(up_val, neg_lim); + } + // sigmoid(gate * alpha) = 1 / (1 + exp(-gate * alpha)) + __m512 neg_ga = _mm512_mul_ps(gate_val, _mm512_set1_ps(-swiglu_alpha)); + neg_ga = _mm512_min_ps(neg_ga, _mm512_set1_ps(88.0f)); + __m512 exp_neg = exp_avx512(neg_ga); + __m512 sigmoid_val = _mm512_div_ps(_mm512_set1_ps(1.0f), + _mm512_add_ps(_mm512_set1_ps(1.0f), exp_neg)); + // gate * sigmoid(gate * alpha) * (up + 1) + __m512 up_plus_1 = _mm512_add_ps(up_val, _mm512_set1_ps(1.0f)); + return _mm512_mul_ps(_mm512_mul_ps(gate_val, sigmoid_val), up_plus_1); + } + return act_fn(gate_val, up_val, swiglu_limit); +} + #define AMX_DISPATCH_QTYPES(QA, QB, ...) \ [&] { \ switch (QB) { \ diff --git a/kt-kernel/operators/amx/moe_base.hpp b/kt-kernel/operators/amx/moe_base.hpp index e56e83af..48c4da9e 100644 --- a/kt-kernel/operators/amx/moe_base.hpp +++ b/kt-kernel/operators/amx/moe_base.hpp @@ -685,8 +685,8 @@ class AMX_MOE_BASE { __m512 gate_val0, gate_val1, up_val0, up_val1; avx512_32xbf16_to_32xfp32((__m512i*)(gate_output_ptr + j), &gate_val0, &gate_val1); avx512_32xbf16_to_32xfp32((__m512i*)(up_output_ptr + j), &up_val0, &up_val1); - __m512 result0 = amx::act_fn(gate_val0, up_val0, config_.swiglu_limit); - __m512 result1 = amx::act_fn(gate_val1, up_val1, config_.swiglu_limit); + __m512 result0 = amx::act_fn(gate_val0, up_val0, config_.swiglu_limit, config_.swiglu_alpha); + __m512 result1 = amx::act_fn(gate_val1, up_val1, config_.swiglu_limit, config_.swiglu_alpha); avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i*)(gate_output_ptr + j)); } } diff --git a/kt-kernel/operators/amx/mxfp8-moe.hpp b/kt-kernel/operators/amx/mxfp8-moe.hpp new file mode 100644 index 00000000..3d571c51 --- /dev/null +++ b/kt-kernel/operators/amx/mxfp8-moe.hpp @@ -0,0 +1,952 @@ +/** + * @Description : MXFP8 MoE operator — FP8 E4M3fn weights × BF16 activations + * @Author : yyj and Claude + * @Date : 2026-06-02 + * @Version : 0.1.0 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + * + * Serves MiniMax M3 Preview (MXFP8 quantized checkpoint). + * Based on fp4-moe.hpp (MXFP4). Key differences from MXFP4: + * Weight: FP8 E4M3fn (1 byte/element, no nibble packing) + * Act: BF16 direct (BufferABF16Impl, same as MXFP4) + * Decode: Reuse GemmKernel224FP8 VBMI LUT tables for FP8→BF16 + * Dot prod: _mm512_dpbf16_ps (BF16×BF16→FP32, same as MXFP4) + * Scale: ue8m0 per-group (group_size=32), converted to FP32 on load + * + * Activation: "swigluoai" — gate * sigmoid(gate * alpha) * (up + 1) + * Handled by act_fn(gate, up, swiglu_limit, swiglu_alpha) in la/amx.hpp. + **/ +#ifndef CPUINFER_OPERATOR_AMX_MXFP8_MOE_H +#define CPUINFER_OPERATOR_AMX_MXFP8_MOE_H + +#include "la/amx_raw_buffers.hpp" // BufferABF16Impl +#include "la/amx_raw_kernels.hpp" // GemmKernel224FP8 LUT tables +#include "moe_base.hpp" + +namespace amx { + +// ============================================================================ +// BufferBMXFP8KGroupImpl — FP8 E4M3fn weight buffer with per-group ue8m0 scale +// +// Memory layout: [FP8 weights (n*k bytes)] [FP32 scales (n * k/gs floats)] +// Weights are row-major, 1 byte per element (no nibble packing). +// Scales are row-major, one FP32 value per group of `k_group_size` elements. +// The ue8m0→FP32 conversion happens during load_weights, not here. +// ============================================================================ +template +struct BufferBMXFP8KGroupImpl { + using dt = uint8_t; + uint8_t* b; // FP8 E4M3fn weights, row-major [n, k], 1 byte per element + float* d; // FP32 per-group scales, row-major [n, k/k_group_size] + int n, k, k_group_size, k_group_count; + + static constexpr int N_STEP = K::N_STEP; + static constexpr int K_STEP = K::K_STEP; + static constexpr bool SCALE = true; + + static size_t required_size(int n, int k, int k_group_size) { + // FP8: 1 byte per element (vs FP4's 0.5 byte) + return static_cast(n) * k + sizeof(float) * n * (k / k_group_size); + } + + BufferBMXFP8KGroupImpl(int n, int k, int k_group_size, void* ptr) + : n(n), k(k), k_group_size(k_group_size) { + assert(reinterpret_cast(ptr) % 64 == 0); + if (n % N_STEP || k % K_STEP || k % k_group_size) { + printf("BufferBMXFP8KGroupImpl: n=%d k=%d N_STEP=%d K_STEP=%d gs=%d\n", + n, k, N_STEP, K_STEP, k_group_size); + throw std::runtime_error("n or k not aligned to N_STEP/K_STEP/group_size"); + } + k_group_count = k / k_group_size; + b = reinterpret_cast(ptr); + // Scale region starts after weight data + d = reinterpret_cast(b + static_cast(n) * k); + } + + // Copy raw FP8 bytes from checkpoint. No repacking needed (already 1 byte/element). + void from_raw_mat(const uint8_t* proj, int ith, int nth) { + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + if (n_start >= n_end) return; + const size_t row_bytes = static_cast(k); // 1 byte per element + const size_t rows = static_cast(n_end - n_start); + std::memcpy(b + n_start * row_bytes, proj + n_start * row_bytes, rows * row_bytes); + } + + // Pointer to FP8 data at row n_begin, column k_begin + uint8_t* get_submat(int /*n*/, int /*k*/, int n_begin, int k_begin) { + return b + static_cast(n_begin) * k + k_begin; + } + + // Pointer to FP32 scale for row n_begin at k-group starting at k_begin + float* get_scale(int /*n*/, int n_begin, int /*k*/, int k_begin) { + return d + static_cast(n_begin) * k_group_count + k_begin / k_group_size; + } + + static std::pair split_range_n(int n, int ith, int nth) { + int n_per_thread = (n + nth - 1) / nth; + n_per_thread = (n_per_thread + N_STEP - 1) / N_STEP * N_STEP; + int n_start = std::min(ith * n_per_thread, n); + int n_end = std::min(n_start + n_per_thread, n); + return {n_start, n_end}; + } +}; + +// ============================================================================ +// GemmKernel224MXFP8SmallKGroup — FP8 E4M3fn × BF16 → FP32 (AVX512 + VBMI) +// +// Per-group scale (group_size=32): each 32 FP8 elements share one ue8m0 scale. +// Inner loop: load 32 FP8 bytes → VBMI LUT decode → 32 BF16 → dpbf16_ps → fmadd(scale). +// ============================================================================ +struct GemmKernel224MXFP8SmallKGroup { + using dt = uint8_t; + using output_t = float; + static constexpr double ELEMENT_SIZE = 1.0; // 1 byte per FP8 value + + static const int M_STEP = 1; + static const int N_STEP = 32; + static const int K_STEP = 32; // = group_size + + static inline const int N_BLOCK = 256; + static inline const int K_BLOCK = 6144; // M3 hidden=6144, fits in one pass + + static std::string name() { return "MXFP8_KGROUP"; } + static int recommended_nth(int n) { return (n + N_BLOCK - 1) / N_BLOCK; } + static std::pair split_range_n(int n, int ith, int nth) { + int n_start = N_BLOCK * ith; + int n_end = std::min(n, N_BLOCK * (ith + 1)); + return {n_start, n_end}; + } + static void config() {} + + // -------------------------------------------------------------------------- + // FP8 E4M3fn → BF16 decode for 32 values + // + // Reuses the VBMI LUT tables from GemmKernel224FP8 (amx_raw_kernels.hpp). + // The 4 tables (bf16_hi_0/1, bf16_lo_0/1) map the 7-bit FP8 magnitude + // (128 entries across two 64-byte tables) to BF16 hi/lo bytes. + // + // Strategy for 32 values (half of what fp8x64_to_bf16x64 handles): + // 1. Zero-extend 32 × uint8 → 32 × uint16 in a __m512i + // 2. Extract the original bytes into lower 32 bytes of a __m512i for LUT + // 3. VBMI permutex2var_epi8 for hi/lo byte lookup (same as fp8x64_to_bf16x64) + // 4. Combine hi (with sign) and lo into BF16 pairs + // + // The tricky part: fp8x64_to_bf16x64 uses unpacklo/hi which interleaves + // within 128-bit lanes, scattering 32 values across 2 registers. Instead, + // we do the lookup on the raw 32 bytes and manually zip lo+hi into uint16. + // -------------------------------------------------------------------------- + + // ActivationBF16: wraps a __m512bh (32 BF16 values = one k-group of activation) + struct ActivationBF16 { + __m512bh a; +#if !defined(__AVX512BF16__) + __m512 a_even; + __m512 a_odd; + inline static const __m512i odd_mask = _mm512_set1_epi32(0xFFFF0000); +#endif + __attribute__((always_inline)) ActivationBF16(__m512bh a_) : a(a_) { +#if !defined(__AVX512BF16__) + a_even = _mm512_castsi512_ps(_mm512_slli_epi32((__m512i)a_, 16)); + a_odd = _mm512_castsi512_ps(_mm512_and_si512((__m512i)a_, odd_mask)); +#endif + } + }; + + // DequantizedWeight: decodes 32 FP8 E4M3fn bytes → 32 BF16 values + struct DequantizedWeight { +#if defined(__AVX512BF16__) + __m512bh d; +#else + __m512 w_even; + __m512 w_odd; +#endif + + __attribute__((always_inline)) DequantizedWeight(__m256i fp8_32) { + // Pad 32 bytes into lower half of __m512i for VBMI LUT lookup + __m512i fp8_64 = _mm512_castsi256_si512(fp8_32); + // zeros in upper 32 bytes → LUT maps them to BF16(0), harmless + + // VBMI lookup: same tables as GemmKernel224FP8 + __m512i b_hi = _mm512_permutex2var_epi8( + GemmKernel224FP8::bf16_hi_0_mask(), fp8_64, GemmKernel224FP8::bf16_hi_1_mask()); + __m512i b_lo = _mm512_permutex2var_epi8( + GemmKernel224FP8::bf16_lo_0_mask(), fp8_64, GemmKernel224FP8::bf16_lo_1_mask()); + // Apply sign from original FP8 bytes + b_hi = _mm512_or_si512(_mm512_and_si512(GemmKernel224FP8::sign_mask(), fp8_64), b_hi); + + // Now b_lo[i] and b_hi[i] are the low/high bytes of BF16 for FP8 byte i. + // We need: bf16[i] = (b_hi[i] << 8) | b_lo[i] as a uint16 at position i. + // + // unpacklo/hi interleave within 128-bit lanes which scrambles the order. + // Instead, use the 32 bytes we care about (positions 0..31) directly: + // - Extract lower 256 bits of b_lo and b_hi (our 32 valid bytes) + // - Interleave with _mm256_unpacklo/hi → 4 chunks of 8 BF16 each + // - Assemble into one __m512i of 32 BF16 values + + __m256i lo_256 = _mm512_castsi512_si256(b_lo); // b_lo bytes 0..31 + __m256i hi_256 = _mm512_castsi512_si256(b_hi); // b_hi bytes 0..31 + + // _mm256_unpacklo/hi_epi8 interleaves within each 128-bit LANE (not whole 256 bits) + // Lane 0 (bytes 0..15): unpacklo → [lo[0],hi[0], lo[1],hi[1], ..., lo[7],hi[7]] = 8 BF16 + // unpackhi → [lo[8],hi[8], ..., lo[15],hi[15]] = 8 BF16 + // Lane 1 (bytes 16..31): unpacklo → [lo[16],hi[16], ..., lo[23],hi[23]] = 8 BF16 + // unpackhi → [lo[24],hi[24], ..., lo[31],hi[31]] = 8 BF16 + __m256i bf16_0_2 = _mm256_unpacklo_epi8(lo_256, hi_256); // lane0: BF16[0..7], lane1: BF16[16..23] + __m256i bf16_1_3 = _mm256_unpackhi_epi8(lo_256, hi_256); // lane0: BF16[8..15], lane1: BF16[24..31] + + // Reorder lanes to get BF16[0..31] in order: + // bf16_0_2 = [BF16[0..7] | BF16[16..23]] (128-bit lanes) + // bf16_1_3 = [BF16[8..15] | BF16[24..31]] + // Target: [BF16[0..7] | BF16[8..15] | BF16[16..23] | BF16[24..31]] + __m512i result = _mm512_inserti64x4(_mm512_castsi256_si512(bf16_0_2), bf16_1_3, 1); + // Now: [BF16[0..7], BF16[16..23], BF16[8..15], BF16[24..31]] + // Need to swap 128-bit lanes 1 and 2: + result = _mm512_permutexvar_epi64( + _mm512_setr_epi64(0, 1, 4, 5, 2, 3, 6, 7), result); + // Now: [BF16[0..7], BF16[8..15], BF16[16..23], BF16[24..31]] ✓ + +#if defined(__AVX512BF16__) + d = (__m512bh)result; +#else + w_even = _mm512_castsi512_ps(_mm512_slli_epi32(result, 16)); + w_odd = _mm512_castsi512_ps(_mm512_and_si512(result, _mm512_set1_epi32(0xFFFF0000))); +#endif + } + }; + + __attribute__((always_inline)) static inline __m512 mxfp8_dot_bf16( + const DequantizedWeight& w, const ActivationBF16& act) { +#if defined(__AVX512BF16__) + return _mm512_dpbf16_ps(_mm512_setzero_ps(), act.a, w.d); +#else + __m512 dot = _mm512_mul_ps(act.a_odd, w.w_odd); + return _mm512_fmadd_ps(act.a_even, w.w_even, dot); +#endif + } + + // Buffer type aliases + using BufferA = BufferABF16Impl; + using BufferB = BufferBMXFP8KGroupImpl; + using BufferC = BufferCReduceImpl; + + __attribute__((always_inline)) static inline void reduce4( + __m512 s0, __m512 s1, __m512 s2, __m512 s3, float* dst) { + dst[0] = _mm512_reduce_add_ps(s0); + dst[1] = _mm512_reduce_add_ps(s1); + dst[2] = _mm512_reduce_add_ps(s2); + dst[3] = _mm512_reduce_add_ps(s3); + } + + // -------------------------------------------------------------------------- + // mat-vec: M tokens × N output rows, 4-wide N unroll. + // Each k-group: load 32 FP8 bytes → decode → BF16 dot → fmadd(scale, dot, acc) + // -------------------------------------------------------------------------- + static void mxfp8_mat_vec_kgroup(int m, int n, int k, int k_group_size, + BufferA* ba, BufferB* bb, BufferC* bc, + int ith, int nth) { + auto [n_start, n_end] = split_range_n(n, ith, nth); + if (n_start >= n_end) return; + const int kg_count = k / 32; + + for (int m_idx = 0; m_idx < m; m_idx++) { + float* c_row = bc->get_submat(m, n, m_idx, n_start); + __m512bh* a_row = (__m512bh*)ba->get_submat(m, k, m_idx, 0); + + int n_pos = n_start; + for (; n_pos + 4 <= n_end; n_pos += 4) { + // FP8 weights: 32 bytes per k-group per N-row (vs 16 bytes for FP4) + __m256i* w0 = (__m256i*)bb->get_submat(n, k, n_pos + 0, 0); + __m256i* w1 = (__m256i*)bb->get_submat(n, k, n_pos + 1, 0); + __m256i* w2 = (__m256i*)bb->get_submat(n, k, n_pos + 2, 0); + __m256i* w3 = (__m256i*)bb->get_submat(n, k, n_pos + 3, 0); + const float* s0 = bb->get_scale(n, n_pos + 0, k, 0); + const float* s1 = bb->get_scale(n, n_pos + 1, k, 0); + const float* s2 = bb->get_scale(n, n_pos + 2, k, 0); + const float* s3 = bb->get_scale(n, n_pos + 3, k, 0); + + __m512 acc0 = _mm512_setzero_ps(); + __m512 acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(); + __m512 acc3 = _mm512_setzero_ps(); + + for (int g = 0; g < kg_count; g++) { + const ActivationBF16 a(a_row[g]); + const DequantizedWeight d0(w0[g]); + const DequantizedWeight d1(w1[g]); + const DequantizedWeight d2(w2[g]); + const DequantizedWeight d3(w3[g]); + acc0 = _mm512_fmadd_ps(_mm512_set1_ps(s0[g]), mxfp8_dot_bf16(d0, a), acc0); + acc1 = _mm512_fmadd_ps(_mm512_set1_ps(s1[g]), mxfp8_dot_bf16(d1, a), acc1); + acc2 = _mm512_fmadd_ps(_mm512_set1_ps(s2[g]), mxfp8_dot_bf16(d2, a), acc2); + acc3 = _mm512_fmadd_ps(_mm512_set1_ps(s3[g]), mxfp8_dot_bf16(d3, a), acc3); + } + reduce4(acc0, acc1, acc2, acc3, c_row + (n_pos - n_start)); + } + // N tail: single-row fallback + for (; n_pos < n_end; n_pos++) { + __m256i* w = (__m256i*)bb->get_submat(n, k, n_pos, 0); + const float* s = bb->get_scale(n, n_pos, k, 0); + __m512 acc = _mm512_setzero_ps(); + for (int g = 0; g < kg_count; g++) { + const ActivationBF16 a(a_row[g]); + const DequantizedWeight d(w[g]); + acc = _mm512_fmadd_ps(_mm512_set1_ps(s[g]), mxfp8_dot_bf16(d, a), acc); + } + c_row[n_pos - n_start] = _mm512_reduce_add_ps(acc); + } + } + } + + // -------------------------------------------------------------------------- + // mat-mat: 4×4 register tile (MB=4 tokens, NB=4 N-rows → 16 accumulators). + // Weight decode shared across MB tokens → amortize VBMI cost by 4×. + // -------------------------------------------------------------------------- + static void mxfp8_mat_mat_kgroup(int m, int n, int k, int k_group_size, + BufferA* ba, BufferB* bb, BufferC* bc, + int ith, int nth) { + auto [n_start, n_end] = split_range_n(n, ith, nth); + if (n_start >= n_end) return; + const int kg_count = k / 32; + constexpr int MB = 4; + constexpr int NB = 4; + + int m_pos = 0; + for (; m_pos + MB <= m; m_pos += MB) { + __m512bh* a_rows[MB] = { + (__m512bh*)ba->get_submat(m, k, m_pos + 0, 0), + (__m512bh*)ba->get_submat(m, k, m_pos + 1, 0), + (__m512bh*)ba->get_submat(m, k, m_pos + 2, 0), + (__m512bh*)ba->get_submat(m, k, m_pos + 3, 0), + }; + + int n_pos = n_start; + for (; n_pos + NB <= n_end; n_pos += NB) { + __m256i* w0 = (__m256i*)bb->get_submat(n, k, n_pos + 0, 0); + __m256i* w1 = (__m256i*)bb->get_submat(n, k, n_pos + 1, 0); + __m256i* w2 = (__m256i*)bb->get_submat(n, k, n_pos + 2, 0); + __m256i* w3 = (__m256i*)bb->get_submat(n, k, n_pos + 3, 0); + const float* s0 = bb->get_scale(n, n_pos + 0, k, 0); + const float* s1 = bb->get_scale(n, n_pos + 1, k, 0); + const float* s2 = bb->get_scale(n, n_pos + 2, k, 0); + const float* s3 = bb->get_scale(n, n_pos + 3, k, 0); + + __m512 acc[MB][NB]; + for (int i = 0; i < MB; i++) + for (int j = 0; j < NB; j++) acc[i][j] = _mm512_setzero_ps(); + + for (int g = 0; g < kg_count; g++) { + const DequantizedWeight d0(w0[g]); + const DequantizedWeight d1(w1[g]); + const DequantizedWeight d2(w2[g]); + const DequantizedWeight d3(w3[g]); + const __m512 sv0 = _mm512_set1_ps(s0[g]); + const __m512 sv1 = _mm512_set1_ps(s1[g]); + const __m512 sv2 = _mm512_set1_ps(s2[g]); + const __m512 sv3 = _mm512_set1_ps(s3[g]); + +#define MXFP8_FMA_ROW(M_I) \ + do { \ + const ActivationBF16 a(a_rows[M_I][g]); \ + acc[M_I][0] = _mm512_fmadd_ps(sv0, mxfp8_dot_bf16(d0, a), acc[M_I][0]); \ + acc[M_I][1] = _mm512_fmadd_ps(sv1, mxfp8_dot_bf16(d1, a), acc[M_I][1]); \ + acc[M_I][2] = _mm512_fmadd_ps(sv2, mxfp8_dot_bf16(d2, a), acc[M_I][2]); \ + acc[M_I][3] = _mm512_fmadd_ps(sv3, mxfp8_dot_bf16(d3, a), acc[M_I][3]); \ + } while (0) + MXFP8_FMA_ROW(0); + MXFP8_FMA_ROW(1); + MXFP8_FMA_ROW(2); + MXFP8_FMA_ROW(3); +#undef MXFP8_FMA_ROW + } + for (int i = 0; i < MB; i++) { + float* c_row = bc->get_submat(m, n, m_pos + i, n_start); + reduce4(acc[i][0], acc[i][1], acc[i][2], acc[i][3], c_row + (n_pos - n_start)); + } + } + // N tail + for (; n_pos < n_end; n_pos++) { + __m256i* w = (__m256i*)bb->get_submat(n, k, n_pos, 0); + const float* s = bb->get_scale(n, n_pos, k, 0); + for (int i = 0; i < MB; i++) { + float* c_row = bc->get_submat(m, n, m_pos + i, n_start); + __m512 acc = _mm512_setzero_ps(); + for (int g = 0; g < kg_count; g++) { + const ActivationBF16 a(a_rows[i][g]); + const DequantizedWeight d(w[g]); + acc = _mm512_fmadd_ps(_mm512_set1_ps(s[g]), mxfp8_dot_bf16(d, a), acc); + } + c_row[n_pos - n_start] = _mm512_reduce_add_ps(acc); + } + } + } + // M tail: remaining tokens that don't fill a 4-token tile + for (int mi = m_pos; mi < m; mi++) { + float* c_row = bc->get_submat(m, n, mi, n_start); + __m512bh* a_row = (__m512bh*)ba->get_submat(m, k, mi, 0); + int n_pos = n_start; + for (; n_pos + 4 <= n_end; n_pos += 4) { + __m256i* w0 = (__m256i*)bb->get_submat(n, k, n_pos + 0, 0); + __m256i* w1 = (__m256i*)bb->get_submat(n, k, n_pos + 1, 0); + __m256i* w2 = (__m256i*)bb->get_submat(n, k, n_pos + 2, 0); + __m256i* w3 = (__m256i*)bb->get_submat(n, k, n_pos + 3, 0); + const float* s0 = bb->get_scale(n, n_pos + 0, k, 0); + const float* s1 = bb->get_scale(n, n_pos + 1, k, 0); + const float* s2 = bb->get_scale(n, n_pos + 2, k, 0); + const float* s3 = bb->get_scale(n, n_pos + 3, k, 0); + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(), + a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + for (int g = 0; g < kg_count; g++) { + const ActivationBF16 a(a_row[g]); + const DequantizedWeight d0(w0[g]); + const DequantizedWeight d1(w1[g]); + const DequantizedWeight d2(w2[g]); + const DequantizedWeight d3(w3[g]); + a0 = _mm512_fmadd_ps(_mm512_set1_ps(s0[g]), mxfp8_dot_bf16(d0, a), a0); + a1 = _mm512_fmadd_ps(_mm512_set1_ps(s1[g]), mxfp8_dot_bf16(d1, a), a1); + a2 = _mm512_fmadd_ps(_mm512_set1_ps(s2[g]), mxfp8_dot_bf16(d2, a), a2); + a3 = _mm512_fmadd_ps(_mm512_set1_ps(s3[g]), mxfp8_dot_bf16(d3, a), a3); + } + reduce4(a0, a1, a2, a3, c_row + (n_pos - n_start)); + } + for (; n_pos < n_end; n_pos++) { + __m256i* w = (__m256i*)bb->get_submat(n, k, n_pos, 0); + const float* s = bb->get_scale(n, n_pos, k, 0); + __m512 acc = _mm512_setzero_ps(); + for (int g = 0; g < kg_count; g++) { + const ActivationBF16 a(a_row[g]); + const DequantizedWeight d(w[g]); + acc = _mm512_fmadd_ps(_mm512_set1_ps(s[g]), mxfp8_dot_bf16(d, a), acc); + } + c_row[n_pos - n_start] = _mm512_reduce_add_ps(acc); + } + } + } +}; + +// Dispatch functions +inline void mxfp8_vec_mul_kgroup(int m, int n, int k, int k_group_size, + std::shared_ptr ba, + std::shared_ptr bb, + std::shared_ptr bc, + int ith, int nth) { + GemmKernel224MXFP8SmallKGroup::mxfp8_mat_vec_kgroup(m, n, k, k_group_size, ba.get(), bb.get(), bc.get(), ith, nth); +} + +inline void mxfp8_mat_mul_kgroup(int m, int n, int k, int k_group_size, + std::shared_ptr ba, + std::shared_ptr bb, + std::shared_ptr bc, + int ith, int nth) { + GemmKernel224MXFP8SmallKGroup::mxfp8_mat_mat_kgroup(m, n, k, k_group_size, ba.get(), bb.get(), bc.get(), ith, nth); +} + +} // namespace amx + +// ============================================================================ +// AMX_MXFP8_MOE_TP — CRTP class for MXFP8 MoE (MiniMax M3) +// ============================================================================ +template +class AMX_MXFP8_MOE_TP : public AMX_MOE_BASE> { + using Base = AMX_MOE_BASE>; + using Base::config_; + using Base::down_ba_; + using Base::down_bb_; + using Base::down_bc_; + using Base::gate_bb_; + using Base::gate_bc_; + using Base::gate_up_ba_; + using Base::m_local_num_; + using Base::tp_part_idx; + using Base::up_bb_; + using Base::up_bc_; + + public: + using typename Base::input_t; + using typename Base::output_t; + + AMX_MXFP8_MOE_TP() = default; + AMX_MXFP8_MOE_TP(GeneralMOEConfig config, int tp_part_idx_ = 0) : Base(config, tp_part_idx_) {} + + void derived_init() { + auto& quant_config = config_.quant_config; + if (quant_config.group_size != 32 || quant_config.zero_point) { + throw std::runtime_error("MXFP8 MoE requires group_size == 32 and no zero_point"); + } + if (config_.hidden_size % quant_config.group_size != 0 || + config_.intermediate_size % quant_config.group_size != 0) { + throw std::runtime_error( + "MXFP8 MoE: hidden_size and intermediate_size must be divisible by group_size"); + } + printf("Creating AMX_MXFP8_MOE_TP %d at numa %d\n", tp_part_idx, numa_node_of_cpu(sched_getcpu())); + } + + ~AMX_MXFP8_MOE_TP() = default; + + size_t buffer_a_required_size_impl(size_t m, size_t k) const { return T::BufferA::required_size(m, k); } + size_t buffer_b_required_size_impl(size_t n, size_t k) const { + return T::BufferB::required_size(n, k, config_.quant_config.group_size); + } + size_t buffer_c_required_size_impl(size_t m, size_t n) const { return T::BufferC::required_size(m, n); } + + std::shared_ptr make_buffer_a_impl(size_t m, size_t k, void* data) const { + return std::make_shared(m, k, data); + } + std::shared_ptr make_buffer_b_impl(size_t n, size_t k, void* data) const { + return std::make_shared(n, k, config_.quant_config.group_size, data); + } + std::shared_ptr make_buffer_c_impl(size_t m, size_t n, void* data) const { + return std::make_shared(m, n, data); + } + + void do_gate_up_gemm(bool do_up, int expert_idx, int ith, int nth, int qlen) { + auto& group_size = config_.quant_config.group_size; + int m = m_local_num_[expert_idx]; + auto& ba = gate_up_ba_[expert_idx]; + auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; + auto& bc = do_up ? up_bc_[expert_idx] : gate_bc_[expert_idx]; + + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mxfp8_mat_mul_kgroup(m, config_.intermediate_size, config_.hidden_size, group_size, ba, bb, bc, ith, nth); + } else { + amx::mxfp8_vec_mul_kgroup(m, config_.intermediate_size, config_.hidden_size, group_size, ba, bb, bc, ith, nth); + } + } + + void do_down_gemm(int expert_idx, int ith, int nth, int qlen) { + auto& group_size = config_.quant_config.group_size; + int m = m_local_num_[expert_idx]; + + if (qlen > 4 * config_.expert_num / config_.num_experts_per_tok) { + amx::mxfp8_mat_mul_kgroup(m, config_.hidden_size, config_.intermediate_size, group_size, + down_ba_[expert_idx], down_bb_[expert_idx], down_bc_[expert_idx], ith, nth); + } else { + amx::mxfp8_vec_mul_kgroup(m, config_.hidden_size, config_.intermediate_size, group_size, + down_ba_[expert_idx], down_bb_[expert_idx], down_bc_[expert_idx], ith, nth); + } + } + + // -------------------------------------------------------------------------- + // ue8m0 → FP32 scale conversion (vectorized) + // + // ue8m0 is an unsigned 8-bit exponent: value = 2^(byte - 127). + // IEEE754 FP32 bit layout: [sign:1][exp:8][mantissa:23]. + // Setting exp field = byte and mantissa = 0 gives 2^(byte - 127) directly. + // Edge case: byte=0 → 2^(-127) ≈ 5.9e-39, but (0 << 23) = 0.0f. + // In practice M3 weights don't have ue8m0=0 scales; we accept 0.0f here. + // -------------------------------------------------------------------------- + static inline void convert_ue8m0_to_fp32(float* __restrict dst, + const uint8_t* __restrict src, + size_t count) { + size_t i = 0; + for (; i + 8 <= count; i += 8) { + __m128i bytes = _mm_loadl_epi64((__m128i const*)(src + i)); + __m256i dwords = _mm256_cvtepu8_epi32(bytes); + __m256 floats = _mm256_castsi256_ps(_mm256_slli_epi32(dwords, 23)); + _mm256_storeu_ps(dst + i, floats); + } + for (; i < count; i++) { + uint32_t bits = static_cast(src[i]) << 23; + std::memcpy(dst + i, &bits, sizeof(float)); + } + } + + void load_weights() { + auto& quant_config = config_.quant_config; + const uint64_t* physical_to_logical_map = (const uint64_t*)config_.physical_to_logical_map; + auto pool = config_.pool->get_subpool(tp_part_idx); + + if (quant_config.group_size == 0 || quant_config.zero_point) + throw std::runtime_error("MXFP8 MoE requires group_size > 0 and no zero_point"); + if (config_.gate_scale == nullptr) + throw std::runtime_error("MXFP8 MoE requires native MXFP8 weights with ue8m0 scales"); + + // --- Load FP8 weights (1 byte per element, no nibble packing) --- + int nth = T::recommended_nth(config_.intermediate_size); + pool->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [this, nth, physical_to_logical_map](int task_id) { + uint64_t expert_idx = task_id / nth; + uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + int ith = task_id % nth; + // FP8: no >> 1 (1 byte/element, not nibble-packed) + size_t weight_offset = logical_expert_id * config_.intermediate_size * config_.hidden_size; + gate_bb_[expert_idx]->from_raw_mat((uint8_t*)config_.gate_proj + weight_offset, ith, nth); + up_bb_[expert_idx]->from_raw_mat((uint8_t*)config_.up_proj + weight_offset, ith, nth); + }, + nullptr); + + nth = T::recommended_nth(config_.hidden_size); + pool->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [this, nth, physical_to_logical_map](int task_id) { + uint64_t expert_idx = task_id / nth; + uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + int ith = task_id % nth; + size_t weight_offset = logical_expert_id * config_.hidden_size * config_.intermediate_size; + down_bb_[expert_idx]->from_raw_mat((uint8_t*)config_.down_proj + weight_offset, ith, nth); + }, + nullptr); + + // --- Convert ue8m0 scales to FP32 --- + // M3 scale dtype is uint8 (ue8m0), not bf16 like V4. Use bit-shift conversion. + pool->do_work_stealing_job( + config_.expert_num, nullptr, + [this, physical_to_logical_map](int task_id) { + uint64_t expert_idx = task_id; + uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + size_t scale_count = + (static_cast(config_.intermediate_size) * config_.hidden_size) / config_.quant_config.group_size; + // gate and up scales: [intermediate_size, hidden_size / group_size] + convert_ue8m0_to_fp32(gate_bb_[expert_idx]->d, + (const uint8_t*)config_.gate_scale + logical_expert_id * scale_count, scale_count); + convert_ue8m0_to_fp32(up_bb_[expert_idx]->d, + (const uint8_t*)config_.up_scale + logical_expert_id * scale_count, scale_count); + // down scale: [hidden_size, intermediate_size / group_size] — same total count + convert_ue8m0_to_fp32(down_bb_[expert_idx]->d, + (const uint8_t*)config_.down_scale + logical_expert_id * scale_count, scale_count); + }, + nullptr); + } + + static inline void fast_memcpy(void* __restrict dst, const void* __restrict src, size_t bytes) { + uint8_t* d = (uint8_t*)dst; + const uint8_t* s = (const uint8_t*)src; + size_t chunks = bytes / 64; + for (size_t i = 0; i < chunks; i++) { + __m512i data = _mm512_loadu_si512((__m512i*)s); + _mm512_storeu_si512((__m512i*)d, data); + d += 64; + s += 64; + } + if (bytes -= chunks * 64) std::memcpy(d, s, bytes); + } + + static inline void fast_fp32_to_bf16(ggml_bf16_t* __restrict dst, const float* __restrict src, size_t count) { + size_t i = 0; + for (; i + 32 <= count; i += 32) { + __m512 v0 = _mm512_loadu_ps(src + i); + __m512 v1 = _mm512_loadu_ps(src + i + 16); + __m512i i0 = _mm512_srli_epi32(_mm512_castps_si512(v0), 16); + __m512i i1 = _mm512_srli_epi32(_mm512_castps_si512(v1), 16); + __m512i packed = _mm512_packus_epi32(i0, i1); + __m512i permuted = _mm512_permutexvar_epi64(_mm512_set_epi64(7, 5, 3, 1, 6, 4, 2, 0), packed); + _mm512_storeu_si512((__m512i*)(dst + i), permuted); + } + for (; i < count; i++) dst[i] = ggml_fp32_to_bf16(src[i]); + } + + // FP32 -> ue8m0 (uint8) — AVX-512 sibling of avx2::fast_fp32_to_ue8m0. + // Used by write_weights_to_buffer for layerwise prefill: CPU stores FP32 + // scales in BufferB.d for fast FMA in inner loop, but GPU MXFP8 buffer + // expects ue8m0 uint8 (1 byte/scale). Extract bits 23-30 of each FP32 + // (the exponent field) and store as uint8. Bit-exact iff the FP32 came + // from convert_ue8m0_to_fp32 (mantissa=0), which is the case for + // layerwise prefill since CPU forward never modifies BufferB.d. + static inline void fast_fp32_to_ue8m0(uint8_t* __restrict dst, + const float* __restrict src, + size_t count) { + size_t i = 0; + for (; i + 16 <= count; i += 16) { + __m512 v = _mm512_loadu_ps(src + i); + __m512i bits = _mm512_castps_si512(v); + // Extract bits 23-30 (FP32 exponent) -> 16 x uint32 in [0,255] + __m512i shifted = _mm512_srli_epi32(bits, 23); + // Truncate 32 -> 8 (AVX512-BW): values already in [0,255] + __m128i p8 = _mm512_cvtepi32_epi8(shifted); + _mm_storeu_si128((__m128i*)(dst + i), p8); + } + for (; i < count; i++) { + uint32_t b; + std::memcpy(&b, &src[i], sizeof(uint32_t)); + dst[i] = (uint8_t)((b >> 23) & 0xFF); + } + } + + // write_weights_to_buffer: copies CPU expert weights to GPU pinned buffer + // for the GPU prefill fallback (sglang/kt_ep_wrapper.py:_prepare_weight_mxfp8). + // + // Unified per-row scatter — works for any (cpu_tp_count, gpu_tp_count) + // relationship. Mirrors avx2/mxfp8-moe.hpp and avx2/fp8-moe.hpp:347-451. + // + // Replaces the earlier `cpu_tp_count >= gpu_tp_count` direct-write branch + // that (a) silently no-op'd for the typical 2-NUMA x tp>=4 deployment and + // (b) wrote bf16 (2 bytes/scale) into a uint8 (1 byte/scale) GPU buffer, + // corrupting adjacent expert slots whenever it did execute. + // + // GPU MXFP8 buffer layout (set by Fp8MoEMethod.create_weights when + // use_mxfp8=True in kt-sglang/fp8.py:872-893): + // w13_weight_scale_inv: (E, 2*intermediate, hidden/group_size) uint8 ue8m0 + // w2_weight_scale_inv: (E, hidden, intermediate/group_size) uint8 ue8m0 + // CPU storage (gate_bb_[e]->d / up_bb_[e]->d / down_bb_[e]->d) is fp32 + // (pre-decoded ue8m0 from disk, for fast FMA in CPU forward); we re-encode + // to ue8m0 via fast_fp32_to_ue8m0 on write. + void write_weights_to_buffer(int gpu_tp_count, [[maybe_unused]] int cpu_tp_count, int expert_id, + const GeneralMOEConfig& full_config, + const std::vector& w13_weight_ptrs, + const std::vector& w13_scale_ptrs, + const std::vector& w2_weight_ptrs, + const std::vector& w2_scale_ptrs) const { + auto& config = config_; + auto pool = config.pool->get_subpool(tp_part_idx); + const int group_size = config.quant_config.group_size; + + // ========= W13 (gate+up): Shape [intermediate, hidden], split by N only ========= + const int cpu_n_w13 = config.intermediate_size; + const int cpu_k_w13 = config.hidden_size; + const int gpu_n_w13 = full_config.intermediate_size / gpu_tp_count; + const int gpu_k_w13 = full_config.hidden_size; + const int global_n_offset_w13 = tp_part_idx * cpu_n_w13; + const size_t gpu_w13_weight_per_mat = (size_t)gpu_n_w13 * gpu_k_w13; + const int scales_per_row_w13 = cpu_k_w13 / group_size; + const size_t gpu_w13_scale_per_mat = (size_t)gpu_n_w13 * scales_per_row_w13; + + // ========= W2 (down): Shape [hidden, intermediate], split by K ========= + const int cpu_n_w2 = config.hidden_size; + const int cpu_k_w2 = config.intermediate_size; + const int gpu_k_w2 = full_config.intermediate_size / gpu_tp_count; + const int global_k_offset_w2 = tp_part_idx * cpu_k_w2; + const int cpu_scales_per_row_w2 = cpu_k_w2 / group_size; + const int gpu_scales_per_row_w2 = gpu_k_w2 / group_size; + + constexpr int NUM_W13_TASKS = 32; // per matrix (gate or up); total 64 W13 tasks + constexpr int NUM_W2_TASKS = 32; + const int total_tasks = NUM_W13_TASKS * 2 + NUM_W2_TASKS; + + pool->do_work_stealing_job( + total_tasks, nullptr, + [=, &w13_weight_ptrs, &w13_scale_ptrs, &w2_weight_ptrs, &w2_scale_ptrs, this](int task_id) { + if (task_id < NUM_W13_TASKS * 2) { + // ---- W13 weight + scale: per-row scatter (one target_gpu per row) ---- + const bool is_up = task_id >= NUM_W13_TASKS; + const int chunk_idx = task_id % NUM_W13_TASKS; + const auto& bb = is_up ? up_bb_[expert_id] : gate_bb_[expert_id]; + + const int rows_per_task = (cpu_n_w13 + NUM_W13_TASKS - 1) / NUM_W13_TASKS; + const int row_start = chunk_idx * rows_per_task; + const int row_end = std::min(row_start + rows_per_task, cpu_n_w13); + if (row_start >= cpu_n_w13) return; + + for (int row = row_start; row < row_end; row++) { + const int global_n = global_n_offset_w13 + row; + const int target_gpu = global_n / gpu_n_w13; + const int n_in_gpu = global_n % gpu_n_w13; + + // Weight row: full K (cpu_k_w13 == gpu_k_w13 for W13). + uint8_t* w_dst = (uint8_t*)w13_weight_ptrs[target_gpu]; + const size_t expert_w_off = is_up ? gpu_w13_weight_per_mat : 0; + fast_memcpy(w_dst + expert_w_off + (size_t)n_in_gpu * gpu_k_w13, + bb->b + (size_t)row * cpu_k_w13, + cpu_k_w13); + + // Scale row: full K/group_size ue8m0 bytes (fp32 -> ue8m0). + uint8_t* s_dst = (uint8_t*)w13_scale_ptrs[target_gpu]; + const size_t expert_s_off = is_up ? gpu_w13_scale_per_mat : 0; + fast_fp32_to_ue8m0( + s_dst + expert_s_off + (size_t)n_in_gpu * scales_per_row_w13, + bb->d + (size_t)row * scales_per_row_w13, + scales_per_row_w13); + } + } else { + // ---- W2 weight + scale: per-row + per-k-slice scatter ---- + const int chunk_idx = task_id - NUM_W13_TASKS * 2; + const auto& bb = down_bb_[expert_id]; + + const int rows_per_task = (cpu_n_w2 + NUM_W2_TASKS - 1) / NUM_W2_TASKS; + const int row_start = chunk_idx * rows_per_task; + const int row_end = std::min(row_start + rows_per_task, cpu_n_w2); + if (row_start >= cpu_n_w2) return; + + for (int row = row_start; row < row_end; row++) { + // Each CPU TP's K range = [global_k_offset_w2, +cpu_k_w2). + // Iterate gpu_k_w2-aligned k-slices within this range; each + // slice goes to its own target_gpu. cpu < gpu case: one slice + // (length cpu_k_w2). cpu == gpu: one slice (length gpu_k_w2). + // cpu > gpu: multiple slices. + for (int k_start = 0; k_start < cpu_k_w2; k_start += gpu_k_w2) { + const int k_slice_len = std::min(gpu_k_w2, cpu_k_w2 - k_start); + const int global_k = global_k_offset_w2 + k_start; + const int target_gpu = global_k / gpu_k_w2; + const int k_in_gpu = global_k % gpu_k_w2; + + // Weight K-slice + uint8_t* w_dst = (uint8_t*)w2_weight_ptrs[target_gpu]; + fast_memcpy(w_dst + (size_t)row * gpu_k_w2 + k_in_gpu, + bb->b + (size_t)row * cpu_k_w2 + k_start, + k_slice_len); + + // Scale K-slice (k_slice_len/group_size ue8m0 bytes) + const int scale_slice_len = k_slice_len / group_size; + const int k_in_gpu_scale = k_in_gpu / group_size; + const int k_start_scale = k_start / group_size; + uint8_t* s_dst = (uint8_t*)w2_scale_ptrs[target_gpu]; + fast_fp32_to_ue8m0( + s_dst + (size_t)row * gpu_scales_per_row_w2 + k_in_gpu_scale, + bb->d + (size_t)row * cpu_scales_per_row_w2 + k_start_scale, + scale_slice_len); + } + } + } + }, + nullptr); + } +}; + +// ============================================================================ +// TP_MOE specialization for AMX_MXFP8_MOE_TP +// ============================================================================ +template +class TP_MOE> : public TP_MOE>> { + public: + using Base = TP_MOE>>; + using Base::Base; + + void load_weights() override { + auto& config = this->config; + auto& tps = this->tps; + auto& tp_count = this->tp_count; + auto pool = config.pool; + const uint64_t* physical_to_logical_map = (const uint64_t*)config.physical_to_logical_map; + + bool use_per_expert_ptrs = !config.gate_projs.empty(); + + if (config.gate_projs.empty() && config.gate_scale == nullptr) + throw std::runtime_error("MXFP8 MoE requires FP8 weights with ue8m0 KGroup Scale"); + + printf("MXFP8 MoE: loading from %s\n", + use_per_expert_ptrs ? "per-expert pointers (gate_projs)" : "flat arrays with KGroup Scale"); + + int& group_size = config.quant_config.group_size; + + pool->dispense_backend()->do_numa_job([&, this](int i) { + auto& tpc = tps[i]->config_; + size_t weight_elem_count = (size_t)tpc.intermediate_size * tpc.hidden_size; + size_t scales_elem_count = (tpc.hidden_size / group_size) * tpc.intermediate_size; + + // FP8: 1 byte per element (no / 2) + tpc.gate_proj = new uint8_t[tpc.expert_num * weight_elem_count]; + tpc.up_proj = new uint8_t[tpc.expert_num * weight_elem_count]; + tpc.down_proj = new uint8_t[tpc.expert_num * weight_elem_count]; + // Scales: uint8 ue8m0 (will be converted to FP32 inside per-TP load_weights) + tpc.gate_scale = new uint8_t[tpc.expert_num * scales_elem_count]; + tpc.up_scale = new uint8_t[tpc.expert_num * scales_elem_count]; + tpc.down_scale = new uint8_t[tpc.expert_num * scales_elem_count]; + + if (use_per_expert_ptrs) { + pool->get_subpool(i)->do_work_stealing_job( + tpc.expert_num, nullptr, + [&, i](int expert_id_) { + size_t expert_id = expert_map(physical_to_logical_map, expert_id_); + + uint8_t* src_gate = (uint8_t*)config.gate_projs[0][expert_id]; + uint8_t* src_up = (uint8_t*)config.up_projs[0][expert_id]; + uint8_t* src_down = (uint8_t*)config.down_projs[0][expert_id]; + uint8_t* src_gate_scale = (uint8_t*)config.gate_scales[0][expert_id]; + uint8_t* src_up_scale = (uint8_t*)config.up_scales[0][expert_id]; + uint8_t* src_down_scale = (uint8_t*)config.down_scales[0][expert_id]; + + // gate/up: row-major [intermediate_size, hidden_size], TP splits along hidden + memcpy((uint8_t*)tpc.gate_proj + expert_id * weight_elem_count, + src_gate + i * weight_elem_count, weight_elem_count); + memcpy((uint8_t*)tpc.up_proj + expert_id * weight_elem_count, + src_up + i * weight_elem_count, weight_elem_count); + memcpy((uint8_t*)tpc.gate_scale + expert_id * scales_elem_count, + src_gate_scale + i * scales_elem_count, scales_elem_count); + memcpy((uint8_t*)tpc.up_scale + expert_id * scales_elem_count, + src_up_scale + i * scales_elem_count, scales_elem_count); + + // down: row-major [hidden_size, intermediate_size], TP splits along intermediate + for (size_t col = 0; col < config.hidden_size; col++) { + memcpy((uint8_t*)tpc.down_proj + expert_id * weight_elem_count + col * tpc.intermediate_size, + src_down + col * config.intermediate_size + i * tpc.intermediate_size, + tpc.intermediate_size); + memcpy((uint8_t*)tpc.down_scale + expert_id * scales_elem_count + + col * (tpc.intermediate_size / group_size), + src_down_scale + col * (config.intermediate_size / group_size) + + i * (tpc.intermediate_size / group_size), + tpc.intermediate_size / group_size); + } + }, + nullptr); + } else { + if (tpc.load == false) { + pool->get_subpool(i)->do_work_stealing_job( + tpc.expert_num, nullptr, + [&, i](int expert_id_) { + size_t expert_id = expert_map(physical_to_logical_map, expert_id_); + + // gate/up: no >> 1 (FP8 = 1 byte/element) + memcpy((uint8_t*)tpc.gate_proj + expert_id * weight_elem_count, + (uint8_t*)config.gate_proj + + expert_id * config.intermediate_size * config.hidden_size + i * weight_elem_count, + weight_elem_count); + memcpy((uint8_t*)tpc.up_proj + expert_id * weight_elem_count, + (uint8_t*)config.up_proj + + expert_id * config.intermediate_size * config.hidden_size + i * weight_elem_count, + weight_elem_count); + // scales (uint8, not bf16) + memcpy((uint8_t*)tpc.gate_scale + expert_id * scales_elem_count, + (uint8_t*)config.gate_scale + + expert_id * (config.hidden_size / group_size) * config.intermediate_size + + i * scales_elem_count, + scales_elem_count); + memcpy((uint8_t*)tpc.up_scale + expert_id * scales_elem_count, + (uint8_t*)config.up_scale + + expert_id * (config.hidden_size / group_size) * config.intermediate_size + + i * scales_elem_count, + scales_elem_count); + + // down: column-wise TP split + for (size_t col = 0; col < config.hidden_size; col++) { + memcpy((uint8_t*)tpc.down_proj + expert_id * weight_elem_count + col * tpc.intermediate_size, + (uint8_t*)config.down_proj + expert_id * config.intermediate_size * config.hidden_size + + col * config.intermediate_size + i * tpc.intermediate_size, + tpc.intermediate_size); + memcpy((uint8_t*)tpc.down_scale + expert_id * scales_elem_count + + col * (tpc.intermediate_size / group_size), + (uint8_t*)config.down_scale + + expert_id * (config.intermediate_size / group_size) * config.hidden_size + + col * (config.intermediate_size / group_size) + + i * (tpc.intermediate_size / group_size), + tpc.intermediate_size / group_size); + } + }, + nullptr); + } + } + printf("MXFP8 TP %d load weight done.\n", i); + }); + + DO_TPS_LOAD_WEIGHTS(pool); + + pool->dispense_backend()->do_numa_job([&, this](int i) { + auto& tpc = tps[i]->config_; + delete[](uint8_t*)(tpc.gate_proj); + delete[](uint8_t*)(tpc.up_proj); + delete[](uint8_t*)(tpc.down_proj); + delete[](uint8_t*)(tpc.gate_scale); + delete[](uint8_t*)(tpc.up_scale); + delete[](uint8_t*)(tpc.down_scale); + }); + + this->weights_loaded = true; + } + + void write_weight_scale_to_buffer(int gpu_tp_count, int expert_id, + const std::vector& w13_weight_ptrs, + const std::vector& w13_scale_ptrs, + const std::vector& w2_weight_ptrs, + const std::vector& w2_scale_ptrs) { + if (!this->weights_loaded) throw std::runtime_error("Not Loaded"); + if (this->tps.empty()) throw std::runtime_error("No TP parts initialized"); + if (w13_weight_ptrs.size() != (size_t)gpu_tp_count || w13_scale_ptrs.size() != (size_t)gpu_tp_count || + w2_weight_ptrs.size() != (size_t)gpu_tp_count || w2_scale_ptrs.size() != (size_t)gpu_tp_count) + throw std::runtime_error("Pointer arrays size must match gpu_tp_count"); + + this->config.pool->dispense_backend()->do_numa_job([&, this](int i) { + this->tps[i]->write_weights_to_buffer(gpu_tp_count, this->tp_count, expert_id, this->config, + w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs); + }); + } +}; + +#endif // CPUINFER_OPERATOR_AMX_MXFP8_MOE_H diff --git a/kt-kernel/operators/avx2/avx2_bf16_utils.hpp b/kt-kernel/operators/avx2/avx2_bf16_utils.hpp index ed9222c4..d39cda9d 100644 --- a/kt-kernel/operators/avx2/avx2_bf16_utils.hpp +++ b/kt-kernel/operators/avx2/avx2_bf16_utils.hpp @@ -127,6 +127,47 @@ static inline __m256 act_fn(__m256 gate_val, __m256 up_val) { return _mm256_mul_ps(act_val, up_val); } +// Overload with swiglu_limit: DeepSeek V4-Flash 2604B asymmetric clamp (no alpha). +// gate = min(gate, limit) (one-sided pre-silu) +// up = clamp(up, -limit, limit) (symmetric) +// Mirrors amx::act_fn(g, u, swiglu_limit). +static inline __m256 act_fn(__m256 gate_val, __m256 up_val, float swiglu_limit) { + if (swiglu_limit > 0.0f) { + const __m256 pos_lim = _mm256_set1_ps(swiglu_limit); + const __m256 neg_lim = _mm256_set1_ps(-swiglu_limit); + gate_val = _mm256_min_ps(gate_val, pos_lim); + up_val = _mm256_min_ps(up_val, pos_lim); + up_val = _mm256_max_ps(up_val, neg_lim); + } + return act_fn(gate_val, up_val); +} + +// MiniMax M3 \"swigluoai\" activation + DeepSeek V4 \"silu\" unified entry point. +// alpha > 0 -> swigluoai: gate * sigmoid(gate * alpha) * (up + 1), symmetric clamp on both +// alpha == 0 -> falls back to silu (with optional one-sided clamp via the overload above) +// Mirrors amx::act_fn(g, u, swiglu_limit, swiglu_alpha). +static inline __m256 act_fn(__m256 gate_val, __m256 up_val, float swiglu_limit, float swiglu_alpha) { + if (swiglu_alpha > 0.0f) { + if (swiglu_limit > 0.0f) { + const __m256 pos_lim = _mm256_set1_ps(swiglu_limit); + const __m256 neg_lim = _mm256_set1_ps(-swiglu_limit); + gate_val = _mm256_min_ps(gate_val, pos_lim); + gate_val = _mm256_max_ps(gate_val, neg_lim); + up_val = _mm256_min_ps(up_val, pos_lim); + up_val = _mm256_max_ps(up_val, neg_lim); + } + // sigmoid(gate * alpha) + __m256 neg_ga = _mm256_mul_ps(gate_val, _mm256_set1_ps(-swiglu_alpha)); + neg_ga = _mm256_min_ps(neg_ga, _mm256_set1_ps(88.0f)); + __m256 exp_neg = exp_avx2(neg_ga); + __m256 sigmoid_val = _mm256_div_ps(_mm256_set1_ps(1.0f), + _mm256_add_ps(_mm256_set1_ps(1.0f), exp_neg)); + __m256 up_plus_1 = _mm256_add_ps(up_val, _mm256_set1_ps(1.0f)); + return _mm256_mul_ps(_mm256_mul_ps(gate_val, sigmoid_val), up_plus_1); + } + return act_fn(gate_val, up_val, swiglu_limit); +} + } // namespace avx2 #endif // CPUINFER_OPERATOR_AVX2_BF16_UTILS_H diff --git a/kt-kernel/operators/avx2/moe_base.hpp b/kt-kernel/operators/avx2/moe_base.hpp index dac66ab0..ae5e17c8 100644 --- a/kt-kernel/operators/avx2/moe_base.hpp +++ b/kt-kernel/operators/avx2/moe_base.hpp @@ -505,6 +505,8 @@ class AVX2_MOE_BASE { int expert_idx = m_expert_id_map_[task_id / nth]; int ith = task_id % nth; auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth); + const float swiglu_limit = config_.swiglu_limit; + const float swiglu_alpha = config_.swiglu_alpha; for (int i = 0; i < m_local_num_[expert_idx]; i++) { ggml_bf16_t* gate_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; ggml_bf16_t* up_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; @@ -512,15 +514,28 @@ class AVX2_MOE_BASE { for (; j + 8 <= n_end; j += 8) { __m256 gate_val = avx2::load_bf16_to_fp32(gate_ptr + j); __m256 up_val = avx2::load_bf16_to_fp32(up_ptr + j); - __m256 result = avx2::act_fn(gate_val, up_val); + __m256 result = avx2::act_fn(gate_val, up_val, swiglu_limit, swiglu_alpha); avx2::store_fp32_to_bf16(gate_ptr + j, result); } - // Scalar tail + // Scalar tail — mirror the vectorized swigluoai / silu paths in avx2::act_fn. for (; j < n_end; j++) { float g = GGML_BF16_TO_FP32(gate_ptr[j]); float u = GGML_BF16_TO_FP32(up_ptr[j]); - float sigmoid_g = 1.0f / (1.0f + expf(-g)); - gate_ptr[j] = GGML_FP32_TO_BF16(g * sigmoid_g * u); + if (swiglu_alpha > 0.0f) { + if (swiglu_limit > 0.0f) { + g = std::min(std::max(g, -swiglu_limit), swiglu_limit); + u = std::min(std::max(u, -swiglu_limit), swiglu_limit); + } + float sigmoid_ga = 1.0f / (1.0f + expf(-g * swiglu_alpha)); + gate_ptr[j] = GGML_FP32_TO_BF16(g * sigmoid_ga * (u + 1.0f)); + } else { + if (swiglu_limit > 0.0f) { + g = std::min(g, swiglu_limit); + u = std::min(std::max(u, -swiglu_limit), swiglu_limit); + } + float sigmoid_g = 1.0f / (1.0f + expf(-g)); + gate_ptr[j] = GGML_FP32_TO_BF16(g * sigmoid_g * u); + } } } }; diff --git a/kt-kernel/operators/avx2/mxfp8-moe.hpp b/kt-kernel/operators/avx2/mxfp8-moe.hpp new file mode 100644 index 00000000..8225fe8e --- /dev/null +++ b/kt-kernel/operators/avx2/mxfp8-moe.hpp @@ -0,0 +1,740 @@ +/** + * @Description : AVX2 MXFP8 MoE operator (AVX2 sibling of amx/mxfp8-moe.hpp) + * @Author : yyj and Claude + * @Date : 2026-06-09 + * @Version : 0.1.0 + * @Copyright (c) 2024 by KVCache.AI, All Rights Reserved. + * + * Serves MiniMax M3 Preview (MXFP8 quantized checkpoint) on AVX2-only hardware + * (Haswell+ / Zen+, no AVX-512 / VBMI / dpbf16_ps required). + * + * Algorithm parity with amx/mxfp8-moe.hpp: + * Weight: FP8 E4M3fn (1 byte/element, row-major [n, k]) + * Scale: ue8m0 per-group (group_size=32), converted to FP32 on load + * Act: BF16 (row-major [m, k]), promoted to FP32 per 8-lane chunk + * Decode: 256-entry FP32 LUT + _mm256_i32gather_ps (vs VBMI permutex2var) + * Dot prod: BF16→FP32 promote + _mm256_fmadd_ps (vs _mm512_dpbf16_ps) + * Activation: swiglu_oai via avx2_bf16_utils.hpp::act_fn (alpha/limit aware) + * + * Inner-loop pattern: **deferred-hsum** — accumulate scale·v_g into a vector + * row accumulator across all k-groups, hsum once per output row. ~10× fewer + * hsums than per-group hsum used by gptq_int4-moe.hpp (which can't defer + * because its scale is baked into dequant). MXFP8 has scale separate from + * dequant, so deferring is math-equivalent: hsum(sum_g s[g]*v_g) = + * sum_g s[g]*hsum(v_g) since hsum is linear. + **/ +#ifndef CPUINFER_OPERATOR_AVX2_MXFP8_MOE_H +#define CPUINFER_OPERATOR_AVX2_MXFP8_MOE_H + +#include +#include +#include + +#include "avx2_bf16_gemm.hpp" +#include "avx2_bf16_utils.hpp" +#include "fp8_dequant.hpp" +#include "moe_base.hpp" + +namespace avx2 { + +// ============================================================================ +// ue8m0 → FP32 scale conversion (vectorized) +// +// ue8m0 byte b ∈ [0, 255] → FP32 value 2^(b - 127). +// IEEE754 FP32 layout: [sign:1][exp:8][mantissa:23]. Setting exp = b and +// mantissa = 0 (with sign = 0) gives exactly 2^(b - 127). +// Edge case: b=0 → (0 << 23) = 0.0f (technically 2^-127 ≈ 5.9e-39, accepted). +// Verbatim copy of amx/mxfp8-moe.hpp::convert_ue8m0_to_fp32 — already pure AVX2. +// ============================================================================ +static inline void convert_ue8m0_to_fp32(float* __restrict dst, + const uint8_t* __restrict src, + size_t count) { + size_t i = 0; + for (; i + 8 <= count; i += 8) { + __m128i bytes = _mm_loadl_epi64((__m128i const*)(src + i)); + __m256i dwords = _mm256_cvtepu8_epi32(bytes); + __m256 floats = _mm256_castsi256_ps(_mm256_slli_epi32(dwords, 23)); + _mm256_storeu_ps(dst + i, floats); + } + for (; i < count; i++) { + uint32_t bits = static_cast(src[i]) << 23; + std::memcpy(dst + i, &bits, sizeof(float)); + } +} + +// ============================================================================ +// FP32 -> ue8m0 (uint8) — reverse of convert_ue8m0_to_fp32 (AVX2 version) +// +// Used by write_weights_to_buffer for layerwise prefill: CPU stores FP32 scales +// internally (in BufferB.d for fast FMA in inner loop), but GPU's MXFP8 buffer +// expects ue8m0 uint8 (1 byte/scale). We extract bits 23-30 of each FP32 (the +// exponent field) and store as uint8. Bit-exact iff the FP32 came from +// convert_ue8m0_to_fp32 (mantissa=0), which is the case for layerwise prefill +// since CPU forward never modifies BufferB.d. +// ============================================================================ +static inline void fast_fp32_to_ue8m0(uint8_t* __restrict dst, + const float* __restrict src, + size_t count) { + size_t i = 0; + for (; i + 8 <= count; i += 8) { + __m256 v = _mm256_loadu_ps(src + i); + __m256i bits = _mm256_castps_si256(v); + // Extract bits 23-30 (the FP32 exponent) -> 8 uint32 values in [0,255] + __m256i shifted = _mm256_srli_epi32(bits, 23); + // Pack 8 uint32 -> 8 uint16 -> 8 uint8 (saturating; values are <= 255 so no-op) + __m128i lo = _mm256_castsi256_si128(shifted); + __m128i hi = _mm256_extracti128_si256(shifted, 1); + __m128i p16 = _mm_packus_epi32(lo, hi); // 8 x uint16 + __m128i p8 = _mm_packus_epi16(p16, p16); // 8 x uint8 in low 64 bits + _mm_storel_epi64((__m128i*)(dst + i), p8); + } + for (; i < count; i++) { + uint32_t bits; + std::memcpy(&bits, &src[i], sizeof(uint32_t)); + dst[i] = (uint8_t)((bits >> 23) & 0xFF); + } +} + +// ============================================================================ +// GemmKernelAVX2MXFP8 — kernel descriptor +// +// M_STEP=1, N_STEP=8, K_STEP=K_GROUP_SIZE=32 +// BufferA: BF16 activations [m, k] (raw, no pre-promotion) +// BufferB: FP8 weights [n, k] + FP32 scales [n, k/group_size] +// BufferC: FP32 output [m, n] +// ============================================================================ +struct GemmKernelAVX2MXFP8 { + using dt = ggml_bf16_t; + using output_t = float; + + static constexpr int M_STEP = 1; + static constexpr int N_STEP = 8; + static constexpr int K_STEP = 32; // == K_GROUP_SIZE + static constexpr int K_GROUP_SIZE = 32; // MXFP8 group size per OCP standard + static constexpr int N_BLOCK = 64; // ~8 N_STEP tiles per task + static constexpr int K_BLOCK = 6144; // M3 hidden=6144 in one pass + static constexpr double ELEMENT_SIZE = 1.0; // FP8 = 1 byte/element + + static void config() {} + + static int recommended_nth(int n) { + return std::max(1, (n + N_BLOCK - 1) / N_BLOCK); + } + + static std::pair split_range_n(int n, int ith, int nth) { + return ::avx2::split_range(n, ith, nth); + } + + // -------------------------------------------------------------------------- + // BufferA: BF16 activations [m, k] — identical to AVX2 BF16/FP8 backends + // -------------------------------------------------------------------------- + struct BufferA { + ggml_bf16_t* data = nullptr; + size_t max_m = 0; + size_t k = 0; + + BufferA() = default; + BufferA(size_t m, size_t k_, void* ptr) : max_m(m), k(k_), data((ggml_bf16_t*)ptr) {} + + static size_t required_size(size_t m, size_t k) { + return m * k * sizeof(ggml_bf16_t); + } + + void set_data(void* ptr) { data = (ggml_bf16_t*)ptr; } + + void from_mat(int m, const ggml_bf16_t* src, int ith, int nth) { + if (ith == 0 && nth == 1) { + std::memcpy(data, src, (size_t)m * k * sizeof(ggml_bf16_t)); + } else { + auto [m_start, m_end] = ::avx2::split_range(m, ith, nth); + std::memcpy(data + m_start * k, src + m_start * k, + (size_t)(m_end - m_start) * k * sizeof(ggml_bf16_t)); + } + } + }; + + // -------------------------------------------------------------------------- + // BufferB: FP8 [n, k] + FP32 scales [n, k/group_size] + // b = ptr; (n*k FP8 bytes) + // d = ptr + n*k; (n * (k/group_size) FP32 scales) + // Identical memory layout to amx::BufferBMXFP8KGroupImpl, so the M3 checkpoint + // loader can target both AMX and AVX2 sides with the same pointer arithmetic. + // -------------------------------------------------------------------------- + struct BufferB { + uint8_t* b = nullptr; + float* d = nullptr; + int n = 0; + int k = 0; + int k_group_size = K_GROUP_SIZE; + int k_group_count = 0; + + BufferB() = default; + BufferB(size_t n_, size_t k_, int gs, void* ptr) + : n((int)n_), k((int)k_), k_group_size(gs) { + if (k % gs != 0) { + printf("BufferB(MXFP8 AVX2): k=%d not divisible by group_size=%d\n", k, gs); + throw std::runtime_error("MXFP8 AVX2: k must be divisible by group_size"); + } + k_group_count = k / gs; + b = (uint8_t*)ptr; + d = (float*)((uint8_t*)ptr + (size_t)n * k); + } + + static size_t required_size(size_t n, size_t k, int gs) { + return n * k + n * (k / gs) * sizeof(float); + } + + // Copy raw FP8 bytes from checkpoint. 1 byte/element, row-major. + void from_raw_mat(const uint8_t* src_weights, int ith, int nth) { + auto [n_start, n_end] = ::avx2::split_range(n, ith, nth); + if (n_start >= n_end) return; + const size_t row_bytes = (size_t)k; + std::memcpy(b + (size_t)n_start * row_bytes, + src_weights + (size_t)n_start * row_bytes, + (size_t)(n_end - n_start) * row_bytes); + } + + // Returns pointer to FP8 byte at (n_begin, k_begin). + uint8_t* get_submat(int n_begin, int k_begin) { + return b + (size_t)n_begin * k + k_begin; + } + + // Returns pointer to FP32 scale for row n_begin starting at k-group k_begin/group_size. + float* get_scale(int n_begin, int k_begin) { + return d + (size_t)n_begin * k_group_count + k_begin / k_group_size; + } + }; + + // -------------------------------------------------------------------------- + // BufferC: FP32 output [m, n] — identical to AVX2 BF16/FP8 backends + // -------------------------------------------------------------------------- + struct BufferC { + float* data = nullptr; + size_t max_m = 0; + size_t n = 0; + + BufferC() = default; + BufferC(size_t m, size_t n_, void* ptr) : max_m(m), n(n_), data((float*)ptr) {} + + static size_t required_size(size_t m, size_t n) { + return m * n * sizeof(float); + } + + void set_data(void* ptr) { data = (float*)ptr; } + + void to_mat(int m, ggml_bf16_t* dst, int ith, int nth) { + auto [n_start, n_end] = ::avx2::split_range((int)n, ith, nth); + for (int mi = 0; mi < m; mi++) { + float* src_row = data + mi * n; + ggml_bf16_t* dst_row = dst + mi * n; + int j = n_start; + for (; j + 8 <= n_end; j += 8) { + __m256 v = _mm256_loadu_ps(src_row + j); + store_fp32_to_bf16(dst_row + j, v); + } + for (; j < n_end; j++) { + dst_row[j] = GGML_FP32_TO_BF16(src_row[j]); + } + } + } + }; +}; + +// ============================================================================ +// gemm_mxfp8_kgroup — AVX2 MXFP8 GEMM (BF16 act × FP8 weight + per-K-group scale) +// +// C[m, n] = sum_g scale[n, g] * sum_{i in group g} A[m, k_base+i] * dequant(B[n, k_base+i]) +// +// Per output row (ni): one vector accumulator `row_acc` lives across all groups; +// per group, do 4× FMA into a fresh `v` (8 lanes), then `row_acc += s[g] * v` +// via broadcast+FMA. After all groups, hsum_avx2(row_acc) once for the scalar. +// ============================================================================ +static inline void gemm_mxfp8_kgroup( + int m, int n, int k, + GemmKernelAVX2MXFP8::BufferA& a, + GemmKernelAVX2MXFP8::BufferB& b, + GemmKernelAVX2MXFP8::BufferC& c, + int ith, int nth) { + + ensure_fp8_lut_initialized(); // idempotent; cheap after first call + + auto [n_start, n_end] = ::avx2::split_range(n, ith, nth); + if (n_start >= n_end) return; + const int kg_count = b.k_group_count; + + for (int ni = n_start; ni < n_end; ni++) { + const uint8_t* w_base = b.get_submat(ni, 0); + const float* s = b.get_scale(ni, 0); + + for (int mi = 0; mi < m; mi++) { + const ggml_bf16_t* a_base = a.data + (size_t)mi * a.k; + __m256 row_acc = _mm256_setzero_ps(); + const uint8_t* w = w_base; + const ggml_bf16_t* a_row = a_base; + + for (int g = 0; g < kg_count; g++) { + // 4-chunk × 8-lane unroll: 32 elements per k-group + __m256 v = _mm256_setzero_ps(); + v = _mm256_fmadd_ps(fp8x8_to_fp32x8(w + 0), load_bf16_to_fp32(a_row + 0), v); + v = _mm256_fmadd_ps(fp8x8_to_fp32x8(w + 8), load_bf16_to_fp32(a_row + 8), v); + v = _mm256_fmadd_ps(fp8x8_to_fp32x8(w + 16), load_bf16_to_fp32(a_row + 16), v); + v = _mm256_fmadd_ps(fp8x8_to_fp32x8(w + 24), load_bf16_to_fp32(a_row + 24), v); + // Defer hsum: broadcast scale, FMA into vector row accumulator. + // Math: hsum(Σ_g s[g] * v_g) = Σ_g s[g] * hsum(v_g) (hsum is linear). + row_acc = _mm256_fmadd_ps(_mm256_set1_ps(s[g]), v, row_acc); + w += 32; + a_row += 32; + } + c.data[mi * n + ni] = hsum_avx2(row_acc); + } + } +} + +} // namespace avx2 + +// ============================================================================ +// AVX2_MXFP8_MOE_TP — CRTP wrapper, mirrors amx::AMX_MXFP8_MOE_TP +// +// Activation (swigluoai) is handled by AVX2_MOE_BASE::apply_activation, which +// dispatches via avx2::act_fn(g, u, swiglu_limit, swiglu_alpha). Our derived +// class just needs derived_init / buffer factories / GEMM dispatch / load. +// ============================================================================ +template +class AVX2_MXFP8_MOE_TP : public AVX2_MOE_BASE> { + using Base = AVX2_MOE_BASE>; + using Base::config_; + using Base::down_ba_; + using Base::down_bb_; + using Base::down_bc_; + using Base::gate_bb_; + using Base::gate_bc_; + using Base::gate_up_ba_; + using Base::m_local_num_; + using Base::tp_part_idx; + using Base::up_bb_; + using Base::up_bc_; + + public: + using typename Base::input_t; + using typename Base::output_t; + + AVX2_MXFP8_MOE_TP() = default; + AVX2_MXFP8_MOE_TP(GeneralMOEConfig config, int tp_part_idx_ = 0) : Base(config, tp_part_idx_) {} + + void derived_init() { + avx2::ensure_fp8_lut_initialized(); // single-threaded init before any forward + auto& quant_config = config_.quant_config; + if (quant_config.group_size != 32 || quant_config.zero_point) { + throw std::runtime_error("AVX2 MXFP8 MoE requires group_size == 32 and no zero_point"); + } + if (config_.hidden_size % quant_config.group_size != 0 || + config_.intermediate_size % quant_config.group_size != 0) { + throw std::runtime_error("AVX2 MXFP8 MoE: hidden_size and intermediate_size must be divisible by group_size"); + } + printf("Created AVX2_MXFP8_MOE_TP %d at numa %d (group_size=%d, swiglu_alpha=%.4f, swiglu_limit=%.4f)\n", + tp_part_idx, numa_node_of_cpu(sched_getcpu()), + quant_config.group_size, config_.swiglu_alpha, config_.swiglu_limit); + } + + ~AVX2_MXFP8_MOE_TP() = default; + + // CRTP buffer creation + size_t buffer_a_required_size_impl(size_t m, size_t k) const { return T::BufferA::required_size(m, k); } + size_t buffer_b_required_size_impl(size_t n, size_t k) const { + return T::BufferB::required_size(n, k, config_.quant_config.group_size); + } + size_t buffer_c_required_size_impl(size_t m, size_t n) const { return T::BufferC::required_size(m, n); } + + std::shared_ptr make_buffer_a_impl(size_t m, size_t k, void* data) const { + return std::make_shared(m, k, data); + } + std::shared_ptr make_buffer_b_impl(size_t n, size_t k, void* data) const { + return std::make_shared(n, k, config_.quant_config.group_size, data); + } + std::shared_ptr make_buffer_c_impl(size_t m, size_t n, void* data) const { + return std::make_shared(m, n, data); + } + + // GEMM dispatch + void do_gate_up_gemm(bool do_up, int expert_idx, int ith, int nth, [[maybe_unused]] int qlen) { + int m = m_local_num_[expert_idx]; + auto& ba = gate_up_ba_[expert_idx]; + auto& bb = do_up ? up_bb_[expert_idx] : gate_bb_[expert_idx]; + auto& bc = do_up ? up_bc_[expert_idx] : gate_bc_[expert_idx]; + avx2::gemm_mxfp8_kgroup(m, config_.intermediate_size, config_.hidden_size, *ba, *bb, *bc, ith, nth); + } + + void do_down_gemm(int expert_idx, int ith, int nth, [[maybe_unused]] int qlen) { + int m = m_local_num_[expert_idx]; + avx2::gemm_mxfp8_kgroup(m, config_.hidden_size, config_.intermediate_size, + *down_ba_[expert_idx], *down_bb_[expert_idx], *down_bc_[expert_idx], ith, nth); + } + + // Load FP8 weights + ue8m0 scales from checkpoint. + // gate_proj/up_proj/down_proj: uint8_t [E, N, K] (FP8 E4M3fn bytes) + // gate_scale/up_scale/down_scale: uint8_t [E, N, K/group_size] (ue8m0) + // We memcpy weights and bit-shift convert scales to FP32 in-place. + void load_weights() { + auto& quant_config = config_.quant_config; + const uint64_t* physical_to_logical_map = (const uint64_t*)config_.physical_to_logical_map; + auto pool = config_.pool->get_subpool(tp_part_idx); + + if (quant_config.group_size == 0 || quant_config.zero_point) + throw std::runtime_error("AVX2 MXFP8 MoE requires group_size > 0 and no zero_point"); + if (config_.gate_scale == nullptr) + throw std::runtime_error("AVX2 MXFP8 MoE requires native MXFP8 weights with ue8m0 scales"); + + // --- Load FP8 weights (1 byte/element) — copy into BufferB.b via from_raw_mat --- + int nth = T::recommended_nth(config_.intermediate_size); + pool->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [this, nth, physical_to_logical_map](int task_id) { + uint64_t expert_idx = task_id / nth; + uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + int ith = task_id % nth; + size_t weight_offset = (size_t)logical_expert_id * config_.intermediate_size * config_.hidden_size; + gate_bb_[expert_idx]->from_raw_mat((uint8_t*)config_.gate_proj + weight_offset, ith, nth); + up_bb_[expert_idx]->from_raw_mat((uint8_t*)config_.up_proj + weight_offset, ith, nth); + }, + nullptr); + + nth = T::recommended_nth(config_.hidden_size); + pool->do_work_stealing_job( + nth * config_.expert_num, nullptr, + [this, nth, physical_to_logical_map](int task_id) { + uint64_t expert_idx = task_id / nth; + uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + int ith = task_id % nth; + size_t weight_offset = (size_t)logical_expert_id * config_.hidden_size * config_.intermediate_size; + down_bb_[expert_idx]->from_raw_mat((uint8_t*)config_.down_proj + weight_offset, ith, nth); + }, + nullptr); + + // --- Convert ue8m0 scales → FP32 in BufferB.d --- + pool->do_work_stealing_job( + config_.expert_num, nullptr, + [this, physical_to_logical_map](int task_id) { + uint64_t expert_idx = task_id; + uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx); + size_t scale_count = + ((size_t)config_.intermediate_size * config_.hidden_size) / config_.quant_config.group_size; + avx2::convert_ue8m0_to_fp32( + gate_bb_[expert_idx]->d, + (const uint8_t*)config_.gate_scale + logical_expert_id * scale_count, scale_count); + avx2::convert_ue8m0_to_fp32( + up_bb_[expert_idx]->d, + (const uint8_t*)config_.up_scale + logical_expert_id * scale_count, scale_count); + avx2::convert_ue8m0_to_fp32( + down_bb_[expert_idx]->d, + (const uint8_t*)config_.down_scale + logical_expert_id * scale_count, scale_count); + }, + nullptr); + } + + // -------------------------------------------------------------------------- + // write_weights_to_buffer: copies CPU expert weights to GPU pinned host buffer + // for layerwise prefill (sglang full-GPU fallback at large prefill token count). + // + // Mirrors amx::AMX_MXFP8_MOE_TP::write_weights_to_buffer (amx/mxfp8-moe.hpp:662) + // with two substitutions for AVX2: + // fast_memcpy → std::memcpy (libc memcpy is AVX2-optimized on Haswell+) + // fast_fp32_to_bf16 → avx2::fast_fp32_to_ue8m0 (writes 1 byte/scale matching + // GPU's torch.uint8 ue8m0 layout) + // + // Scale dtype on the GPU side is torch.uint8 (ue8m0), see + // kt-sglang/python/sglang/srt/layers/quantization/fp8.py:872. + // -------------------------------------------------------------------------- + void write_weights_to_buffer(int gpu_tp_count, [[maybe_unused]] int cpu_tp_count, int expert_id, + const GeneralMOEConfig& full_config, + const std::vector& w13_weight_ptrs, + const std::vector& w13_scale_ptrs, + const std::vector& w2_weight_ptrs, + const std::vector& w2_scale_ptrs) const { + // Unified per-row scatter that works for any (cpu_tp_count, gpu_tp_count) + // relationship (mirrors avx2/fp8-moe.hpp:347-451). Replaces the earlier + // `cpu_tp_count >= gpu_tp_count` direct-write branch which was a silent + // no-op for the realistic 2-NUMA × tp>=4 case (typical M3 deployment: + // 2 NUMA sockets × tp=8). Each task processes a row chunk of this CPU + // TP's slice; per row we compute target_gpu = global_n / gpu_n_w13 (W13) + // or scatter across multiple gpu_tp k-slices (W2) and write to the + // corresponding GPU TP staging buffer at the right offset. + // + // MXFP8 specifics vs FP8 block: + // - Scale dtype on GPU side is uint8 ue8m0 (1 byte per `group_size` + // weights, per row, not per block). + // - Source scale (bb->d) is stored as FP32 in kt-kernel BufferB and is + // converted to ue8m0 via avx2::fast_fp32_to_ue8m0 at write time. + // - Scale layout matches GPU: (E, 2*intermediate, hidden/group_size) + // for W13 and (E, hidden, intermediate/group_size) for W2. + + auto& config = config_; + auto pool = config.pool->get_subpool(tp_part_idx); + const int group_size = config.quant_config.group_size; + + // ========= W13 (gate+up): Shape [intermediate, hidden], split by N only ========= + const int cpu_n_w13 = config.intermediate_size; + const int cpu_k_w13 = config.hidden_size; + const int gpu_n_w13 = full_config.intermediate_size / gpu_tp_count; + const int gpu_k_w13 = full_config.hidden_size; + const int global_n_offset_w13 = tp_part_idx * cpu_n_w13; + const size_t gpu_w13_weight_per_mat = (size_t)gpu_n_w13 * gpu_k_w13; + const int scales_per_row_w13 = cpu_k_w13 / group_size; + const size_t gpu_w13_scale_per_mat = (size_t)gpu_n_w13 * scales_per_row_w13; + + // ========= W2 (down): Shape [hidden, intermediate], split by K ========= + const int cpu_n_w2 = config.hidden_size; + const int cpu_k_w2 = config.intermediate_size; + const int gpu_k_w2 = full_config.intermediate_size / gpu_tp_count; + const int global_k_offset_w2 = tp_part_idx * cpu_k_w2; + const int cpu_scales_per_row_w2 = cpu_k_w2 / group_size; + const int gpu_scales_per_row_w2 = gpu_k_w2 / group_size; + + constexpr int NUM_W13_TASKS = 32; // per matrix (gate or up); total 64 W13 tasks + constexpr int NUM_W2_TASKS = 32; + const int total_tasks = NUM_W13_TASKS * 2 + NUM_W2_TASKS; + + pool->do_work_stealing_job( + total_tasks, nullptr, + [=, &w13_weight_ptrs, &w13_scale_ptrs, &w2_weight_ptrs, &w2_scale_ptrs, this](int task_id) { + if (task_id < NUM_W13_TASKS * 2) { + // ---- W13 weight + scale: per-row scatter (one target_gpu per row) ---- + const bool is_up = task_id >= NUM_W13_TASKS; + const int chunk_idx = task_id % NUM_W13_TASKS; + const auto& bb = is_up ? up_bb_[expert_id] : gate_bb_[expert_id]; + + const int rows_per_task = (cpu_n_w13 + NUM_W13_TASKS - 1) / NUM_W13_TASKS; + const int row_start = chunk_idx * rows_per_task; + const int row_end = std::min(row_start + rows_per_task, cpu_n_w13); + if (row_start >= cpu_n_w13) return; + + for (int row = row_start; row < row_end; row++) { + const int global_n = global_n_offset_w13 + row; + const int target_gpu = global_n / gpu_n_w13; + const int n_in_gpu = global_n % gpu_n_w13; + + // Weight row: full K (cpu_k_w13 == gpu_k_w13 for W13). + uint8_t* w_dst = (uint8_t*)w13_weight_ptrs[target_gpu]; + const size_t expert_w_off = is_up ? gpu_w13_weight_per_mat : 0; + std::memcpy(w_dst + expert_w_off + (size_t)n_in_gpu * gpu_k_w13, + bb->b + (size_t)row * cpu_k_w13, + cpu_k_w13); + + // Scale row: full K/group_size ue8m0 bytes (fp32 → ue8m0 conversion). + uint8_t* s_dst = (uint8_t*)w13_scale_ptrs[target_gpu]; + const size_t expert_s_off = is_up ? gpu_w13_scale_per_mat : 0; + avx2::fast_fp32_to_ue8m0( + s_dst + expert_s_off + (size_t)n_in_gpu * scales_per_row_w13, + bb->d + (size_t)row * scales_per_row_w13, + scales_per_row_w13); + } + } else { + // ---- W2 weight + scale: per-row + per-k-slice scatter ---- + const int chunk_idx = task_id - NUM_W13_TASKS * 2; + const auto& bb = down_bb_[expert_id]; + + const int rows_per_task = (cpu_n_w2 + NUM_W2_TASKS - 1) / NUM_W2_TASKS; + const int row_start = chunk_idx * rows_per_task; + const int row_end = std::min(row_start + rows_per_task, cpu_n_w2); + if (row_start >= cpu_n_w2) return; + + for (int row = row_start; row < row_end; row++) { + // CPU's K range = [global_k_offset_w2, global_k_offset_w2 + cpu_k_w2). + // Each gpu_k_w2-aligned slice within this range goes to its own + // target_gpu. Loop covers cpu_k_w2/gpu_k_w2 GPU TPs (or 1 if + // cpu_k_w2 < gpu_k_w2 — i.e. cpu_tp_count > gpu_tp_count case). + for (int k_start = 0; k_start < cpu_k_w2; k_start += gpu_k_w2) { + const int k_slice_len = std::min(gpu_k_w2, cpu_k_w2 - k_start); + const int global_k = global_k_offset_w2 + k_start; + const int target_gpu = global_k / gpu_k_w2; + const int k_in_gpu = global_k % gpu_k_w2; + + // Weight K-slice + uint8_t* w_dst = (uint8_t*)w2_weight_ptrs[target_gpu]; + std::memcpy(w_dst + (size_t)row * gpu_k_w2 + k_in_gpu, + bb->b + (size_t)row * cpu_k_w2 + k_start, + k_slice_len); + + // Scale K-slice (k_slice_len/group_size ue8m0 bytes) + const int scale_slice_len = k_slice_len / group_size; + const int k_in_gpu_scale = k_in_gpu / group_size; + const int k_start_scale = k_start / group_size; + uint8_t* s_dst = (uint8_t*)w2_scale_ptrs[target_gpu]; + avx2::fast_fp32_to_ue8m0( + s_dst + (size_t)row * gpu_scales_per_row_w2 + k_in_gpu_scale, + bb->d + (size_t)row * cpu_scales_per_row_w2 + k_start_scale, + scale_slice_len); + } + } + } + }, + nullptr); + } +}; + +// ============================================================================ +// TP_MOE specialization — handles per-expert FP8 weight + ue8m0 scale loading +// across TP parts. Mirrors AVX2_FP8_MOE_TP's TP_MOE but with MXFP8 layout: +// weights: [E, N, K] FP8 bytes (no nibble pack, no block scale) +// scales: [E, N, K/group_size] ue8m0 bytes (NOT pre-converted) +// +// Inside per-TP load_weights() we point config to the staged TP-sliced bytes, +// then derived_class::load_weights() reads from there and runs ue8m0→FP32. +// ============================================================================ +template +class TP_MOE> : public TP_MOE>> { + public: + using Base = TP_MOE>>; + using Base::Base; + + void load_weights() override { + auto& config = this->config; + auto& tps = this->tps; + auto pool = config.pool; + const uint64_t* physical_to_logical_map = (const uint64_t*)config.physical_to_logical_map; + + const int group_size = config.quant_config.group_size; + if (group_size == 0 || config.quant_config.zero_point) { + throw std::runtime_error("MXFP8 MoE only supports group-wise (group_size > 0, zero_point=false)"); + } + + if (config.gate_projs.empty() && config.gate_proj == nullptr) { + throw std::runtime_error("no weight source"); + } + const bool use_per_expert_ptrs = !config.gate_projs.empty(); + + // Full dimensions + const size_t full_weight_elems = (size_t)config.intermediate_size * config.hidden_size; + const size_t full_scale_elems = full_weight_elems / (size_t)group_size; // ue8m0 = 1 byte each + + pool->dispense_backend()->do_numa_job([&, this](int i) { + auto& tpc = tps[i]->config_; + const size_t tp_weight_elems = (size_t)tpc.intermediate_size * tpc.hidden_size; + const size_t tp_scale_elems = tp_weight_elems / (size_t)group_size; + + // Allocate temporary buffers for TP-sliced FP8 + ue8m0 scales + tpc.gate_proj = new uint8_t[tpc.expert_num * tp_weight_elems]; + tpc.up_proj = new uint8_t[tpc.expert_num * tp_weight_elems]; + tpc.down_proj = new uint8_t[tpc.expert_num * tp_weight_elems]; + tpc.gate_scale = new uint8_t[tpc.expert_num * tp_scale_elems]; + tpc.up_scale = new uint8_t[tpc.expert_num * tp_scale_elems]; + tpc.down_scale = new uint8_t[tpc.expert_num * tp_scale_elems]; + + // gate/up: split N=intermediate, each expert is [intermediate, hidden] FP8 + [intermediate, hidden/gs] ue8m0 + const size_t gate_up_w_src_off = i * tp_weight_elems; // bytes + const size_t gate_up_s_src_off = i * tp_scale_elems; // bytes (ue8m0=1B) + + // down: split K=intermediate (columns of [hidden, intermediate] FP8 + [hidden, intermediate/gs] ue8m0) + const size_t down_col_off = (size_t)i * tpc.intermediate_size; + const size_t down_scale_col_off = down_col_off / (size_t)group_size; + + pool->get_subpool(i)->do_work_stealing_job( + tpc.expert_num, nullptr, + [&, &tpc](int expert_id_) { + const size_t expert_id = expert_map(physical_to_logical_map, expert_id_); + + uint8_t* gate_dst = (uint8_t*)tpc.gate_proj + expert_id * tp_weight_elems; + uint8_t* up_dst = (uint8_t*)tpc.up_proj + expert_id * tp_weight_elems; + uint8_t* down_dst = (uint8_t*)tpc.down_proj + expert_id * tp_weight_elems; + uint8_t* gate_s_dst = (uint8_t*)tpc.gate_scale + expert_id * tp_scale_elems; + uint8_t* up_s_dst = (uint8_t*)tpc.up_scale + expert_id * tp_scale_elems; + uint8_t* down_s_dst = (uint8_t*)tpc.down_scale + expert_id * tp_scale_elems; + + const uint8_t* gate_src; + const uint8_t* up_src; + const uint8_t* down_src; + const uint8_t* gate_s_src; + const uint8_t* up_s_src; + const uint8_t* down_s_src; + + if (use_per_expert_ptrs) { + gate_src = (const uint8_t*)config.gate_projs[0][expert_id] + gate_up_w_src_off; + up_src = (const uint8_t*)config.up_projs[0][expert_id] + gate_up_w_src_off; + down_src = (const uint8_t*)config.down_projs[0][expert_id]; + gate_s_src = (const uint8_t*)config.gate_scales[0][expert_id] + gate_up_s_src_off; + up_s_src = (const uint8_t*)config.up_scales[0][expert_id] + gate_up_s_src_off; + down_s_src = (const uint8_t*)config.down_scales[0][expert_id]; + } else { + gate_src = (const uint8_t*)config.gate_proj + expert_id * full_weight_elems + gate_up_w_src_off; + up_src = (const uint8_t*)config.up_proj + expert_id * full_weight_elems + gate_up_w_src_off; + down_src = (const uint8_t*)config.down_proj + expert_id * full_weight_elems; + gate_s_src = (const uint8_t*)config.gate_scale + expert_id * full_scale_elems + gate_up_s_src_off; + up_s_src = (const uint8_t*)config.up_scale + expert_id * full_scale_elems + gate_up_s_src_off; + down_s_src = (const uint8_t*)config.down_scale + expert_id * full_scale_elems; + } + + // gate/up weights + scales: contiguous N slice + std::memcpy(gate_dst, gate_src, tp_weight_elems); + std::memcpy(up_dst, up_src, tp_weight_elems); + std::memcpy(gate_s_dst, gate_s_src, tp_scale_elems); + std::memcpy(up_s_dst, up_s_src, tp_scale_elems); + + // down weights: column slice within each of `hidden` rows + // src row = [hidden_row, full_intermediate]; dst row = [hidden_row, tp_intermediate] + for (int row = 0; row < config.hidden_size; row++) { + std::memcpy(down_dst + (size_t)row * tpc.intermediate_size, + down_src + (size_t)row * config.intermediate_size + down_col_off, + tpc.intermediate_size); + } + // down scales: column slice within each of `hidden` scale rows + const int full_kg = config.intermediate_size / group_size; + const int tp_kg = tpc.intermediate_size / group_size; + for (int row = 0; row < config.hidden_size; row++) { + std::memcpy(down_s_dst + (size_t)row * tp_kg, + down_s_src + (size_t)row * full_kg + down_scale_col_off, + tp_kg); + } + }, + nullptr); + }); + + // Call per-TP load_weights (does FP8 memcpy into BufferB.b + ue8m0→FP32 into BufferB.d) + pool->dispense_backend()->do_numa_job([&, this](int i) { + tps[i]->load_weights(); + }); + + // Free temporary buffers + pool->dispense_backend()->do_numa_job([&, this](int i) { + auto& tpc = tps[i]->config_; + delete[] (uint8_t*)tpc.gate_proj; + delete[] (uint8_t*)tpc.up_proj; + delete[] (uint8_t*)tpc.down_proj; + delete[] (uint8_t*)tpc.gate_scale; + delete[] (uint8_t*)tpc.up_scale; + delete[] (uint8_t*)tpc.down_scale; + }); + + this->weights_loaded = true; + } + + // -------------------------------------------------------------------------- + // write_weight_scale_to_buffer: orchestrator for layerwise prefill. + // Called once per expert by Python (kt-sglang/.../kt_ep_wrapper.py:_prepare_weight_fp8). + // Dispatches across all NUMA TP parts; each part runs its own + // AVX2_MXFP8_MOE_TP::write_weights_to_buffer to mirror its slice of the expert + // into the pre-allocated GPU pinned staging buffer. + // + // Mirrors amx/mxfp8-moe.hpp:897-912. + // SFINAE in ext_bindings.cpp:435 auto-detects this method and exposes + // `moe.write_weight_scale_to_buffer_task(...)` to Python — no manual binding needed. + // -------------------------------------------------------------------------- + void write_weight_scale_to_buffer(int gpu_tp_count, int expert_id, + const std::vector& w13_weight_ptrs, + const std::vector& w13_scale_ptrs, + const std::vector& w2_weight_ptrs, + const std::vector& w2_scale_ptrs) { + if (!this->weights_loaded) throw std::runtime_error("Not Loaded"); + if (this->tps.empty()) throw std::runtime_error("No TP parts initialized"); + if (w13_weight_ptrs.size() != (size_t)gpu_tp_count || + w13_scale_ptrs.size() != (size_t)gpu_tp_count || + w2_weight_ptrs.size() != (size_t)gpu_tp_count || + w2_scale_ptrs.size() != (size_t)gpu_tp_count) + throw std::runtime_error("Pointer arrays size must match gpu_tp_count"); + + this->config.pool->dispense_backend()->do_numa_job([&, this](int i) { + this->tps[i]->write_weights_to_buffer(gpu_tp_count, this->tp_count, expert_id, this->config, + w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs); + }); + } +}; + +#endif // CPUINFER_OPERATOR_AVX2_MXFP8_MOE_H diff --git a/kt-kernel/operators/common.hpp b/kt-kernel/operators/common.hpp index da7bbd4b..86f3464c 100644 --- a/kt-kernel/operators/common.hpp +++ b/kt-kernel/operators/common.hpp @@ -319,6 +319,11 @@ struct GeneralMOEConfig { // Origin: kt-sglang 耦合 (carries the V4-2604B limit set by sglang side). float swiglu_limit = 0.0f; + // MiniMax M3 "swigluoai" activation: gate * sigmoid(gate * alpha) * (up + 1). + // When alpha > 0, act_fn uses the swigluoai formula with symmetric clamp on + // both gate and up (±swiglu_limit). 0.0f = disabled (standard silu path). + float swiglu_alpha = 0.0f; + GeneralMOEConfig() {} GeneralMOEConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size) diff --git a/kt-kernel/python/experts.py b/kt-kernel/python/experts.py index 2e99bf33..318fe317 100644 --- a/kt-kernel/python/experts.py +++ b/kt-kernel/python/experts.py @@ -41,6 +41,7 @@ INFERENCE_METHODS = frozenset( "FP8_PERCHANNEL", # Per-channel FP8 "GPTQ_INT4", # GPTQ INT4 "MXFP4", # MXFP4 (E2M1 nibble + ue8m0 group scale, e.g. DeepSeek-V4-Flash routed experts) + "MXFP8", # MXFP8 (E4M3fn byte + ue8m0 group scale, e.g. MiniMax-M3-Preview) "LLAMAFILE", # GGUF format "MOE_INT4", "MOE_INT8", # General kernel @@ -149,6 +150,9 @@ class KTMoEWrapper: # _apply_swiglu_limit). Plumbed into MOEConfig.swiglu_limit and # consumed by amx::act_fn. Origin: kt-sglang 耦合. swiglu_limit: float = 0.0, + # MiniMax M3 swigluoai sigmoid alpha. 0.0 = standard silu (default). + # Non-zero triggers gate * sigmoid(gate * alpha) * (up + 1) in act_fn. + swiglu_alpha: float = 0.0, ): """ Factory method to create the appropriate backend implementation. @@ -220,6 +224,7 @@ class KTMoEWrapper: method=method, numa_nodes=numa_nodes, swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, ) else: # mode == "sft" # SFT factory does not plumb swiglu_limit; reject non-zero @@ -316,6 +321,7 @@ def _create_inference_wrapper( method: str, numa_nodes: Optional[List[int]] = None, swiglu_limit: float = 0.0, + swiglu_alpha: float = 0.0, ) -> BaseMoEWrapper: """ Create an inference wrapper based on the method. @@ -329,7 +335,7 @@ def _create_inference_wrapper( # Select backend based on method if method in ["AMXINT4", "AMXINT8"]: backend_cls = AMXMoEWrapper - elif method in ["RAWINT4", "FP8", "BF16", "FP8_PERCHANNEL", "GPTQ_INT4", "MXFP4"]: + elif method in ["RAWINT4", "FP8", "BF16", "FP8_PERCHANNEL", "GPTQ_INT4", "MXFP4", "MXFP8"]: backend_cls = NativeMoEWrapper elif method == "LLAMAFILE": backend_cls = LlamafileMoEWrapper @@ -347,15 +353,16 @@ def _create_inference_wrapper( # into a non-MXFP4 backend; act_fn would then clamp gate/up to ±10 with # no warning. Gate strictly on method instead. Origin: kt-sglang 耦合. extra_kwargs = {} - if method == "MXFP4": + if method in ("MXFP4", "MXFP8"): extra_kwargs["swiglu_limit"] = swiglu_limit + extra_kwargs["swiglu_alpha"] = swiglu_alpha elif swiglu_limit != 0.0: raise ValueError( - f"swiglu_limit={swiglu_limit} is only supported on method='MXFP4', " + f"swiglu_limit={swiglu_limit} is only supported on method='MXFP4'/'MXFP8', " f"got method={method!r} (backend={backend_cls.__name__}). This " f"usually means SGLANG_DSV4_2604_SUBMODE=2604B is set in the " f"environment while the current launch does not actually use " - f"MXFP4 weights — either unset the env or pass --kt-method MXFP4." + f"MXFP4/MXFP8 weights — either unset the env or pass --kt-method MXFP4/MXFP8." ) return backend_cls( layer_idx=layer_idx, diff --git a/kt-kernel/python/utils/amx.py b/kt-kernel/python/utils/amx.py index 97cb9955..1b04e4de 100644 --- a/kt-kernel/python/utils/amx.py +++ b/kt-kernel/python/utils/amx.py @@ -16,6 +16,7 @@ from .loader import ( BF16SafeTensorLoader, GPTQSafeTensorLoader, MXFP4SafeTensorLoader, + MXFP8SafeTensorLoader, ) from kt_kernel_ext.moe import MOEConfig import kt_kernel_ext.moe as _moe_mod @@ -24,6 +25,7 @@ AMXInt4_MOE = getattr(_moe_mod, "AMXInt4_MOE", None) AMXInt8_MOE = getattr(_moe_mod, "AMXInt8_MOE", None) AMXInt4_KGroup_MOE = getattr(_moe_mod, "AMXInt4_KGroup_MOE", None) AMXFP4_KGroup_MOE = getattr(_moe_mod, "AMXFP4_KGroup_MOE", None) +AMXMXFP8_KGroup_MOE = getattr(_moe_mod, "AMXMXFP8_KGroup_MOE", None) AMXFP8_MOE = getattr(_moe_mod, "AMXFP8_MOE", None) AMXBF16_MOE = getattr(_moe_mod, "AMXBF16_MOE", None) AMXFP8PerChannel_MOE = getattr(_moe_mod, "AMXFP8PerChannel_MOE", None) @@ -32,6 +34,7 @@ AVX2FP8_MOE = getattr(_moe_mod, "AVX2FP8_MOE", None) AVX2GPTQInt4_MOE = getattr(_moe_mod, "AVX2GPTQInt4_MOE", None) AVX2RawInt4_MOE = getattr(_moe_mod, "AVX2RawInt4_MOE", None) AVX2MXFP4_MOE = getattr(_moe_mod, "AVX2MXFP4_MOE", None) +AVX2MXFP8_MOE = getattr(_moe_mod, "AVX2MXFP8_MOE", None) AVXVNNI256GPTQInt4_MOE = getattr(_moe_mod, "AVXVNNI256GPTQInt4_MOE", None) AVXVNNI256RawInt4_MOE = getattr(_moe_mod, "AVXVNNI256RawInt4_MOE", None) @@ -39,6 +42,7 @@ _HAS_AMXINT4_SUPPORT = AMXInt4_MOE is not None _HAS_AMXINT8_SUPPORT = AMXInt8_MOE is not None _HAS_RAWINT4_SUPPORT = AMXInt4_KGroup_MOE is not None _HAS_MXFP4_SUPPORT = AMXFP4_KGroup_MOE is not None +_HAS_MXFP8_SUPPORT = AMXMXFP8_KGroup_MOE is not None _HAS_FP8_SUPPORT = AMXFP8_MOE is not None _HAS_BF16_SUPPORT = AMXBF16_MOE is not None _HAS_FP8_PERCHANNEL_SUPPORT = AMXFP8PerChannel_MOE is not None @@ -47,6 +51,7 @@ _HAS_AVX2_FP8_SUPPORT = AVX2FP8_MOE is not None _HAS_AVX2_GPTQ_INT4_SUPPORT = AVX2GPTQInt4_MOE is not None _HAS_AVX2_RAWINT4_SUPPORT = AVX2RawInt4_MOE is not None _HAS_AVX2_MXFP4_SUPPORT = AVX2MXFP4_MOE is not None +_HAS_AVX2_MXFP8_SUPPORT = AVX2MXFP8_MOE is not None _HAS_AVXVNNI256_GPTQ_INT4_SUPPORT = AVXVNNI256GPTQInt4_MOE is not None _HAS_AVXVNNI256_RAW_INT4_SUPPORT = AVXVNNI256RawInt4_MOE is not None _AVXVNNI256_GPTQ_INT4_MAX_GROUP_SIZE = 256 @@ -176,6 +181,44 @@ def _select_mxfp4_backend(): return None +def _select_mxfp8_backend(): + """Select MXFP8 backend: AMX/AVX-512 (preferred) > AVX2 (fallback). + + Override with KT_MXFP8_BACKEND=avx2|amx. + Returns None if no MXFP8 backend is available. + """ + forced = os.getenv("KT_MXFP8_BACKEND", "").strip().lower() + + if forced == "amx": + if not _HAS_MXFP8_SUPPORT: + raise RuntimeError( + "KT_MXFP8_BACKEND=amx requested, but AMXMXFP8_KGroup_MOE is not compiled in. " + "Recompile with AVX512F + AVX512BW + AVX512_BF16 + AVX512_VBMI enabled." + ) + if not _host_has_cpu_flag("amx_tile", "amx_bf16"): + raise RuntimeError( + "KT_MXFP8_BACKEND=amx requested, but the host CPU lacks AMX (amx_tile / amx_bf16). " + "This would SIGILL at first forward. Unset the env to fall back to AVX2." + ) + return AMXMXFP8_KGroup_MOE + + if forced == "avx2": + if not _HAS_AVX2_MXFP8_SUPPORT: + raise RuntimeError( + "KT_MXFP8_BACKEND=avx2 requested, but AVX2MXFP8_MOE is not compiled in. " + "Recompile with AVX2 + FMA enabled." + ) + return AVX2MXFP8_MOE + + # Auto-select: prefer AMX iff the .so was built with it AND the runtime CPU has AMX. + # Compile-time-only check would SIGILL on AVX-512 CPUs lacking AMX (pre-Sapphire Rapids). + if _HAS_MXFP8_SUPPORT and _host_has_cpu_flag("amx_tile", "amx_bf16"): + return AMXMXFP8_KGroup_MOE + if _HAS_AVX2_MXFP8_SUPPORT: + return AVX2MXFP8_MOE + return None + + class AMXMoEWrapper(BaseMoEWrapper): """ AMX-based MoE wrapper implementation. @@ -499,14 +542,16 @@ class NativeMoEWrapper(BaseMoEWrapper): method: str = "RAWINT4", numa_nodes: Optional[List[int]] = None, swiglu_limit: float = 0.0, + swiglu_alpha: float = 0.0, ): - # Defence in depth: reject swiglu_limit on non-MXFP4 methods even + self._swiglu_alpha = float(swiglu_alpha) + # Defence in depth: reject swiglu_limit on non-MXFP4/MXFP8 methods even # if the experts.py guard is bypassed (e.g., by a future caller # that constructs NativeMoEWrapper directly). Origin: kt-sglang 耦合. - if swiglu_limit != 0.0 and method != "MXFP4": + if swiglu_limit != 0.0 and method not in ("MXFP4", "MXFP8"): raise ValueError( f"NativeMoEWrapper received swiglu_limit={swiglu_limit} with " - f"method={method!r}; the V4-2604B clamp only applies to MXFP4. " + f"method={method!r}; the clamp only applies to MXFP4/MXFP8. " f"This indicates a missing guard in the caller." ) if method == "RAWINT4" and not ( @@ -552,6 +597,13 @@ class NativeMoEWrapper(BaseMoEWrapper): " - AVX2 + FMA (for AVX2 fallback backend)\n" "Please recompile kt_kernel_ext with one of the above enabled." ) + if method == "MXFP8" and not (_HAS_MXFP8_SUPPORT or _HAS_AVX2_MXFP8_SUPPORT): + raise RuntimeError( + "MXFP8 backend not available. Required ISA (any one of):\n" + " - AVX512F + AVX512BW + AVX512_BF16 + AVX512_VBMI (for AMX/AVX-512 backend)\n" + " - AVX2 + FMA (for AVX2 fallback backend)\n" + "Please recompile kt_kernel_ext with one of the above enabled." + ) super().__init__( layer_idx=layer_idx, @@ -596,6 +648,8 @@ class NativeMoEWrapper(BaseMoEWrapper): return GPTQSafeTensorLoader(weight_path) elif method == "MXFP4": return MXFP4SafeTensorLoader(weight_path) + elif method == "MXFP8": + return MXFP8SafeTensorLoader(weight_path) else: raise NotImplementedError(f"Unsupported method for NativeMoEWrapper: {method}") @@ -645,13 +699,22 @@ class NativeMoEWrapper(BaseMoEWrapper): self.loader = NativeMoEWrapper._native_loader_instance t0 = time.time() - base_key = f"model.layers.{self.layer_idx}" - try: - weights = self.loader.load_experts(base_key) - except (ValueError, KeyError): - # For VL/multimodal models (e.g. Qwen3.5) with 'language_model' prefix - base_key = f"model.language_model.layers.{self.layer_idx}" - weights = self.loader.load_experts(base_key) + _candidates = [ + f"model.layers.{self.layer_idx}", + f"language_model.model.layers.{self.layer_idx}", + f"model.language_model.layers.{self.layer_idx}", + ] + weights = None + for base_key in _candidates: + try: + weights = self.loader.load_experts(base_key) + break + except (ValueError, KeyError): + continue + if weights is None: + raise ValueError( + f"No experts found for layer {self.layer_idx} under any prefix: {_candidates}" + ) t1 = time.time() # Keep individual tensors instead of stacking - avoid expensive memory copy @@ -692,6 +755,9 @@ class NativeMoEWrapper(BaseMoEWrapper): # ue8m0 is losslessly representable in bf16 (8-bit exponent, 0 mantissa); # the loader has already done that conversion. assert self.gate_scales[0].dtype == torch.bfloat16, "Expected bf16 scales for MXFP4" + elif self.method == "MXFP8": + # ue8m0 scales stay as uint8; C++ convert_ue8m0_to_fp32 handles conversion. + assert self.gate_scales[0].dtype == torch.uint8, "Expected uint8 (ue8m0) scales for MXFP8" t2 = time.time() @@ -729,11 +795,11 @@ class NativeMoEWrapper(BaseMoEWrapper): # experts.py and the __init__ guards still cannot apply the clamp # on RAWINT4 / FP8 / BF16 / FP8_PERCHANNEL / GPTQ_INT4 paths. # Origin: kt-sglang 耦合. - if self.swiglu_limit != 0.0 and self.method != "MXFP4": + if self.swiglu_limit != 0.0 and self.method not in ("MXFP4", "MXFP8"): raise ValueError( f"NativeMoEWrapper.load_weights: swiglu_limit=" f"{self.swiglu_limit} with method={self.method!r}; clamp is " - f"only valid for MXFP4." + f"only valid for MXFP4/MXFP8." ) moe_config.swiglu_limit = self.swiglu_limit @@ -777,6 +843,21 @@ class NativeMoEWrapper(BaseMoEWrapper): "Compile with AVX512_BF16 (AMXFP4_KGroup_MOE) or AVX2 (AVX2MXFP4_MOE)." ) self.moe = backend_cls(moe_config) + elif self.method == "MXFP8": + # MXFP8: FP8 E4M3fn byte weights, ue8m0/uint8 per-32 group scale + # (e.g. MiniMax-M3-Preview) + group_size = self.hidden_size // self.gate_scales[0].shape[1] + moe_config.quant_config.bits = 8 + moe_config.quant_config.group_size = group_size + moe_config.quant_config.zero_point = False + moe_config.swiglu_alpha = getattr(self, "_swiglu_alpha", 0.0) + backend_cls = _select_mxfp8_backend() + if backend_cls is None: + raise RuntimeError( + "No MXFP8 backend available after runtime selection. " + "Compile with AVX512+VBMI (AMXMXFP8_KGroup_MOE) or AVX2 (AVX2MXFP8_MOE)." + ) + self.moe = backend_cls(moe_config) elif self.method == "FP8": moe_config.quant_config.bits = 8 moe_config.quant_config.group_size = 128 diff --git a/kt-kernel/python/utils/loader.py b/kt-kernel/python/utils/loader.py index b85e4c80..25178229 100644 --- a/kt-kernel/python/utils/loader.py +++ b/kt-kernel/python/utils/loader.py @@ -1222,3 +1222,82 @@ class MXFP4SafeTensorLoader(SafeTensorLoader): "up_scale": up_scales, "down_scale": down_scales, } + + +class MXFP8SafeTensorLoader(SafeTensorLoader): + """Loader for native MXFP8 expert weights (MiniMax M3 Preview format). + + Per expert layout: + {base}.block_sparse_moe.experts.{i}.w1.weight F8_E4M3 [N, K] gate + {base}.block_sparse_moe.experts.{i}.w1.weight_scale_inv U8 [N, K/32] ue8m0 + {base}.block_sparse_moe.experts.{i}.w3.{weight,weight_scale_inv} up + {base}.block_sparse_moe.experts.{i}.w2.{weight,weight_scale_inv} down + + M3 keys are prefixed with ``language_model.model.layers.{L}``; we also probe + the stripped form. Scales stay as uint8 — the C++ kernel converts ue8m0→FP32 + via bit-shift during load_weights. + """ + + EXPERTS_PATH_TPL = "{base}.block_sparse_moe.experts" + PROJ_NAMES = ("w1", "w3", "w2") # (gate, up, down) + + def _experts_prefix_candidates(self, base_key: str) -> list[str]: + candidates = [self.EXPERTS_PATH_TPL.format(base=base_key)] + for strip in ("language_model.model.", "language_model.", "model."): + if base_key.startswith(strip): + candidates.append(self.EXPERTS_PATH_TPL.format(base=base_key[len(strip):])) + return list(dict.fromkeys(candidates)) + + def load_experts(self, base_key: str, device: str = "cpu"): + gate_name, up_name, down_name = self.PROJ_NAMES + prefix = None + expert_count = 0 + for cand in self._experts_prefix_candidates(base_key): + expert_count = 0 + while self.has_tensor(f"{cand}.{expert_count}.{gate_name}.weight"): + expert_count += 1 + if expert_count > 0: + prefix = cand + break + if prefix is None: + raise ValueError( + f"No MXFP8 experts found under any of: {self._experts_prefix_candidates(base_key)}" + ) + + gate_weights = [None] * expert_count + up_weights = [None] * expert_count + down_weights = [None] * expert_count + gate_scales = [None] * expert_count + up_scales = [None] * expert_count + down_scales = [None] * expert_count + + for exp_id in range(expert_count): + for proj, dst in ( + (gate_name, gate_weights), + (up_name, up_weights), + (down_name, down_weights), + ): + w = self.load_tensor(f"{prefix}.{exp_id}.{proj}.weight", device).contiguous() + if w.dtype != torch.uint8: + w = w.view(torch.uint8) + dst[exp_id] = w + + for proj, dst in ( + (gate_name, gate_scales), + (up_name, up_scales), + (down_name, down_scales), + ): + s = self.load_tensor(f"{prefix}.{exp_id}.{proj}.weight_scale_inv", device).contiguous() + if s.dtype != torch.uint8: + s = s.view(torch.uint8) + dst[exp_id] = s + + print(f"[MXFP8SafeTensorLoader] Loaded {expert_count} experts from {prefix}") + return { + "gate": gate_weights, + "up": up_weights, + "down": down_weights, + "gate_scale": gate_scales, + "up_scale": up_scales, + "down_scale": down_scales, + }