diff --git a/kt-kernel/examples/test_fp4_moe_avx2.py b/kt-kernel/examples/test_fp4_moe_avx2.py new file mode 100644 index 00000000..4bf7f0d2 --- /dev/null +++ b/kt-kernel/examples/test_fp4_moe_avx2.py @@ -0,0 +1,236 @@ +import os +import sys +from typing import Dict + +sys.path.insert(0, os.path.dirname(__file__) + "/../build") + +import torch +from kt_kernel import kt_kernel_ext + +torch.manual_seed(42) + +hidden_size = 7168 +intermediate_size = 2048 +max_len = 25600 + +expert_num = 16 +num_experts_per_tok = 8 + +layer_num = 1 +CPUInfer = kt_kernel_ext.CPUInfer(40) +validation_iter = 3 +k_group_size = 32 +debug_print_count = 16 + +QLEN_LIST = [1, 32] +DISPATCH_THRESHOLD = 4 * expert_num / num_experts_per_tok + +physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous() + +E2M1_VALUES = torch.tensor([ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, +], dtype=torch.float32) + + +def act_fn(x): + return x / (1.0 + torch.exp(-x)) + + +def mlp_torch(input, gate_proj, up_proj, down_proj): + gate_buf = torch.mm(input, gate_proj.t()) + up_buf = torch.mm(input, up_proj.t()) + intermediate = act_fn(gate_buf) * up_buf + return torch.mm(intermediate, down_proj.t()) + + +def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj): + cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num)) + cnts.scatter_(1, expert_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = expert_ids.view(-1).argsort() + sorted_tokens = input[idxs // expert_ids.shape[1]] + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert_out = mlp_torch(sorted_tokens[start_idx:end_idx], gate_proj[i], up_proj[i], down_proj[i]) + outputs.append(expert_out) + start_idx = end_idx + outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0) + new_x = torch.empty_like(outs) + new_x[idxs] = outs + return ( + new_x.view(*expert_ids.shape, -1) + .type(weights.dtype) + .mul_(weights.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + + +def quantize_mxfp4_tensor(weights: torch.Tensor, group_size: int): + weights_f32 = weights.to(torch.float32) + e, rows, cols = weights_f32.shape + reshaped = weights_f32.view(e, rows, cols // group_size, group_size) + max_abs = reshaped.abs().amax(dim=-1, keepdim=True) + max_abs = torch.clamp(max_abs, min=1e-8) + scales = (max_abs / 6.0).squeeze(-1) + normalized = reshaped / scales.unsqueeze(-1) + e2m1_vals = E2M1_VALUES.view(1, 1, 1, 1, 16) + normalized_expanded = normalized.unsqueeze(-1) + distances = torch.abs(normalized_expanded - e2m1_vals) + closest_indices = distances.argmin(dim=-1) + dequant = E2M1_VALUES[closest_indices].to(torch.float32) * scales.unsqueeze(-1) + dequant = dequant.view(e, rows, cols) + nibbles = closest_indices.to(torch.uint8) + nibbles = nibbles.view(e, rows, cols // 2, 2) + lo = nibbles[..., 0] + hi = nibbles[..., 1] + packed_bytes = (hi << 4) | lo + bytes_view = packed_bytes.view(e, rows, cols // 8, 4) + packed_int32 = ( + bytes_view[..., 0].to(torch.int32) | + (bytes_view[..., 1].to(torch.int32) << 8) | + (bytes_view[..., 2].to(torch.int32) << 16) | + (bytes_view[..., 3].to(torch.int32) << 24) + ) + packed_int32 = packed_int32.view(e, rows, cols // 8).contiguous() + scales = scales.to(torch.bfloat16).contiguous().view(e, rows, cols // group_size).contiguous() + return packed_int32, scales, dequant + + +WEIGHT_PATTERNS = { + "uniform_scale": ("All k-groups share the same abs max / scale", lambda g: torch.full((g,), 0.02, dtype=torch.float32)), + "alternating_scale": ("Alternate small / large abs max per k-group", lambda g: torch.where(torch.arange(g) % 2 == 0, torch.full((g,), 0.015), torch.full((g,), 0.03))), + "ramp_scale": ("Linearly increasing abs max per k-group", lambda g: torch.linspace(0.005, 0.04, steps=g, dtype=torch.float32)), + "random": ("Random bf16 weights (baseline)", None), +} + + +def build_structured_tensor(shape, pattern): + if pattern == "random": + torch.manual_seed(42) + return (torch.randn(shape, dtype=torch.bfloat16) / 100.0).contiguous() + e, rows, cols = shape + groups = cols // k_group_size + group_vals = WEIGHT_PATTERNS[pattern][1](groups).to(torch.float32) + block = group_vals.view(1, 1, groups, 1).expand(e, rows, groups, k_group_size).clone() + row_signs = torch.where( + (torch.arange(rows) % 2 == 0), + torch.ones(rows, dtype=torch.float32), + -torch.ones(rows, dtype=torch.float32), + ).view(1, rows, 1, 1) + col_offsets = torch.linspace(-0.0005, 0.0005, steps=k_group_size, dtype=torch.float32).view(1, 1, 1, k_group_size) + block = block * row_signs + col_offsets + return block.reshape(shape).to(torch.bfloat16).contiguous() + + +def prepare_weights(pattern): + gate_proj = build_structured_tensor((expert_num, intermediate_size, hidden_size), pattern) + up_proj = build_structured_tensor((expert_num, intermediate_size, hidden_size), pattern) + down_proj = build_structured_tensor((expert_num, hidden_size, intermediate_size), pattern) + gate_q, gate_s, gate_dq = quantize_mxfp4_tensor(gate_proj, k_group_size) + up_q, up_s, up_dq = quantize_mxfp4_tensor(up_proj, k_group_size) + down_q, down_s, down_dq = quantize_mxfp4_tensor(down_proj, k_group_size) + return { + "gate_qweight": gate_q.contiguous(), "up_qweight": up_q.contiguous(), "down_qweight": down_q.contiguous(), + "gate_scales": gate_s.contiguous(), "up_scales": up_s.contiguous(), "down_scales": down_s.contiguous(), + "dequantized": {"gate_proj": gate_dq.to(torch.bfloat16), "up_proj": up_dq.to(torch.bfloat16), "down_proj": down_dq.to(torch.bfloat16)}, + } + + +def build_moes(quant_data): + AVX2MXFP4_MOE = getattr(kt_kernel_ext.moe, "AVX2MXFP4_MOE", None) + if AVX2MXFP4_MOE is None: + raise RuntimeError("AVX2MXFP4_MOE not found — rebuild kt-kernel with the AVX2 MXFP4 path") + moes = [] + with torch.inference_mode(mode=True): + for _ in range(layer_num): + config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size, 0) + config.max_len = max_len + config.quant_config.bits = 4 + config.quant_config.group_size = k_group_size + config.quant_config.zero_point = False + config.gate_proj = quant_data["gate_qweight"].data_ptr() + config.up_proj = quant_data["up_qweight"].data_ptr() + config.down_proj = quant_data["down_qweight"].data_ptr() + config.gate_scale = quant_data["gate_scales"].data_ptr() + config.up_scale = quant_data["up_scales"].data_ptr() + config.down_scale = quant_data["down_scales"].data_ptr() + config.pool = CPUInfer.backend_ + moe = AVX2MXFP4_MOE(config) + CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr())) + CPUInfer.sync() + moes.append(moe) + return moes + + +def run_case(pattern, qlen): + print("\n" + "=" * 70) + desc = WEIGHT_PATTERNS[pattern][0] + path = "mat-vec" if qlen <= DISPATCH_THRESHOLD else "mat-mat" + print(f"Running case: {pattern} -> {desc} (qlen={qlen}, path={path})") + print("=" * 70) + + quant_data = prepare_weights(pattern) + moes = build_moes(quant_data) + dq = quant_data["dequantized"] + + diffs = [] + with torch.inference_mode(mode=True): + for i in range(validation_iter): + torch.manual_seed(100 + i) + bsz_tensor = torch.tensor([qlen], device="cpu") + expert_ids = torch.stack( + [torch.randperm(expert_num)[:num_experts_per_tok] for _ in range(qlen)] + ).contiguous() + weights = torch.randn((qlen, num_experts_per_tok), dtype=torch.float32).contiguous() + input_tensor = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous() / 100 + output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous() + + moe = moes[i % layer_num] + CPUInfer.submit(moe.forward_task( + bsz_tensor.data_ptr(), num_experts_per_tok, + expert_ids.data_ptr(), weights.data_ptr(), + input_tensor.data_ptr(), output.data_ptr(), False, + )) + CPUInfer.sync() + + t_output = moe_torch(input_tensor.to(torch.bfloat16), expert_ids, weights, + dq["gate_proj"], dq["up_proj"], dq["down_proj"]).to(torch.bfloat16) + + diff = torch.mean(torch.abs(output.float() - t_output.float())) / ( + torch.mean(torch.abs(t_output.float())) + 1e-12) + diffs.append(diff.item()) + print(f"[{pattern}] iter {i}: rel-L1 = {diff:.4f}") + print(f" output {output.flatten()[:debug_print_count]}") + print(f" t_output {t_output.flatten()[:debug_print_count]}") + + return {"case": pattern, "description": desc, + "mean": sum(diffs)/len(diffs) if diffs else 0.0, + "max": max(diffs) if diffs else 0.0, + "min": min(diffs) if diffs else 0.0} + + +def run_fp4_moe_avx2_test(): + summary = [] + for qlen in QLEN_LIST: + path = "mat-vec" if qlen <= DISPATCH_THRESHOLD else "mat-mat" + print(f"\n##### qlen={qlen} path={path} #####") + for pattern in WEIGHT_PATTERNS: + r = run_case(pattern, qlen) + r.update({"qlen": qlen, "path": path}) + summary.append(r) + + print("\n=== AVX2 MXFP4 — Relative Error Summary ===") + print(f"{'Case':<20} {'qlen':>5} {'path':<8} {'Mean':>10} {'Max':>10} {'Min':>10}") + for r in summary: + print(f"{r['case']:<20} {r['qlen']:>5} {r['path']:<8} " + f"{r['mean']*100:9.2f}% {r['max']*100:9.2f}% {r['min']*100:9.2f}%") + + +if __name__ == "__main__": + run_fp4_moe_avx2_test() diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 66f6fd31..81ca2577 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -59,6 +59,7 @@ static const bool _is_plain_ = false; #include "operators/avx2/fp8-moe.hpp" #include "operators/avx2/gptq_int4-moe.hpp" #include "operators/avx2/gptq_int4_avxvnni-moe.hpp" +#include "operators/avx2/mxfp4-moe.hpp" #include "operators/avx2/rawint4-moe.hpp" #include "operators/avx2/rawint4_avxvnni-moe.hpp" #endif @@ -821,6 +822,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { bind_moe_module>(moe_module, "AVX2FP8_MOE"); bind_moe_module>(moe_module, "AVX2GPTQInt4_MOE"); bind_moe_module>(moe_module, "AVX2RawInt4_MOE"); + bind_moe_module>(moe_module, "AVX2MXFP4_MOE"); bind_moe_module>(moe_module, "AVXVNNI256GPTQInt4_MOE"); bind_moe_module>(moe_module, diff --git a/kt-kernel/operators/avx2/moe_base.hpp b/kt-kernel/operators/avx2/moe_base.hpp index c61a9367..088fcf7d 100644 --- a/kt-kernel/operators/avx2/moe_base.hpp +++ b/kt-kernel/operators/avx2/moe_base.hpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -61,6 +62,8 @@ class AVX2_MOE_BASE { std::vector> down_bb_; std::vector> down_bc_; + std::vector owned_aligned_allocs_; + size_t pool_count_ = 0; size_t gate_up_ba_pool_bytes_ = 0; size_t gate_bc_pool_bytes_ = 0; @@ -117,14 +120,21 @@ class AVX2_MOE_BASE { down_bc_.push_back(make_buffer_c(config_.max_len, config_.hidden_size, nullptr)); void* gate_bb_ptr = - std::aligned_alloc(64, buffer_b_required_size(config_.intermediate_size, config_.hidden_size)); + std::aligned_alloc(64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL); + if (!gate_bb_ptr) throw std::runtime_error("aligned_alloc failed for gate BufferB"); + owned_aligned_allocs_.push_back(gate_bb_ptr); gate_bb_.push_back(make_buffer_b(config_.intermediate_size, config_.hidden_size, gate_bb_ptr)); - void* up_bb_ptr = std::aligned_alloc(64, buffer_b_required_size(config_.intermediate_size, config_.hidden_size)); + void* up_bb_ptr = + std::aligned_alloc(64, (buffer_b_required_size(config_.intermediate_size, config_.hidden_size) + 63) & ~63ULL); + if (!up_bb_ptr) throw std::runtime_error("aligned_alloc failed for up BufferB"); + owned_aligned_allocs_.push_back(up_bb_ptr); up_bb_.push_back(make_buffer_b(config_.intermediate_size, config_.hidden_size, up_bb_ptr)); void* down_bb_ptr = - std::aligned_alloc(64, buffer_b_required_size(config_.hidden_size, config_.intermediate_size)); + std::aligned_alloc(64, (buffer_b_required_size(config_.hidden_size, config_.intermediate_size) + 63) & ~63ULL); + if (!down_bb_ptr) throw std::runtime_error("aligned_alloc failed for down BufferB"); + owned_aligned_allocs_.push_back(down_bb_ptr); down_bb_.push_back(make_buffer_b(config_.hidden_size, config_.intermediate_size, down_bb_ptr)); } @@ -145,7 +155,9 @@ class AVX2_MOE_BASE { shared_mem_buffer_numa.alloc(tp_part_idx, this, mem_requests); } - ~AVX2_MOE_BASE() = default; + ~AVX2_MOE_BASE() { + for (void* p : owned_aligned_allocs_) std::free(p); + } void warm_up() { int qlen = config_.max_len; diff --git a/kt-kernel/operators/avx2/mxfp4-moe.hpp b/kt-kernel/operators/avx2/mxfp4-moe.hpp new file mode 100644 index 00000000..394301ca --- /dev/null +++ b/kt-kernel/operators/avx2/mxfp4-moe.hpp @@ -0,0 +1,839 @@ +/** + * @Description : AVX2 MXFP4 MoE operator for DeepSeek V4 native inference + * SPDX-License-Identifier: Apache-2.0 + * + * MXFP4 stores FP4 E2M1 weights nibble-packed (2 values per byte), plus + * per-group-32 FP32 scales. This AVX2 backend dequantizes FP4→FP32 via SSSE3 + * PSHUFB lookup tables, then accumulates with BF16 activations using + * _mm256_fmadd_ps. No AVX-512 required. + * + * FP4 E2M1 values: {0, ±0.5, ±1.0, ±1.5, ±2.0, ±3.0, ±4.0, ±6.0} + **/ +#ifndef CPUINFER_OPERATOR_AVX2_MXFP4_MOE_H +#define CPUINFER_OPERATOR_AVX2_MXFP4_MOE_H + +#include +#include + +#include +#include +#include +#include +#include + +#include "avx2_bf16_gemm.hpp" +#include "avx2_bf16_utils.hpp" +#include "moe_base.hpp" + +namespace avx2 { + +// ============================================================================ +// AVX2 MXFP4 GemmKernel +// Weights: FP4 E2M1 nibble-packed, per-group-32 FP32 scales +// Activations: BF16; Output: FP32 +// ============================================================================ +struct GemmKernelAVX2MXFP4 { + using dt = uint8_t; + using output_t = float; + static constexpr int M_STEP = 1; + static constexpr int N_STEP = 8; + static constexpr int K_STEP = 32; + static constexpr int N_BLOCK = 64; + static constexpr int K_BLOCK = 128; + static constexpr double ELEMENT_SIZE = 0.5; + + static void config() {} + static int recommended_nth(int n) { return std::max(1, div_up(n, N_BLOCK)); } + static std::pair split_range_n(int n, int ith, int nth) { return split_range(n, ith, nth); } + static std::string name() { return "AVX2_MXFP4"; } + + // FP4 E2M1 → BF16 lookup tables (low byte and high byte of BF16) + // Identical to amx/fp4-moe.hpp GemmKernel224MXFP4SmallKGroup LUTs + alignas(16) static constexpr uint8_t fp4_bf16_lo[16] = {0x00, 0x00, 0x80, 0xC0, 0x00, 0x40, 0x80, 0xC0, + 0x00, 0x00, 0x80, 0xC0, 0x00, 0x40, 0x80, 0xC0}; + alignas(16) static constexpr uint8_t fp4_bf16_hi[16] = {0x00, 0x3F, 0x3F, 0x3F, 0x40, 0x40, 0x40, 0x40, + 0x80, 0xBF, 0xBF, 0xBF, 0xC0, 0xC0, 0xC0, 0xC0}; + + // Dequantize 4 packed bytes (= 8 FP4 E2M1 nibbles) → 8 FP32 values. + // + // Nibble order: packed[j] = {lo_nib = column 2j, hi_nib = column 2j+1} + // Output order: columns 0, 1, 2, ..., 7 (sequential). + // + // Algorithm: SSSE3 PSHUFB table lookup → 8 BF16 values → BF16-to-FP32 + // (zero-extend uint16 to uint32, shift left 16). + __attribute__((always_inline)) static inline __m256 fp4x8_to_fp32(const uint8_t* packed, __m128i lut_lo, + __m128i lut_hi) { + int32_t raw; + std::memcpy(&raw, packed, 4); + __m128i data = _mm_cvtsi32_si128(raw); + __m128i lo_mask = _mm_set1_epi8(0x0F); + + // Extract lo nibbles (bits 0-3) and hi nibbles (bits 4-7) + __m128i lo_nib = _mm_and_si128(data, lo_mask); + __m128i hi_nib = _mm_and_si128(_mm_srli_epi16(data, 4), lo_mask); + + // PSHUFB: BF16 byte lookup for each set of nibbles + __m128i lo_bf16 = _mm_unpacklo_epi8(_mm_shuffle_epi8(lut_lo, lo_nib), + _mm_shuffle_epi8(lut_hi, lo_nib)); // 4 BF16: cols {0,2,4,6} + __m128i hi_bf16 = _mm_unpacklo_epi8(_mm_shuffle_epi8(lut_lo, hi_nib), + _mm_shuffle_epi8(lut_hi, hi_nib)); // 4 BF16: cols {1,3,5,7} + + // Interleave: {col0,col1,col2,...,col7} in sequential order + __m128i bf16_8 = _mm_unpacklo_epi16(lo_bf16, hi_bf16); // 8 × BF16 in 128 bits + + // BF16 → FP32: zero-extend u16 → u32, then shift left 16 + return _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(bf16_8), 16)); + } + + // Scalar dequant of a single nibble index (for tail handling) + static inline float fp4_scalar(uint8_t nib) { + const uint32_t bf16_bits = (uint32_t)fp4_bf16_lo[nib] | ((uint32_t)fp4_bf16_hi[nib] << 8); + const uint32_t fp32_bits = bf16_bits << 16; + float result; + std::memcpy(&result, &fp32_bits, sizeof(float)); + return result; + } + + // ---- Buffer types -------------------------------------------------------- + + struct BufferA { + ggml_bf16_t* data = nullptr; + size_t max_m = 0, k = 0; + + BufferA() = default; + BufferA(size_t m, size_t k_, int /*group_size*/, void* ptr) : data((ggml_bf16_t*)ptr), max_m(m), k(k_) {} + + static size_t required_size(size_t m, size_t k, int /*group_size*/) { 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] = 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)); + } + } + }; + + struct BufferB { + uint8_t* b = nullptr; // nibble-packed FP4 (may be nullptr in scale-only mode) + float* d = nullptr; // FP32 group scales + int n = 0, k = 0, k_group_size = 0, k_group_count = 0; + + BufferB() = default; + + // Full allocation: b and d packed into a single aligned block. + BufferB(int n_, int k_, int k_group_size_, void* ptr) + : b((uint8_t*)ptr), n(n_), k(k_), k_group_size(k_group_size_) { + if (k_group_size <= 0 || k % k_group_size != 0 || k % 8 != 0) + throw std::runtime_error("MXFP4 AVX2 requires k % group_size == 0 and k % 8 == 0"); + k_group_count = k / k_group_size; + d = (float*)((uint8_t*)ptr + (size_t)n * k / 2); + } + + // Scale-only: b is set externally from mmap'd safetensor data. + BufferB(int n_, int k_, int k_group_size_, void* scale_ptr, std::nullptr_t) + : b(nullptr), n(n_), k(k_), k_group_size(k_group_size_) { + if (k_group_size <= 0 || k % k_group_size != 0 || k % 8 != 0) + throw std::runtime_error("MXFP4 AVX2 requires k % group_size == 0 and k % 8 == 0"); + k_group_count = k / k_group_size; + d = (float*)scale_ptr; + } + + static size_t required_size(size_t n, size_t k, int k_group_size) { + return n * k / 2 + n * (k / k_group_size) * sizeof(float); + } + static size_t required_size_scale_only(size_t n, size_t k, int k_group_size) { + return n * (k / k_group_size) * sizeof(float); + } + + void from_raw_mat(const uint8_t* proj, int ith, int nth) { + if (b == nullptr) return; + auto [n_start, n_end] = split_range(n, ith, nth); + const size_t row_bytes = (size_t)k / 2; + std::memcpy(b + n_start * row_bytes, proj + n_start * row_bytes, (size_t)(n_end - n_start) * row_bytes); + } + }; + + struct BufferC { + float* data = nullptr; + size_t max_m = 0, n = 0; + + BufferC() = default; + BufferC(size_t m, size_t n_, void* ptr) : data((float*)ptr), max_m(m), n(n_) {} + + 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] = split_range((int)n, ith, nth); + for (int mi = 0; mi < m; mi++) { + float* src_row = data + (size_t)mi * n; + ggml_bf16_t* dst_row = dst + (size_t)mi * n; + int j = n_start; + for (; j + 8 <= n_end; j += 8) store_fp32_to_bf16(dst_row + j, _mm256_loadu_ps(src_row + j)); + for (; j < n_end; j++) dst_row[j] = GGML_FP32_TO_BF16(src_row[j]); + } + } + }; +}; + +// ============================================================================ +// Compute: M tokens × N output neurons, K inner dimension, group_size scale +// +// N-outer loop: each weight row loaded once and shared across M tokens. +// 4-token M-blocking for prefill: weight decode amortized over 4 activations. +// ============================================================================ +static void gemm_mxfp4(int m, int n, int k, GemmKernelAVX2MXFP4::BufferA& a, GemmKernelAVX2MXFP4::BufferB& b, + GemmKernelAVX2MXFP4::BufferC& c, int ith, int nth) { + // H3: Check weight buffer is not null + if (b.b == nullptr) { + throw std::runtime_error("gemm_mxfp4: weight buffer (b.b) is null"); + } + auto [n_start, n_end] = split_range(n, ith, nth); + const int group_count = b.k_group_count; + const int group_size = b.k_group_size; + const size_t row_bytes = (size_t)k / 2; + + const __m128i lut_lo = _mm_load_si128((const __m128i*)GemmKernelAVX2MXFP4::fp4_bf16_lo); + const __m128i lut_hi = _mm_load_si128((const __m128i*)GemmKernelAVX2MXFP4::fp4_bf16_hi); + + for (int ni = n_start; ni < n_end; ni++) { + const uint8_t* b_row = b.b + (size_t)ni * row_bytes; + const float* b_scales = b.d + (size_t)ni * group_count; + + if (ni + 1 < n_end) { + _mm_prefetch((const char*)(b.b + (size_t)(ni + 1) * row_bytes), _MM_HINT_T0); + _mm_prefetch((const char*)(b.b + (size_t)(ni + 1) * row_bytes + 64), _MM_HINT_T0); + _mm_prefetch((const char*)(b.b + (size_t)(ni + 1) * row_bytes + 128), _MM_HINT_T0); + _mm_prefetch((const char*)(b.b + (size_t)(ni + 1) * row_bytes + 192), _MM_HINT_T0); + } + + // 4-token blocked path: decode weight once, fmadd into 4 token accumulators + int mi = 0; + for (; mi + 4 <= m; mi += 4) { + const ggml_bf16_t* a0 = a.data + (size_t)(mi + 0) * a.k; + const ggml_bf16_t* a1 = a.data + (size_t)(mi + 1) * a.k; + const ggml_bf16_t* a2 = a.data + (size_t)(mi + 2) * a.k; + const ggml_bf16_t* a3 = a.data + (size_t)(mi + 3) * a.k; + __m256 tot0 = _mm256_setzero_ps(), tot1 = _mm256_setzero_ps(); + __m256 tot2 = _mm256_setzero_ps(), tot3 = _mm256_setzero_ps(); + + for (int g = 0; g < group_count; g++) { + const int k_base = g * group_size; + __m256 g0 = _mm256_setzero_ps(), g1 = _mm256_setzero_ps(); + __m256 g2 = _mm256_setzero_ps(), g3 = _mm256_setzero_ps(); + + int ki = 0; + for (; ki + 32 <= group_size; ki += 32) { + const uint8_t* w = b_row + (k_base + ki) / 2; + __m256 wv0 = GemmKernelAVX2MXFP4::fp4x8_to_fp32(w + 0, lut_lo, lut_hi); + __m256 wv1 = GemmKernelAVX2MXFP4::fp4x8_to_fp32(w + 4, lut_lo, lut_hi); + __m256 wv2 = GemmKernelAVX2MXFP4::fp4x8_to_fp32(w + 8, lut_lo, lut_hi); + __m256 wv3 = GemmKernelAVX2MXFP4::fp4x8_to_fp32(w + 12, lut_lo, lut_hi); + + g0 = _mm256_fmadd_ps(load_bf16_to_fp32(a0 + k_base + ki), wv0, g0); + g0 = _mm256_fmadd_ps(load_bf16_to_fp32(a0 + k_base + ki + 8), wv1, g0); + g0 = _mm256_fmadd_ps(load_bf16_to_fp32(a0 + k_base + ki + 16), wv2, g0); + g0 = _mm256_fmadd_ps(load_bf16_to_fp32(a0 + k_base + ki + 24), wv3, g0); + + g1 = _mm256_fmadd_ps(load_bf16_to_fp32(a1 + k_base + ki), wv0, g1); + g1 = _mm256_fmadd_ps(load_bf16_to_fp32(a1 + k_base + ki + 8), wv1, g1); + g1 = _mm256_fmadd_ps(load_bf16_to_fp32(a1 + k_base + ki + 16), wv2, g1); + g1 = _mm256_fmadd_ps(load_bf16_to_fp32(a1 + k_base + ki + 24), wv3, g1); + + g2 = _mm256_fmadd_ps(load_bf16_to_fp32(a2 + k_base + ki), wv0, g2); + g2 = _mm256_fmadd_ps(load_bf16_to_fp32(a2 + k_base + ki + 8), wv1, g2); + g2 = _mm256_fmadd_ps(load_bf16_to_fp32(a2 + k_base + ki + 16), wv2, g2); + g2 = _mm256_fmadd_ps(load_bf16_to_fp32(a2 + k_base + ki + 24), wv3, g2); + + g3 = _mm256_fmadd_ps(load_bf16_to_fp32(a3 + k_base + ki), wv0, g3); + g3 = _mm256_fmadd_ps(load_bf16_to_fp32(a3 + k_base + ki + 8), wv1, g3); + g3 = _mm256_fmadd_ps(load_bf16_to_fp32(a3 + k_base + ki + 16), wv2, g3); + g3 = _mm256_fmadd_ps(load_bf16_to_fp32(a3 + k_base + ki + 24), wv3, g3); + } + for (; ki + 8 <= group_size; ki += 8) { + const uint8_t* w = b_row + (k_base + ki) / 2; + __m256 wv = GemmKernelAVX2MXFP4::fp4x8_to_fp32(w, lut_lo, lut_hi); + g0 = _mm256_fmadd_ps(load_bf16_to_fp32(a0 + k_base + ki), wv, g0); + g1 = _mm256_fmadd_ps(load_bf16_to_fp32(a1 + k_base + ki), wv, g1); + g2 = _mm256_fmadd_ps(load_bf16_to_fp32(a2 + k_base + ki), wv, g2); + g3 = _mm256_fmadd_ps(load_bf16_to_fp32(a3 + k_base + ki), wv, g3); + } + + __m256 sv = _mm256_broadcast_ss(&b_scales[g]); + tot0 = _mm256_fmadd_ps(g0, sv, tot0); + tot1 = _mm256_fmadd_ps(g1, sv, tot1); + tot2 = _mm256_fmadd_ps(g2, sv, tot2); + tot3 = _mm256_fmadd_ps(g3, sv, tot3); + } + c.data[(size_t)(mi + 0) * n + ni] = hsum_avx2(tot0); + c.data[(size_t)(mi + 1) * n + ni] = hsum_avx2(tot1); + c.data[(size_t)(mi + 2) * n + ni] = hsum_avx2(tot2); + c.data[(size_t)(mi + 3) * n + ni] = hsum_avx2(tot3); + } + + // M-tail: single token fallback + for (; mi < m; mi++) { + const ggml_bf16_t* a_row = a.data + (size_t)mi * a.k; + __m256 total_acc = _mm256_setzero_ps(); + float scalar_tail = 0.0f; + + for (int g = 0; g < group_count; g++) { + const float scale = b_scales[g]; + const int k_base = g * group_size; + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + __m256 acc2 = _mm256_setzero_ps(), acc3 = _mm256_setzero_ps(); + + int ki = 0; + for (; ki + 32 <= group_size; ki += 32) { + const uint8_t* w = b_row + (k_base + ki) / 2; + acc0 = _mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki), + GemmKernelAVX2MXFP4::fp4x8_to_fp32(w + 0, lut_lo, lut_hi), acc0); + acc1 = _mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki + 8), + GemmKernelAVX2MXFP4::fp4x8_to_fp32(w + 4, lut_lo, lut_hi), acc1); + acc2 = _mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki + 16), + GemmKernelAVX2MXFP4::fp4x8_to_fp32(w + 8, lut_lo, lut_hi), acc2); + acc3 = _mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki + 24), + GemmKernelAVX2MXFP4::fp4x8_to_fp32(w + 12, lut_lo, lut_hi), acc3); + } + __m256 g_acc = _mm256_add_ps(_mm256_add_ps(acc0, acc2), _mm256_add_ps(acc1, acc3)); + + for (; ki + 8 <= group_size; ki += 8) { + const uint8_t* w = b_row + (k_base + ki) / 2; + g_acc = _mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki), + GemmKernelAVX2MXFP4::fp4x8_to_fp32(w, lut_lo, lut_hi), g_acc); + } + + total_acc = _mm256_fmadd_ps(g_acc, _mm256_broadcast_ss(&scale), total_acc); + + for (; ki < group_size; ki++) { + const int pos = k_base + ki; + const uint8_t packed = b_row[pos / 2]; + const uint8_t nib = (pos & 1) ? (packed >> 4) : (packed & 0x0F); + scalar_tail += GGML_BF16_TO_FP32(a_row[pos]) * GemmKernelAVX2MXFP4::fp4_scalar(nib) * scale; + } + } + + c.data[(size_t)mi * n + ni] = hsum_avx2(total_acc) + scalar_tail; + } + } +} + +} // namespace avx2 + +// ============================================================================ +// CRTP MoE wrapper — mirrors AVX2_RAW_INT4_MOE_TP +// ============================================================================ +template +class AVX2_MXFP4_MOE_TP : public AVX2_MOE_BASE> { + using Base = AVX2_MOE_BASE>; + using Base::config_; + using Base::down_ba_; + using Base::down_bb_; + using Base::down_bc_; + using Base::gate_bb_; + using Base::gate_bc_; + using Base::gate_up_ba_; + using Base::m_local_num_; + using Base::tp_part_idx; + using Base::up_bb_; + using Base::up_bc_; + + public: + using typename Base::input_t; + using typename Base::output_t; + + AVX2_MXFP4_MOE_TP() = default; + AVX2_MXFP4_MOE_TP(GeneralMOEConfig config, int tp_part_idx_ = 0) : Base(config, tp_part_idx_) {} + + void derived_init() { + if (config_.quant_config.group_size == 0 || config_.quant_config.zero_point) + throw std::runtime_error("MXFP4 AVX2 MoE requires KGroup FP4 without zero point"); + printf("Created AVX2_MXFP4_MOE_TP %d at numa %d (group_size=%d)\n", tp_part_idx, numa_node_of_cpu(sched_getcpu()), + config_.quant_config.group_size); + } + + size_t buffer_a_required_size_impl(size_t m, size_t k) const { + return T::BufferA::required_size(m, k, config_.quant_config.group_size); + } + size_t buffer_b_required_size_impl(size_t n, size_t k) const { + // Always allocate full buffer (weights + scales) for both per-expert and flat-buffer modes. + // Per-expert mode used to only allocate scales and point b into mmap, but that caused + // use-after-free when Python releases the mmap after load_weights(). + return T::BufferB::required_size(n, k, config_.quant_config.group_size); + } + size_t buffer_c_required_size_impl(size_t m, size_t n) const { return T::BufferC::required_size(m, n); } + + std::shared_ptr make_buffer_a_impl(size_t m, size_t k, void* data) const { + return std::make_shared(m, k, config_.quant_config.group_size, data); + } + std::shared_ptr make_buffer_b_impl(size_t n, size_t k, void* data) const { + // Always create full BufferB with owned weight buffer for both modes. + return std::make_shared((int)n, (int)k, config_.quant_config.group_size, data); + } + std::shared_ptr make_buffer_c_impl(size_t m, size_t n, void* data) const { + return std::make_shared(m, n, data); + } + + void do_gate_up_gemm(bool do_up, int expert_idx, int ith, int nth, int) { + int m = m_local_num_[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_mxfp4(m, config_.intermediate_size, config_.hidden_size, *gate_up_ba_[expert_idx], *bb, *bc, ith, nth); + } + + void do_down_gemm(int expert_idx, int ith, int nth, int) { + int m = m_local_num_[expert_idx]; + avx2::gemm_mxfp4(m, config_.hidden_size, config_.intermediate_size, *down_ba_[expert_idx], *down_bb_[expert_idx], + *down_bc_[expert_idx], ith, nth); + } + + void load_weights() { + int group_size = config_.quant_config.group_size; + const uint64_t* physical_to_logical_map = (const uint64_t*)config_.physical_to_logical_map; + auto pool = config_.pool->get_subpool(tp_part_idx); + + const bool use_per_expert = !config_.gate_projs.empty(); + if (!use_per_expert && config_.gate_proj == nullptr) + throw std::runtime_error("MXFP4 AVX2 MoE requires weight pointers"); + if (!use_per_expert && config_.gate_scale == nullptr) + throw std::runtime_error("MXFP4 AVX2 MoE requires scale pointers"); + + if (use_per_expert) { + // C1: Per-expert mode with TP > 1 must use TP_MOE wrapper which has full_intermediate_size. + // Direct use of AVX2_MXFP4_MOE_TP with tp_part_idx > 0 in per-expert mode is not supported + // because we cannot compute correct offsets without knowing full_intermediate_size. + if (tp_part_idx > 0) { + throw std::runtime_error( + "AVX2_MXFP4_MOE_TP::load_weights() per-expert mode with tp_part_idx > 0 is not supported. " + "Use TP_MOE wrapper for multi-TP scenarios."); + } + + // Copy weights from mmap into owned buffers (aligned with AMX behavior). + // Previously used direct pointers into mmap, but Python releases mmap after load_weights(), + // causing use-after-free. Now we memcpy to decouple from mmap lifecycle. + pool->do_work_stealing_job( + config_.expert_num, nullptr, + [this, physical_to_logical_map](int expert_idx) { + if (expert_idx < 0 || expert_idx >= config_.expert_num || gate_bb_[expert_idx] == nullptr) return; + uint64_t lid = expert_map(physical_to_logical_map, expert_idx); + // H1 & H2: Validate source pointers and lid bounds + if (lid >= config_.gate_projs[0].size()) { + throw std::runtime_error("load_weights: lid " + std::to_string(lid) + + " out of bounds (size=" + std::to_string(config_.gate_projs[0].size()) + ")"); + } + if (config_.gate_projs[0][lid] == nullptr || config_.up_projs[0][lid] == nullptr || + config_.down_projs[0][lid] == nullptr) { + throw std::runtime_error("load_weights: null weight pointer for expert " + std::to_string(lid)); + } + // tp_part_idx == 0 guaranteed here, so offset is 0. Copy full expert weights. + gate_bb_[expert_idx]->from_raw_mat((const uint8_t*)config_.gate_projs[0][lid], 0, 1); + up_bb_[expert_idx]->from_raw_mat((const uint8_t*)config_.up_projs[0][lid], 0, 1); + down_bb_[expert_idx]->from_raw_mat((const uint8_t*)config_.down_projs[0][lid], 0, 1); + }, + nullptr); + + pool->do_work_stealing_job( + config_.expert_num, nullptr, + [this, physical_to_logical_map, group_size](int task_id) { + uint64_t expert_idx = task_id; + if (expert_idx >= (uint64_t)config_.expert_num || gate_bb_[expert_idx] == nullptr) return; + uint64_t lid = expert_map(physical_to_logical_map, expert_idx); + // H2: Bounds check (already validated in weight loading, but be safe) + if (lid >= config_.gate_scales[0].size()) return; + if (config_.gate_scales[0][lid] == nullptr || config_.up_scales[0][lid] == nullptr || + config_.down_scales[0][lid] == nullptr) return; + size_t scale_elem_count = ((size_t)config_.hidden_size * config_.intermediate_size) / group_size; + // tp_part_idx == 0 guaranteed here, so offset is 0 + convert_or_copy(gate_bb_[expert_idx]->d, (const ggml_bf16_t*)config_.gate_scales[0][lid], scale_elem_count); + convert_or_copy(up_bb_[expert_idx]->d, (const ggml_bf16_t*)config_.up_scales[0][lid], scale_elem_count); + convert_or_copy(down_bb_[expert_idx]->d, (const ggml_bf16_t*)config_.down_scales[0][lid], scale_elem_count); + }, + nullptr); + } else { + // Flat-buffer mode: copy TP-sliced weights, convert BF16 scales → FP32. + int nth_gate = T::recommended_nth(config_.intermediate_size); + pool->do_work_stealing_job( + nth_gate * config_.expert_num, nullptr, + [this, nth_gate, physical_to_logical_map](int task_id) { + uint64_t expert_idx = task_id / nth_gate; + if (config_.should_skip_expert(expert_idx)) return; + uint64_t lid = expert_map(physical_to_logical_map, expert_idx); + int ith = task_id % nth_gate; + size_t weight_offset = ((size_t)lid * config_.intermediate_size * config_.hidden_size) / 2; + gate_bb_[expert_idx]->from_raw_mat((const uint8_t*)config_.gate_proj + weight_offset, ith, nth_gate); + up_bb_[expert_idx]->from_raw_mat((const uint8_t*)config_.up_proj + weight_offset, ith, nth_gate); + }, + nullptr); + + int nth_down = T::recommended_nth(config_.hidden_size); + pool->do_work_stealing_job( + nth_down * config_.expert_num, nullptr, + [this, nth_down, physical_to_logical_map](int task_id) { + uint64_t expert_idx = task_id / nth_down; + if (config_.should_skip_expert(expert_idx)) return; + uint64_t lid = expert_map(physical_to_logical_map, expert_idx); + int ith = task_id % nth_down; + size_t weight_offset = ((size_t)lid * config_.hidden_size * config_.intermediate_size) / 2; + down_bb_[expert_idx]->from_raw_mat((const uint8_t*)config_.down_proj + weight_offset, ith, nth_down); + }, + nullptr); + + pool->do_work_stealing_job( + config_.expert_num, nullptr, + [this, physical_to_logical_map, group_size](int task_id) { + uint64_t expert_idx = task_id; + if (config_.should_skip_expert(expert_idx)) return; + uint64_t lid = expert_map(physical_to_logical_map, expert_idx); + size_t scale_elem_count = ((size_t)config_.hidden_size * config_.intermediate_size) / group_size; + convert_or_copy(gate_bb_[expert_idx]->d, (ggml_bf16_t*)config_.gate_scale + lid * scale_elem_count, + scale_elem_count); + convert_or_copy(up_bb_[expert_idx]->d, (ggml_bf16_t*)config_.up_scale + lid * scale_elem_count, + scale_elem_count); + convert_or_copy(down_bb_[expert_idx]->d, (ggml_bf16_t*)config_.down_scale + lid * scale_elem_count, + scale_elem_count); + }, + nullptr); + } + } + + static inline void fp32_to_bf16(ggml_bf16_t* dst, const float* src, size_t count) { + convert_or_copy(dst, src, count); + } + + void write_weights_to_buffer(int gpu_tp_count, int cpu_tp_count, int expert_id, const GeneralMOEConfig& full_config, + const std::vector& w13_weight_ptrs, + const std::vector& w13_scale_ptrs, + const std::vector& w2_weight_ptrs, + const std::vector& w2_scale_ptrs) const { + if (expert_id < 0 || expert_id >= config_.expert_num || gate_bb_[expert_id] == nullptr || + up_bb_[expert_id] == nullptr || down_bb_[expert_id] == nullptr) + throw std::runtime_error("MXFP4 AVX2 write_weights_to_buffer: invalid expert"); + + const int group_size = config_.quant_config.group_size; + auto pool = config_.pool->get_subpool(tp_part_idx); + + size_t cpu_tp_weight_bytes = (size_t)config_.intermediate_size * config_.hidden_size / 2; + size_t cpu_tp_scale_elem_count = (size_t)config_.intermediate_size * config_.hidden_size / group_size; + size_t gpu_tp_weight_bytes = (size_t)full_config.intermediate_size * full_config.hidden_size / gpu_tp_count / 2; + size_t gpu_tp_scale_elem_count = + (size_t)full_config.intermediate_size * full_config.hidden_size / gpu_tp_count / group_size; + + if (cpu_tp_count >= gpu_tp_count) { + int target_gpu_tp = tp_part_idx / (cpu_tp_count / gpu_tp_count); + int local_idx = tp_part_idx % (cpu_tp_count / gpu_tp_count); + uint8_t* w13_wdst = (uint8_t*)w13_weight_ptrs[target_gpu_tp]; + ggml_bf16_t* w13_sdst = (ggml_bf16_t*)w13_scale_ptrs[target_gpu_tp]; + uint8_t* w2_wdst = (uint8_t*)w2_weight_ptrs[target_gpu_tp]; + ggml_bf16_t* w2_sdst = (ggml_bf16_t*)w2_scale_ptrs[target_gpu_tp]; + size_t w_off = (size_t)local_idx * cpu_tp_weight_bytes; + size_t s_off = (size_t)local_idx * cpu_tp_scale_elem_count; + + constexpr int NWT = 8; + int ndt = std::min(std::max(1, config_.hidden_size / 128), 32); + size_t wchunk = ((cpu_tp_weight_bytes + NWT - 1) / NWT + 63) & ~63ULL; + pool->do_work_stealing_job( + NWT * 2 + ndt + 2, nullptr, + [=, this](int tid) { + if (tid < NWT) { + size_t s = (size_t)tid * wchunk, e = std::min(s + wchunk, cpu_tp_weight_bytes); + if (s < e) std::memcpy(w13_wdst + w_off + s, gate_bb_[expert_id]->b + s, e - s); + } else if (tid < NWT * 2) { + size_t s = (size_t)(tid - NWT) * wchunk, e = std::min(s + wchunk, cpu_tp_weight_bytes); + if (s < e) std::memcpy(w13_wdst + w_off + gpu_tp_weight_bytes + s, up_bb_[expert_id]->b + s, e - s); + } else if (tid < NWT * 2 + ndt) { + size_t cols_per = (config_.hidden_size + ndt - 1) / ndt; + size_t cs = (size_t)(tid - NWT * 2) * cols_per; + size_t ce = std::min(cs + cols_per, (size_t)config_.hidden_size); + size_t wpc = config_.intermediate_size >> 1; + size_t spc = config_.intermediate_size / group_size; + size_t gws = (full_config.intermediate_size / gpu_tp_count) >> 1; + size_t gss = (full_config.intermediate_size / gpu_tp_count) / group_size; + size_t gwo = (size_t)local_idx * wpc, gso = (size_t)local_idx * spc; + for (size_t col = cs; col < ce; col++) { + std::memcpy(w2_wdst + col * gws + gwo, down_bb_[expert_id]->b + col * wpc, wpc); + fp32_to_bf16(w2_sdst + col * gss + gso, down_bb_[expert_id]->d + col * spc, spc); + } + } else if (tid == NWT * 2 + ndt) { + fp32_to_bf16(w13_sdst + s_off, gate_bb_[expert_id]->d, cpu_tp_scale_elem_count); + } else { + fp32_to_bf16(w13_sdst + s_off + gpu_tp_scale_elem_count, up_bb_[expert_id]->d, cpu_tp_scale_elem_count); + } + }, + nullptr); + } else { + // cpu_tp_count < gpu_tp_count: one CPU partition feeds multiple GPU TPs + int gpu_per_cpu = gpu_tp_count / cpu_tp_count; + int start_gtp = tp_part_idx * gpu_per_cpu; + size_t dw = cpu_tp_weight_bytes / gpu_per_cpu; + size_t ds = cpu_tp_scale_elem_count / gpu_per_cpu; + constexpr int NWT = 8; + int ndt = std::min(std::max(1, config_.hidden_size / 128), 32); + size_t wchunk = ((dw + NWT - 1) / NWT + 63) & ~63ULL; + int tpg = NWT * 2 + ndt + 2; + pool->do_work_stealing_job( + tpg * gpu_per_cpu, nullptr, + [=, this, &w13_weight_ptrs, &w13_scale_ptrs, &w2_weight_ptrs, &w2_scale_ptrs](int task_id) { + int lg = task_id / tpg; + int tt = task_id % tpg; + int gtp = start_gtp + lg; + uint8_t* w13_wdst = (uint8_t*)w13_weight_ptrs[gtp]; + ggml_bf16_t* w13_sdst = (ggml_bf16_t*)w13_scale_ptrs[gtp]; + uint8_t* w2_wdst = (uint8_t*)w2_weight_ptrs[gtp]; + ggml_bf16_t* w2_sdst = (ggml_bf16_t*)w2_scale_ptrs[gtp]; + size_t cow = (size_t)lg * dw, cos = (size_t)lg * ds; + if (tt < NWT) { + size_t s = (size_t)tt * wchunk, e = std::min(s + wchunk, dw); + if (s < e) std::memcpy(w13_wdst + s, gate_bb_[expert_id]->b + cow + s, e - s); + } else if (tt < NWT * 2) { + size_t s = (size_t)(tt - NWT) * wchunk, e = std::min(s + wchunk, dw); + if (s < e) std::memcpy(w13_wdst + gpu_tp_weight_bytes + s, up_bb_[expert_id]->b + cow + s, e - s); + } else if (tt < NWT * 2 + ndt) { + size_t cols_per = (config_.hidden_size + ndt - 1) / ndt; + size_t cs = (size_t)(tt - NWT * 2) * cols_per; + size_t ce = std::min(cs + cols_per, (size_t)config_.hidden_size); + size_t wgc = (config_.intermediate_size / gpu_per_cpu) >> 1; + size_t sgc = (config_.intermediate_size / gpu_per_cpu) / group_size; + for (size_t col = cs; col < ce; col++) { + size_t cwoff = col * config_.intermediate_size / 2 + lg * dw / config_.hidden_size; + size_t csoff = col * (config_.intermediate_size / group_size) + lg * ds / config_.hidden_size; + std::memcpy(w2_wdst + col * wgc, down_bb_[expert_id]->b + cwoff, wgc); + fp32_to_bf16(w2_sdst + col * sgc, down_bb_[expert_id]->d + csoff, sgc); + } + } else if (tt == NWT * 2 + ndt) { + fp32_to_bf16(w13_sdst, gate_bb_[expert_id]->d + cos, ds); + } else { + fp32_to_bf16(w13_sdst + gpu_tp_scale_elem_count, up_bb_[expert_id]->d + cos, ds); + } + }, + nullptr); + } + } +}; + +// TP_MOE specialization (boilerplate, identical to rawint4 pattern) +template +class TP_MOE> : public TP_MOE>> { + std::vector tp_owned_down_bufs_; + std::vector tp_owned_gate_bufs_; + std::vector tp_owned_up_bufs_; + + public: + using Base = TP_MOE>>; + using Base::Base; + + ~TP_MOE() { + for (void* p : tp_owned_down_bufs_) + if (p) std::free(p); + for (void* p : tp_owned_gate_bufs_) + if (p) std::free(p); + for (void* p : tp_owned_up_bufs_) + if (p) std::free(p); + } + + 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; + bool use_per_expert = !config.gate_projs.empty(); + if (config.gate_projs.empty() && config.gate_scale == nullptr) + throw std::runtime_error("MXFP4 AVX2 MoE only supports packed FP4 with KGroup scales"); + + if (use_per_expert) { + int group_size = config.quant_config.group_size; + int full_interm = config.intermediate_size; + + // Allocate per-partition repack buffers for all projections. + // Previously gate/up used direct mmap pointers, causing use-after-free when Python + // releases mmap. Now we memcpy all weights to decouple from mmap lifecycle (aligned with AMX). + tp_owned_down_bufs_.resize(this->tp_count, nullptr); + tp_owned_gate_bufs_.resize(this->tp_count, nullptr); + tp_owned_up_bufs_.resize(this->tp_count, nullptr); + pool->dispense_backend()->do_numa_job([&, this](int i) { + auto* tp = tps[i].get(); + auto& tpc = tp->config_; + tpc.physical_to_logical_map = config.physical_to_logical_map; + int per_tp_interm = tpc.intermediate_size; + if (per_tp_interm % 2 != 0) + throw std::runtime_error("MXFP4 TP: per_tp_interm must be even for nibble-aligned addressing, got " + + std::to_string(per_tp_interm)); + // gate/up: [intermediate × hidden] = [N × K], TP slices along N → contiguous rows + size_t gate_up_wt_per_expert = (size_t)per_tp_interm * tpc.hidden_size / 2; + size_t gate_up_alloc = ((size_t)tpc.expert_num * gate_up_wt_per_expert + 63) & ~(size_t)63; + uint8_t* gate_buf = (uint8_t*)std::aligned_alloc(64, gate_up_alloc); + uint8_t* up_buf = (uint8_t*)std::aligned_alloc(64, gate_up_alloc); + if (!gate_buf || !up_buf) throw std::runtime_error("aligned_alloc failed for MXFP4 gate/up_buf"); + tp_owned_gate_bufs_[i] = gate_buf; + tp_owned_up_bufs_[i] = up_buf; + // down: [hidden × intermediate], TP slices along K → non-contiguous + size_t down_wt_per_expert = (size_t)tpc.hidden_size * per_tp_interm / 2; + size_t alloc_size = ((size_t)tpc.expert_num * down_wt_per_expert + 63) & ~(size_t)63; + uint8_t* down_buf = (uint8_t*)std::aligned_alloc(64, alloc_size); + if (!down_buf) throw std::runtime_error("aligned_alloc failed for MXFP4 down_buf"); + tp_owned_down_bufs_[i] = down_buf; + + auto subpool = pool->get_subpool(i); + subpool->do_work_stealing_job( + tpc.expert_num, nullptr, + [&, i, per_tp_interm, full_interm, gate_buf, up_buf, down_buf, gate_up_wt_per_expert, down_wt_per_expert](int eid) { + if (tpc.should_skip_expert(eid)) return; + uint64_t lid = expert_map(physical_to_logical_map, eid); + + // H1 & H2: Validate source pointers and lid bounds + if (lid >= config.gate_projs[0].size()) { + throw std::runtime_error("TP_MOE load_weights: lid " + std::to_string(lid) + + " out of bounds (size=" + std::to_string(config.gate_projs[0].size()) + ")"); + } + if (config.gate_projs[0][lid] == nullptr || config.up_projs[0][lid] == nullptr || + config.down_projs[0][lid] == nullptr) { + throw std::runtime_error("TP_MOE load_weights: null weight pointer for expert " + std::to_string(lid)); + } + + // gate/up weights: N-split, contiguous rows → memcpy TP slice into owned buffer + // Each TP partition handles per_tp_interm rows starting at row i*per_tp_interm + size_t n_byte_off = (size_t)i * per_tp_interm * tpc.hidden_size / 2; + uint8_t* dst_gate = gate_buf + (size_t)eid * gate_up_wt_per_expert; + uint8_t* dst_up = up_buf + (size_t)eid * gate_up_wt_per_expert; + std::memcpy(dst_gate, (const uint8_t*)config.gate_projs[0][lid] + n_byte_off, gate_up_wt_per_expert); + std::memcpy(dst_up, (const uint8_t*)config.up_projs[0][lid] + n_byte_off, gate_up_wt_per_expert); + tp->gate_bb_[eid]->b = dst_gate; + tp->up_bb_[eid]->b = dst_up; + + // gate/up scales: contiguous → convert BF16→FP32 + size_t scale_count = (size_t)(tpc.hidden_size / group_size) * per_tp_interm; + size_t scale_off = (size_t)i * per_tp_interm * (tpc.hidden_size / group_size); + if (config.gate_scales[0][lid] == nullptr || config.up_scales[0][lid] == nullptr) { + throw std::runtime_error("TP_MOE load_weights: null scale pointer for expert " + std::to_string(lid)); + } + convert_or_copy(tp->gate_bb_[eid]->d, (const ggml_bf16_t*)config.gate_scales[0][lid] + scale_off, + scale_count); + convert_or_copy(tp->up_bb_[eid]->d, (const ggml_bf16_t*)config.up_scales[0][lid] + scale_off, + scale_count); + + // down weights: K-split, non-contiguous → per-row memcpy into repack buf + uint8_t* src_down = (uint8_t*)config.down_projs[0][lid]; + uint8_t* dst_down = down_buf + (size_t)eid * down_wt_per_expert; + for (int row = 0; row < tpc.hidden_size; row++) { + std::memcpy(dst_down + (size_t)row * per_tp_interm / 2, + src_down + (size_t)row * full_interm / 2 + (size_t)i * per_tp_interm / 2, + per_tp_interm / 2); + } + tp->down_bb_[eid]->b = dst_down; + + // down scales: K-split, non-contiguous → per-row convert + if (config.down_scales[0][lid] == nullptr) { + throw std::runtime_error("TP_MOE load_weights: null down_scale pointer for expert " + std::to_string(lid)); + } + int full_interm_groups = full_interm / group_size; + int per_tp_groups = per_tp_interm / group_size; + const ggml_bf16_t* src_ds = (const ggml_bf16_t*)config.down_scales[0][lid]; + float* dst_ds = tp->down_bb_[eid]->d; + for (int row = 0; row < tpc.hidden_size; row++) { + convert_or_copy(dst_ds + (size_t)row * per_tp_groups, + src_ds + (size_t)row * full_interm_groups + (size_t)i * per_tp_groups, per_tp_groups); + } + }, + nullptr); + }); + } else { + int group_size = config.quant_config.group_size; + pool->dispense_backend()->do_numa_job([&, this](int i) { + auto& tpc = tps[i]->config_; + if (tpc.intermediate_size % 2 != 0) + throw std::runtime_error("MXFP4 TP flat-buffer: intermediate_size must be even for nibble-aligned addressing, got " + + std::to_string(tpc.intermediate_size)); + size_t weight_elem_count = (size_t)tpc.intermediate_size * tpc.hidden_size; + size_t scales_elem_count = ((size_t)tpc.hidden_size / group_size) * tpc.intermediate_size; + tpc.gate_proj = new uint8_t[(tpc.expert_num * weight_elem_count) / 2]; + tpc.up_proj = new uint8_t[(tpc.expert_num * weight_elem_count) / 2]; + tpc.down_proj = new uint8_t[(tpc.expert_num * weight_elem_count) / 2]; + tpc.gate_scale = new ggml_bf16_t[tpc.expert_num * scales_elem_count]; + tpc.up_scale = new ggml_bf16_t[tpc.expert_num * scales_elem_count]; + tpc.down_scale = new ggml_bf16_t[tpc.expert_num * scales_elem_count]; + }); + + // Fill the TP-sliced flat buffers from full-model data, then load. + // TP slicing is along the intermediate dimension. gate/up are + // [intermediate × hidden] with intermediate as outer (N) dim — contiguous + // memcpy slice works. down is [hidden × intermediate] with intermediate + // as inner (K) dim — must loop per hidden row. + 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 = ((size_t)tpc.hidden_size / group_size) * tpc.intermediate_size; + size_t n_per_tp = tpc.intermediate_size; + size_t hidden_groups_per_row = (size_t)tpc.hidden_size / group_size; + size_t interm_groups_per_row = (size_t)tpc.intermediate_size / group_size; + size_t full_interm_groups_per_row = (size_t)config.intermediate_size / group_size; + pool->get_subpool(i)->do_work_stealing_job( + tpc.expert_num, nullptr, + [&, i](int eid) { + uint64_t lid = expert_map(physical_to_logical_map, eid); + uint8_t* src_gate = (uint8_t*)config.gate_proj + + (lid * (size_t)config.intermediate_size + i * n_per_tp) * config.hidden_size / 2; + uint8_t* src_up = (uint8_t*)config.up_proj + + (lid * (size_t)config.intermediate_size + i * n_per_tp) * config.hidden_size / 2; + std::memcpy((uint8_t*)tpc.gate_proj + eid * weight_elem_count / 2, src_gate, weight_elem_count / 2); + std::memcpy((uint8_t*)tpc.up_proj + eid * weight_elem_count / 2, src_up, weight_elem_count / 2); + ggml_bf16_t* src_gate_scale = + (ggml_bf16_t*)config.gate_scale + + (lid * (size_t)config.intermediate_size + i * n_per_tp) * hidden_groups_per_row; + ggml_bf16_t* src_up_scale = + (ggml_bf16_t*)config.up_scale + + (lid * (size_t)config.intermediate_size + i * n_per_tp) * hidden_groups_per_row; + std::memcpy((ggml_bf16_t*)tpc.gate_scale + eid * scales_elem_count, src_gate_scale, + scales_elem_count * sizeof(ggml_bf16_t)); + std::memcpy((ggml_bf16_t*)tpc.up_scale + eid * scales_elem_count, src_up_scale, + scales_elem_count * sizeof(ggml_bf16_t)); + uint8_t* src_down = + (uint8_t*)config.down_proj + lid * (size_t)config.hidden_size * config.intermediate_size / 2; + ggml_bf16_t* src_down_scale = + (ggml_bf16_t*)config.down_scale + lid * (size_t)config.hidden_size * full_interm_groups_per_row; + for (size_t col = 0; col < (size_t)tpc.hidden_size; col++) { + std::memcpy((uint8_t*)tpc.down_proj + eid * weight_elem_count / 2 + col * n_per_tp / 2, + src_down + (col * config.intermediate_size + i * n_per_tp) / 2, n_per_tp / 2); + std::memcpy((ggml_bf16_t*)tpc.down_scale + eid * scales_elem_count + col * interm_groups_per_row, + src_down_scale + col * full_interm_groups_per_row + i * interm_groups_per_row, + interm_groups_per_row * sizeof(ggml_bf16_t)); + } + }, + nullptr); + }); + + 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[] (ggml_bf16_t*)tpc.gate_scale; + delete[] (ggml_bf16_t*)tpc.up_scale; + delete[] (ggml_bf16_t*)tpc.down_scale; + }); + } + + this->weights_loaded = true; + } + + void write_weight_scale_to_buffer(int gpu_tp_count, int expert_id, + const std::vector& w13_weight_ptrs, + const std::vector& w13_scale_ptrs, + const std::vector& w2_weight_ptrs, + const std::vector& w2_scale_ptrs) { + if (!this->weights_loaded) throw std::runtime_error("Not Loaded"); + if ((int)w13_weight_ptrs.size() != gpu_tp_count || (int)w13_scale_ptrs.size() != gpu_tp_count || + (int)w2_weight_ptrs.size() != gpu_tp_count || (int)w2_scale_ptrs.size() != 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_MXFP4_MOE_H diff --git a/kt-kernel/python/utils/amx.py b/kt-kernel/python/utils/amx.py index 53c553dd..97cb9955 100644 --- a/kt-kernel/python/utils/amx.py +++ b/kt-kernel/python/utils/amx.py @@ -31,6 +31,7 @@ AVX2BF16_MOE = getattr(_moe_mod, "AVX2BF16_MOE", None) 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) AVXVNNI256GPTQInt4_MOE = getattr(_moe_mod, "AVXVNNI256GPTQInt4_MOE", None) AVXVNNI256RawInt4_MOE = getattr(_moe_mod, "AVXVNNI256RawInt4_MOE", None) @@ -45,6 +46,7 @@ _HAS_AVX2_BF16_SUPPORT = AVX2BF16_MOE is not None _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_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 @@ -143,6 +145,37 @@ def _select_rawint4_backend(group_size: Optional[int] = None): return None +def _select_mxfp4_backend(): + """Select MXFP4 backend: AMX/AVX-512 (preferred) > AVX2 (fallback). + + Override with KT_MXFP4_BACKEND=avx2|amx. + Returns None if no MXFP4 backend is available. + """ + forced = os.getenv("KT_MXFP4_BACKEND", "").strip().lower() + + if forced == "amx": + if not _HAS_MXFP4_SUPPORT: + raise RuntimeError( + "KT_MXFP4_BACKEND=amx requested, but AMXFP4_KGroup_MOE is not compiled in. " + "Recompile with AVX512F + AVX512BW + AVX512_BF16 enabled." + ) + return AMXFP4_KGroup_MOE + + if forced == "avx2": + if not _HAS_AVX2_MXFP4_SUPPORT: + raise RuntimeError( + "KT_MXFP4_BACKEND=avx2 requested, but AVX2MXFP4_MOE is not compiled in. " + "Recompile with AVX2 + FMA enabled." + ) + return AVX2MXFP4_MOE + + if _HAS_MXFP4_SUPPORT: + return AMXFP4_KGroup_MOE + if _HAS_AVX2_MXFP4_SUPPORT: + return AVX2MXFP4_MOE + return None + + class AMXMoEWrapper(BaseMoEWrapper): """ AMX-based MoE wrapper implementation. @@ -512,11 +545,12 @@ class NativeMoEWrapper(BaseMoEWrapper): "Please recompile kt_kernel_ext with GPTQ INT4 support enabled.\n" "AVX-VNNI-256 will be selected automatically when available on the current CPU." ) - if method == "MXFP4" and not _HAS_MXFP4_SUPPORT: + if method == "MXFP4" and not (_HAS_MXFP4_SUPPORT or _HAS_AVX2_MXFP4_SUPPORT): raise RuntimeError( - "MXFP4 backend not available. Required ISA:\n" - " - AVX512F + AVX512BW + AVX512_BF16\n" - "Please recompile kt_kernel_ext with AVX512 + BF16 enabled." + "MXFP4 backend not available. Required ISA (any one of):\n" + " - AVX512F + AVX512BW + AVX512_BF16 (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__( @@ -736,7 +770,13 @@ class NativeMoEWrapper(BaseMoEWrapper): moe_config.quant_config.bits = 4 moe_config.quant_config.group_size = group_size moe_config.quant_config.zero_point = False - self.moe = AMXFP4_KGroup_MOE(moe_config) + backend_cls = _select_mxfp4_backend() + if backend_cls is None: + raise RuntimeError( + "No MXFP4 backend available after runtime selection. " + "Compile with AVX512_BF16 (AMXFP4_KGroup_MOE) or AVX2 (AVX2MXFP4_MOE)." + ) + self.moe = backend_cls(moe_config) elif self.method == "FP8": moe_config.quant_config.bits = 8 moe_config.quant_config.group_size = 128 @@ -769,7 +809,7 @@ class NativeMoEWrapper(BaseMoEWrapper): elif self.method == "BF16": # BF16 has no quantization config needed # Prefer AMX backend, fall back to AVX2 - if _HAS_BF16_SUPPORT and _host_has_cpu_flag("amx_bf16"): + if _HAS_BF16_SUPPORT: self.moe = AMXBF16_MOE(moe_config) else: self.moe = AVX2BF16_MOE(moe_config)