mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-19 10:13:32 +00:00
[feat](kt-kernel): AVX2 MXFP4 MoE MXFP4 dispatch (#2015)
* [feat](kt-kernel): AVX2 MXFP4 MoE MXFP4 dispatch
- Add AVX2 MXFP4 MoE kernel (mxfp4-moe.hpp) with 4-token M-blocking
- Wire AVX2MXFP4_MOE binding in ext_bindings.cpp
- Support TP_MOE down_proj slicing and multi-pool per-expert loading
- Add test_fp4_moe_avx2.py integration test
* [fix](kt-kernel): address PR #2010 review — memory leaks, alignment, dynamic expert update
- Track aligned_alloc pointers in AVX2_MOE_BASE::owned_aligned_allocs_ and
free them in the destructor (fixes BufferB backing memory leak on destroy).
- Track per-TP down_buf allocations in TP_MOE::tp_owned_down_bufs_ with
nullptr checks and size rounding to alignment boundary.
- Add nibble-alignment runtime check for per_tp_interm in MXFP4 TP K-split.
- Add write_weight_scale_to_buffer override to TP_MOE<AVX2_MXFP4_MOE_TP>,
enabling dynamic expert update with kt-threadpool-count>=2.
- Guard against ZeroDivisionError in test_fp4_moe_avx2.py.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [fix](kt-kernel): add intermediate_size parity check in MXFP4 TP flat-buffer path
The per-expert path validates that intermediate_size is even (required for
nibble-aligned FP4 addressing), but the flat-buffer path was missing this
check — an odd value would silently truncate /2 divisions, corrupting
memcpy sizes and offsets.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* fix(avx2-moe): fix TP offset calculation and add safety checks
C1-C4: Fix incorrect TP offset calculations in load_weights()
- Per-expert mode used per_tp_interm instead of full_interm for offsets
- This caused segfault when TP > 1 due to invalid pointer arithmetic
H1-H3: Add safety checks
- H1: Validate source weight pointers are not null
- H2: Check lid index is within bounds
- H3: Check BufferB.b is not null in gemm_mxfp4
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
* fix(avx2-moe): revert incorrect C2/C4 offset changes, keep safety checks
Reverts the incorrect offset calculation changes from previous commit.
The original per_tp_interm-based offsets were correct:
- gate/up weights are N-split (along intermediate dim)
- Each TP partition handles per_tp_interm rows
- Offset = i * per_tp_interm * hidden / 2 (not full_interm)
Keeps H1-H3 safety checks:
- H1: Validate source weight pointers are not null
- H2: Check lid index is within bounds
- H3: Check BufferB.b is not null in gemm_mxfp4
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
* fix(avx2): copy weights to owned buffers in per-expert mode
Previously, AVX2 MXFP4 MoE per-expert mode directly pointed BufferB.b
into mmap'd safetensor data. This caused use-after-free when Python
layer releases the mmap after load_weights() returns.
Now AVX2 copies weights into owned buffers via memcpy/from_raw_mat(),
matching AMX behavior. This decouples the MoE weights from mmap lifecycle.
Changes:
- buffer_b_required_size_impl: always allocate full buffer (weights + scales)
- make_buffer_b_impl: always create full BufferB with owned storage
- Single-TP per-expert: use from_raw_mat() instead of direct pointer
- TP_MOE per-expert: add gate/up owned buffers with memcpy
- Destructor: free gate/up buffers alongside down
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
* Revert "[fix] Add runtime AMX BF16 check to prevent SIGILL on pre-Sapphire Rapids CPUs (#2018)"
This reverts commit f1e2b82c74.
* Remove AMX tile MXFP4 kernel (GemmKernel224MXFP4)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Benjamin
GitHub
parent
f1e2b82c74
commit
ef6c47f9d2
5 changed files with 1139 additions and 10 deletions
236
kt-kernel/examples/test_fp4_moe_avx2.py
Normal file
236
kt-kernel/examples/test_fp4_moe_avx2.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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<AVX2_FP8_MOE_TP<avx2::GemmKernelAVX2FP8>>(moe_module, "AVX2FP8_MOE");
|
||||
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<AVXVNNI256_GPTQ_INT4_MOE_TP<avxvnni::GemmKernelAVXVNNI256GPTQInt4>>(moe_module,
|
||||
"AVXVNNI256GPTQInt4_MOE");
|
||||
bind_moe_module<AVXVNNI256_RAW_INT4_MOE_TP<avxvnni_rawint4::GemmKernelAVXVNNI256RawInt4>>(moe_module,
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@
|
|||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
|
@ -61,6 +62,8 @@ class AVX2_MOE_BASE {
|
|||
std::vector<std::shared_ptr<typename T::BufferB>> down_bb_;
|
||||
std::vector<std::shared_ptr<typename T::BufferC>> down_bc_;
|
||||
|
||||
std::vector<void*> 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;
|
||||
|
|
|
|||
839
kt-kernel/operators/avx2/mxfp4-moe.hpp
Normal file
839
kt-kernel/operators/avx2/mxfp4-moe.hpp
Normal file
|
|
@ -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 <immintrin.h>
|
||||
#include <tmmintrin.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#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<int, int> 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 T = avx2::GemmKernelAVX2MXFP4>
|
||||
class AVX2_MXFP4_MOE_TP : public AVX2_MOE_BASE<T, AVX2_MXFP4_MOE_TP<T>> {
|
||||
using Base = AVX2_MOE_BASE<T, AVX2_MXFP4_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_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<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, config_.quant_config.group_size, data);
|
||||
}
|
||||
std::shared_ptr<typename T::BufferB> 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<typename T::BufferB>((int)n, (int)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) {
|
||||
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<AVX2_MXFP4_MOE_TP> 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<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 {
|
||||
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 <typename K>
|
||||
class TP_MOE<AVX2_MXFP4_MOE_TP<K>> : public TP_MOE<AVX2_MOE_BASE<K, AVX2_MXFP4_MOE_TP<K>>> {
|
||||
std::vector<void*> tp_owned_down_bufs_;
|
||||
std::vector<void*> tp_owned_gate_bufs_;
|
||||
std::vector<void*> tp_owned_up_bufs_;
|
||||
|
||||
public:
|
||||
using Base = TP_MOE<AVX2_MOE_BASE<K, AVX2_MXFP4_MOE_TP<K>>>;
|
||||
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<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 ((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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue