[feat] MXFP8 MoE support (#2041)

Add MXFP8 kernel for Minimax-M3-MXFP8 Day 0

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())

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())

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<AMX_FP8_MOE_TP<amx::GemmKernel224FP8>>(moe_module, "AMXFP8_MOE");
   bind_moe_module<AMX_FP8_PERCHANNEL_MOE_TP<amx::GemmKernel224FP8PerChannel>>(moe_module, "AMXFP8PerChannel_MOE");
   bind_moe_module<AMX_FP4_MOE_TP<amx::GemmKernel224MXFP4SmallKGroup>>(moe_module, "AMXFP4_KGroup_MOE");
+  bind_moe_module<AMX_MXFP8_MOE_TP<amx::GemmKernel224MXFP8SmallKGroup>>(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<AVX2_GPTQ_INT4_MOE_TP<avx2::GemmKernelAVX2GPTQInt4>>(moe_module, "AVX2GPTQInt4_MOE");
   bind_moe_module<AVX2_RAW_INT4_MOE_TP<avx2::GemmKernelAVX2RawInt4>>(moe_module, "AVX2RawInt4_MOE");
   bind_moe_module<AVX2_MXFP4_MOE_TP<avx2::GemmKernelAVX2MXFP4>>(moe_module, "AVX2MXFP4_MOE");
+  bind_moe_module<AVX2_MXFP8_MOE_TP<avx2::GemmKernelAVX2MXFP8>>(moe_module, "AVX2MXFP8_MOE");
   bind_moe_module<AVXVNNI256_GPTQ_INT4_MOE_TP<avxvnni::GemmKernelAVXVNNI256GPTQInt4>>(moe_module,
                                                                                       "AVXVNNI256GPTQInt4_MOE");
   bind_moe_module<AVXVNNI256_RAW_INT4_MOE_TP<avxvnni_rawint4::GemmKernelAVXVNNI256RawInt4>>(moe_module,

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) {                                                        \

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));
         }
       }

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 <typename K>
+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<size_t>(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<intptr_t>(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<uint8_t*>(ptr);
+    // Scale region starts after weight data
+    d = reinterpret_cast<float*>(b + static_cast<size_t>(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<size_t>(k);  // 1 byte per element
+    const size_t rows = static_cast<size_t>(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<size_t>(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<size_t>(n_begin) * k_group_count + k_begin / k_group_size;
+  }
+
+  static std::pair<int, int> 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<int, int> 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<GemmKernel224MXFP8SmallKGroup>;
+  using BufferB = BufferBMXFP8KGroupImpl<GemmKernel224MXFP8SmallKGroup>;
+  using BufferC = BufferCReduceImpl<GemmKernel224MXFP8SmallKGroup>;
+
+  __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<GemmKernel224MXFP8SmallKGroup::BufferA> ba,
+                                 std::shared_ptr<GemmKernel224MXFP8SmallKGroup::BufferB> bb,
+                                 std::shared_ptr<GemmKernel224MXFP8SmallKGroup::BufferC> 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<GemmKernel224MXFP8SmallKGroup::BufferA> ba,
+                                 std::shared_ptr<GemmKernel224MXFP8SmallKGroup::BufferB> bb,
+                                 std::shared_ptr<GemmKernel224MXFP8SmallKGroup::BufferC> 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 T = amx::GemmKernel224MXFP8SmallKGroup>
+class AMX_MXFP8_MOE_TP : public AMX_MOE_BASE<T, AMX_MXFP8_MOE_TP<T>> {
+  using Base = AMX_MOE_BASE<T, AMX_MXFP8_MOE_TP<T>>;
+  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<typename T::BufferA> make_buffer_a_impl(size_t m, size_t k, void* data) const {
+    return std::make_shared<typename T::BufferA>(m, k, data);
+  }
+  std::shared_ptr<typename T::BufferB> make_buffer_b_impl(size_t n, size_t k, void* data) const {
+    return std::make_shared<typename T::BufferB>(n, k, config_.quant_config.group_size, data);
+  }
+  std::shared_ptr<typename T::BufferC> make_buffer_c_impl(size_t m, size_t n, void* data) const {
+    return std::make_shared<typename T::BufferC>(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<uint32_t>(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<size_t>(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<uintptr_t>& w13_weight_ptrs,
+                               const std::vector<uintptr_t>& w13_scale_ptrs,
+                               const std::vector<uintptr_t>& w2_weight_ptrs,
+                               const std::vector<uintptr_t>& 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 <typename K>
+class TP_MOE<AMX_MXFP8_MOE_TP<K>> : public TP_MOE<AMX_MOE_BASE<K, AMX_MXFP8_MOE_TP<K>>> {
+ public:
+  using Base = TP_MOE<AMX_MOE_BASE<K, AMX_MXFP8_MOE_TP<K>>>;
+  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<uintptr_t>& w13_weight_ptrs,
+                                    const std::vector<uintptr_t>& w13_scale_ptrs,
+                                    const std::vector<uintptr_t>& w2_weight_ptrs,
+                                    const std::vector<uintptr_t>& 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

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

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);
+          }
         }
       }
     };

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 <cassert>
+#include <cstdint>
+#include <cstring>
+
+#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<uint32_t>(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<int, int> 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 T = avx2::GemmKernelAVX2MXFP8>
+class AVX2_MXFP8_MOE_TP : public AVX2_MOE_BASE<T, AVX2_MXFP8_MOE_TP<T>> {
+  using Base = AVX2_MOE_BASE<T, AVX2_MXFP8_MOE_TP<T>>;
+  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<typename T::BufferA> make_buffer_a_impl(size_t m, size_t k, void* data) const {
+    return std::make_shared<typename T::BufferA>(m, k, data);
+  }
+  std::shared_ptr<typename T::BufferB> make_buffer_b_impl(size_t n, size_t k, void* data) const {
+    return std::make_shared<typename T::BufferB>(n, k, config_.quant_config.group_size, data);
+  }
+  std::shared_ptr<typename T::BufferC> make_buffer_c_impl(size_t m, size_t n, void* data) const {
+    return std::make_shared<typename T::BufferC>(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<uintptr_t>& w13_weight_ptrs,
+                               const std::vector<uintptr_t>& w13_scale_ptrs,
+                               const std::vector<uintptr_t>& w2_weight_ptrs,
+                               const std::vector<uintptr_t>& 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 <typename K>
+class TP_MOE<AVX2_MXFP8_MOE_TP<K>> : public TP_MOE<AVX2_MOE_BASE<K, AVX2_MXFP8_MOE_TP<K>>> {
+ public:
+  using Base = TP_MOE<AVX2_MOE_BASE<K, AVX2_MXFP8_MOE_TP<K>>>;
+  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<uintptr_t>& w13_weight_ptrs,
+                                    const std::vector<uintptr_t>& w13_scale_ptrs,
+                                    const std::vector<uintptr_t>& w2_weight_ptrs,
+                                    const std::vector<uintptr_t>& 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

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)

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,

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

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,
+        }