wehub-resource-sync
ec436095dd
Book-CI / test (macos-latest) (push) Has been cancelled
Book-CI / test (ubuntu-latest) (push) Has been cancelled
Book-CI / test (windows-latest) (push) Has been cancelled
Release Fake Tag / publish (push) Has been cancelled
Deploy / deploy (macos-latest) (push) Has been cancelled
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Deploy / deploy (windows-latest) (push) Has been cancelled
Release to PyPI / Build & publish sglang-kt (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.11) (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.12) (push) Has been cancelled
Release to PyPI / Publish kt-kernel to PyPI (push) Has been cancelled
681 lines
35 KiB
C++
681 lines
35 KiB
C++
/**
|
|
* @Description : AVX2 RAWINT4 MoE operator for Kimi native INT4 weights
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* RAWINT4 stores signed int4 weights packed two values per byte, plus BF16
|
|
* K-group scales. This AVX2 backend keeps the native layout so GPU layerwise
|
|
* prefill can reuse the same weight buffers, and uses BF16 activations with
|
|
* FP32 accumulation for CPU decode.
|
|
**/
|
|
#ifndef CPUINFER_OPERATOR_AVX2_RAW_INT4_MOE_H
|
|
#define CPUINFER_OPERATOR_AVX2_RAW_INT4_MOE_H
|
|
|
|
#include "avx2_bf16_gemm.hpp"
|
|
#include "avx2_bf16_utils.hpp"
|
|
#include "gptq_int4_dequant.hpp"
|
|
#include "moe_base.hpp"
|
|
|
|
namespace avx2 {
|
|
|
|
// Local, RAWINT4-only variant of gptq_sym_dequant_8x4bit that skips the
|
|
// per-group scale multiply. Callers fold the scale in via a single FMA after
|
|
// the K-loop so the inner path avoids a broadcast+mul per packed int32.
|
|
static inline __m256 rawint4_dequant_8x4bit_unscaled(uint32_t packed_weight) {
|
|
const __m256i shifts = _mm256_set_epi32(28, 24, 20, 16, 12, 8, 4, 0);
|
|
__m256i packed_v = _mm256_set1_epi32(packed_weight);
|
|
__m256i nibbles = _mm256_and_si256(_mm256_srlv_epi32(packed_v, shifts), _mm256_set1_epi32(0xF));
|
|
__m256 w = _mm256_cvtepi32_ps(nibbles);
|
|
return _mm256_sub_ps(w, _mm256_set1_ps(8.0f));
|
|
}
|
|
|
|
struct GemmKernelAVX2RawInt4 {
|
|
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 = 8;
|
|
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); }
|
|
|
|
struct BufferA {
|
|
ggml_bf16_t* data = nullptr;
|
|
size_t max_m = 0;
|
|
size_t k = 0;
|
|
|
|
BufferA() = default;
|
|
BufferA(size_t m, size_t k_, int, void* ptr) : data((ggml_bf16_t*)ptr), max_m(m), k(k_) {}
|
|
|
|
static size_t required_size(size_t m, size_t k, int) { 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;
|
|
float* d = nullptr;
|
|
int n = 0;
|
|
int k = 0;
|
|
int k_group_size = 0;
|
|
int 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("RAWINT4 requires k aligned to group_size and 8");
|
|
}
|
|
k_group_count = k / k_group_size;
|
|
d = (float*)((uint8_t*)ptr + ((size_t)n * k / 2));
|
|
}
|
|
|
|
// Scale-only allocation: b points to external (mmap'd) weight data; d owns scale_ptr.
|
|
// Used when weights are consumed directly from safetensor mmap without copying.
|
|
BufferB(int n_, int k_, int k_group_size_, void* scale_ptr, std::nullptr_t /*scale_only*/)
|
|
: 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("RAWINT4 requires k aligned to group_size and 8");
|
|
}
|
|
k_group_count = k / k_group_size;
|
|
d = (float*)scale_ptr;
|
|
}
|
|
|
|
// Full size: packed INT4 weights + float32 scales.
|
|
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);
|
|
}
|
|
|
|
// Scale-only size: only float32 scales (b will point to external weight data).
|
|
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; // scale-only mode: b is an external pointer set later
|
|
auto [n_start, n_end] = split_range_n(n, ith, nth);
|
|
const size_t row_bytes = (size_t)k / 2;
|
|
std::memcpy(b + (size_t)n_start * row_bytes, proj + (size_t)n_start * row_bytes,
|
|
(size_t)(n_end - n_start) * row_bytes);
|
|
}
|
|
};
|
|
|
|
struct BufferC {
|
|
float* data = nullptr;
|
|
size_t max_m = 0;
|
|
size_t n = 0;
|
|
|
|
BufferC() = default;
|
|
BufferC(size_t m, size_t n_, void* ptr) : 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]);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
};
|
|
|
|
static inline void gemm_rawint4(int m, int n, int k, GemmKernelAVX2RawInt4::BufferA& a,
|
|
GemmKernelAVX2RawInt4::BufferB& b, GemmKernelAVX2RawInt4::BufferC& c, int ith,
|
|
int nth) {
|
|
auto [n_start, n_end] = split_range(n, ith, nth);
|
|
const int group_size = b.k_group_size;
|
|
const int group_count = b.k_group_count;
|
|
const size_t row_bytes = (size_t)k / 2;
|
|
|
|
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;
|
|
|
|
// Prefetch the head of the next B row while we chew through this one.
|
|
// Streams the decoupled weight pages for large N without blocking.
|
|
if (ni + 1 < n_end) {
|
|
const uint8_t* b_next = b.b + (size_t)(ni + 1) * row_bytes;
|
|
_mm_prefetch((const char*)b_next, _MM_HINT_T0);
|
|
_mm_prefetch((const char*)(b_next + 64), _MM_HINT_T0);
|
|
}
|
|
|
|
for (int mi = 0; mi < m; mi++) {
|
|
const ggml_bf16_t* a_row = a.data + (size_t)mi * a.k;
|
|
|
|
// Running vector total: each group folds in via a single FMA with its
|
|
// scale broadcast, so we only do one horizontal reduction per (mi, ni).
|
|
__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 acc1 = _mm256_setzero_ps();
|
|
__m256 acc2 = _mm256_setzero_ps();
|
|
__m256 acc3 = _mm256_setzero_ps();
|
|
__m256 acc4 = _mm256_setzero_ps();
|
|
|
|
int ki = 0;
|
|
for (; ki + 32 <= group_size; ki += 32) {
|
|
uint32_t p0, p1, p2, p3;
|
|
std::memcpy(&p0, b_row + (k_base + ki) / 2, sizeof(uint32_t));
|
|
std::memcpy(&p1, b_row + (k_base + ki + 8) / 2, sizeof(uint32_t));
|
|
std::memcpy(&p2, b_row + (k_base + ki + 16) / 2, sizeof(uint32_t));
|
|
std::memcpy(&p3, b_row + (k_base + ki + 24) / 2, sizeof(uint32_t));
|
|
acc1 = _mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki), rawint4_dequant_8x4bit_unscaled(p0), acc1);
|
|
acc2 = _mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki + 8), rawint4_dequant_8x4bit_unscaled(p1), acc2);
|
|
acc3 =
|
|
_mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki + 16), rawint4_dequant_8x4bit_unscaled(p2), acc3);
|
|
acc4 =
|
|
_mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki + 24), rawint4_dequant_8x4bit_unscaled(p3), acc4);
|
|
}
|
|
__m256 g_acc = _mm256_add_ps(_mm256_add_ps(acc1, acc3), _mm256_add_ps(acc2, acc4));
|
|
|
|
for (; ki + 8 <= group_size; ki += 8) {
|
|
uint32_t packed;
|
|
std::memcpy(&packed, b_row + (k_base + ki) / 2, sizeof(uint32_t));
|
|
g_acc =
|
|
_mm256_fmadd_ps(load_bf16_to_fp32(a_row + k_base + ki), rawint4_dequant_8x4bit_unscaled(packed), g_acc);
|
|
}
|
|
|
|
// Fold this group's unscaled accumulator into the running total with
|
|
// one scale-broadcast FMA — saves (group_count - 1) hsum reductions
|
|
// compared to reducing per group.
|
|
total_acc = _mm256_fmadd_ps(g_acc, _mm256_broadcast_ss(&scale), total_acc);
|
|
|
|
for (; ki < group_size; ki++) {
|
|
const uint8_t packed = b_row[(k_base + ki) / 2];
|
|
const int nibble = ((k_base + ki) & 1) ? (packed >> 4) : (packed & 0x0F);
|
|
scalar_tail += GGML_BF16_TO_FP32(a_row[k_base + ki]) * (float)(nibble - 8) * scale;
|
|
}
|
|
}
|
|
|
|
c.data[(size_t)mi * n + ni] = hsum_avx2(total_acc) + scalar_tail;
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace avx2
|
|
|
|
template <class T = avx2::GemmKernelAVX2RawInt4>
|
|
class AVX2_RAW_INT4_MOE_TP : public AVX2_MOE_BASE<T, AVX2_RAW_INT4_MOE_TP<T>> {
|
|
using Base = AVX2_MOE_BASE<T, AVX2_RAW_INT4_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_RAW_INT4_MOE_TP() = default;
|
|
AVX2_RAW_INT4_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("RAWINT4 AVX2 MoE only supports KGroup signed INT4 without zero point");
|
|
}
|
|
printf("Created AVX2_RAW_INT4_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 {
|
|
// When per-expert source pointers are available, only allocate float32 scales.
|
|
// Weights will be served directly from the mmap'd safetensor data (no copy).
|
|
if (!config_.gate_projs.empty()) {
|
|
return T::BufferB::required_size_scale_only(n, k, config_.quant_config.group_size);
|
|
}
|
|
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 {
|
|
// Scale-only mode: b is nullptr here; set externally in load_weights().
|
|
if (!config_.gate_projs.empty()) {
|
|
return std::make_shared<typename T::BufferB>((int)n, (int)k, config_.quant_config.group_size, data, nullptr);
|
|
}
|
|
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_rawint4(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_rawint4(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("RAWINT4 AVX2 MoE requires weight pointers");
|
|
}
|
|
if (!use_per_expert && config_.gate_scale == nullptr) {
|
|
throw std::runtime_error("RAWINT4 AVX2 MoE requires scale pointers");
|
|
}
|
|
|
|
if (use_per_expert) {
|
|
// Direct-pointer mode: BufferB.b is set to point into the mmap'd safetensor data
|
|
// (no weight copy). Only float32 scales are allocated and converted.
|
|
//
|
|
// For gate/up: source shape [intermediate_size_full, hidden_size/2] row-major.
|
|
// TP partition tp_part_idx handles rows [tp_part_idx * n_per_tp, (tp_part_idx+1) * n_per_tp).
|
|
// These are contiguous in memory → simple byte offset: tp_part_idx * n_per_tp * (k/2).
|
|
// With kt_threadpool_count=1 (tp_part_idx=0): offset = 0.
|
|
//
|
|
// For down: source shape [hidden_size, intermediate_size_full/2] row-major.
|
|
// TP partition tp_part_idx handles columns [tp_part_idx * n_per_tp/2, ...) per row.
|
|
// These are NOT contiguous across rows for tp_count > 1.
|
|
// This mode therefore requires kt_threadpool_count=1 (enforced in outer TP wrapper).
|
|
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 ||
|
|
up_bb_[expert_idx] == nullptr || down_bb_[expert_idx] == nullptr) {
|
|
return;
|
|
}
|
|
uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx);
|
|
// Gate/Up row offset for this TP partition (bytes).
|
|
size_t gate_tp_byte_offset = (size_t)tp_part_idx * config_.intermediate_size * config_.hidden_size / 2;
|
|
gate_bb_[expert_idx]->b = (uint8_t*)config_.gate_projs[0][logical_expert_id] + gate_tp_byte_offset;
|
|
up_bb_[expert_idx]->b = (uint8_t*)config_.up_projs[0][logical_expert_id] + gate_tp_byte_offset;
|
|
// Down column offset (bytes per row start). Correct only for tp_count=1.
|
|
size_t down_tp_byte_offset = (size_t)tp_part_idx * config_.intermediate_size / 2;
|
|
down_bb_[expert_idx]->b = (uint8_t*)config_.down_projs[0][logical_expert_id] + down_tp_byte_offset;
|
|
},
|
|
nullptr);
|
|
|
|
// Scale conversion: BF16 → float32.
|
|
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 ||
|
|
up_bb_[expert_idx] == nullptr || down_bb_[expert_idx] == nullptr) {
|
|
return;
|
|
}
|
|
uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx);
|
|
size_t scale_elem_count = ((size_t)config_.hidden_size * config_.intermediate_size) / group_size;
|
|
// Gate/Up scale offset: rows [tp_part_idx * n_per_tp, ...) of scale[n_total, k/gs].
|
|
size_t gate_scale_tp_offset =
|
|
(size_t)tp_part_idx * config_.intermediate_size * (config_.hidden_size / group_size);
|
|
// Down scale offset: cols [tp_part_idx * (n_per_tp/gs), ...) per row. Correct for tp_count=1.
|
|
size_t down_scale_tp_offset = (size_t)tp_part_idx * (config_.intermediate_size / group_size);
|
|
convert_or_copy(gate_bb_[expert_idx]->d,
|
|
(const ggml_bf16_t*)config_.gate_scales[0][logical_expert_id] + gate_scale_tp_offset,
|
|
scale_elem_count);
|
|
convert_or_copy(up_bb_[expert_idx]->d,
|
|
(const ggml_bf16_t*)config_.up_scales[0][logical_expert_id] + gate_scale_tp_offset,
|
|
scale_elem_count);
|
|
convert_or_copy(down_bb_[expert_idx]->d,
|
|
(const ggml_bf16_t*)config_.down_scales[0][logical_expert_id] + down_scale_tp_offset,
|
|
scale_elem_count);
|
|
},
|
|
nullptr);
|
|
} else {
|
|
// Flat-buffer mode: copy TP-sliced weights from flat buffer into allocated BufferB.b.
|
|
int nth = T::recommended_nth(config_.intermediate_size);
|
|
pool->do_work_stealing_job(
|
|
nth * config_.expert_num, nullptr,
|
|
[this, nth, physical_to_logical_map](int task_id) {
|
|
uint64_t expert_idx = task_id / nth;
|
|
if (config_.should_skip_expert(expert_idx)) return;
|
|
uint64_t logical_expert_id = expert_map(physical_to_logical_map, expert_idx);
|
|
int ith = task_id % nth;
|
|
size_t weight_offset = ((size_t)logical_expert_id * config_.intermediate_size * config_.hidden_size) / 2;
|
|
gate_bb_[expert_idx]->from_raw_mat((const uint8_t*)config_.gate_proj + weight_offset, ith, nth);
|
|
up_bb_[expert_idx]->from_raw_mat((const uint8_t*)config_.up_proj + weight_offset, ith, nth);
|
|
},
|
|
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 logical_expert_id = expert_map(physical_to_logical_map, expert_idx);
|
|
int ith = task_id % nth_down;
|
|
size_t weight_offset = ((size_t)logical_expert_id * 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);
|
|
|
|
// Scale conversion in flat-buffer mode.
|
|
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 logical_expert_id = 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 + logical_expert_id * scale_elem_count, scale_elem_count);
|
|
convert_or_copy(up_bb_[expert_idx]->d,
|
|
(ggml_bf16_t*)config_.up_scale + logical_expert_id * scale_elem_count, scale_elem_count);
|
|
convert_or_copy(down_bb_[expert_idx]->d,
|
|
(ggml_bf16_t*)config_.down_scale + logical_expert_id * 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("RAWINT4 write_weights_to_buffer requested an expert without loaded weights");
|
|
}
|
|
const int group_size = config_.quant_config.group_size;
|
|
auto pool = config_.pool->get_subpool(tp_part_idx);
|
|
|
|
size_t cpu_tp_weight_elem_count = (size_t)config_.intermediate_size * config_.hidden_size;
|
|
size_t cpu_tp_weight_bytes = cpu_tp_weight_elem_count / 2;
|
|
size_t cpu_tp_scale_elem_count = cpu_tp_weight_elem_count / group_size;
|
|
size_t gpu_tp_weight_elem_count = (size_t)full_config.intermediate_size * full_config.hidden_size / gpu_tp_count;
|
|
size_t gpu_tp_weight_bytes = gpu_tp_weight_elem_count / 2;
|
|
size_t gpu_tp_scale_elem_count = gpu_tp_weight_elem_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_weight_dst = (uint8_t*)w13_weight_ptrs[target_gpu_tp];
|
|
ggml_bf16_t* w13_scale_dst = (ggml_bf16_t*)w13_scale_ptrs[target_gpu_tp];
|
|
uint8_t* w2_weight_dst = (uint8_t*)w2_weight_ptrs[target_gpu_tp];
|
|
ggml_bf16_t* w2_scale_dst = (ggml_bf16_t*)w2_scale_ptrs[target_gpu_tp];
|
|
size_t offset_in_gpu_weight = local_idx * cpu_tp_weight_bytes;
|
|
size_t offset_in_gpu_scale = local_idx * cpu_tp_scale_elem_count;
|
|
|
|
constexpr int NUM_WEIGHT_TASKS = 8;
|
|
constexpr int MIN_COLS_PER_TASK = 128;
|
|
int num_down_tasks = std::min(std::max(1, config_.hidden_size / MIN_COLS_PER_TASK), 32);
|
|
int total_tasks = NUM_WEIGHT_TASKS * 2 + num_down_tasks + 2;
|
|
size_t weight_chunk_size = (cpu_tp_weight_bytes + NUM_WEIGHT_TASKS - 1) / NUM_WEIGHT_TASKS;
|
|
weight_chunk_size = (weight_chunk_size + 63) & ~63ULL;
|
|
|
|
pool->do_work_stealing_job(
|
|
total_tasks, nullptr,
|
|
[=, this](int task_id) {
|
|
if (task_id < NUM_WEIGHT_TASKS) {
|
|
size_t start = (size_t)task_id * weight_chunk_size;
|
|
size_t end = std::min(start + weight_chunk_size, cpu_tp_weight_bytes);
|
|
if (start < end)
|
|
std::memcpy(w13_weight_dst + offset_in_gpu_weight + start, gate_bb_[expert_id]->b + start, end - start);
|
|
} else if (task_id < NUM_WEIGHT_TASKS * 2) {
|
|
int chunk_idx = task_id - NUM_WEIGHT_TASKS;
|
|
size_t start = (size_t)chunk_idx * weight_chunk_size;
|
|
size_t end = std::min(start + weight_chunk_size, cpu_tp_weight_bytes);
|
|
if (start < end)
|
|
std::memcpy(w13_weight_dst + offset_in_gpu_weight + gpu_tp_weight_bytes + start,
|
|
up_bb_[expert_id]->b + start, end - start);
|
|
} else if (task_id < NUM_WEIGHT_TASKS * 2 + num_down_tasks) {
|
|
int chunk_idx = task_id - NUM_WEIGHT_TASKS * 2;
|
|
size_t cols_per_chunk = (config_.hidden_size + num_down_tasks - 1) / num_down_tasks;
|
|
size_t col_start = (size_t)chunk_idx * cols_per_chunk;
|
|
size_t col_end = std::min(col_start + cols_per_chunk, (size_t)config_.hidden_size);
|
|
size_t weight_per_col = config_.intermediate_size >> 1;
|
|
size_t scale_per_col = config_.intermediate_size / group_size;
|
|
size_t gpu_weight_stride = (full_config.intermediate_size / gpu_tp_count) >> 1;
|
|
size_t gpu_scale_stride = (full_config.intermediate_size / gpu_tp_count) / group_size;
|
|
size_t gpu_weight_slice_offset = local_idx * weight_per_col;
|
|
size_t gpu_scale_slice_offset = local_idx * scale_per_col;
|
|
for (size_t col = col_start; col < col_end; col++) {
|
|
std::memcpy(w2_weight_dst + col * gpu_weight_stride + gpu_weight_slice_offset,
|
|
down_bb_[expert_id]->b + col * weight_per_col, weight_per_col);
|
|
fp32_to_bf16(w2_scale_dst + col * gpu_scale_stride + gpu_scale_slice_offset,
|
|
down_bb_[expert_id]->d + col * scale_per_col, scale_per_col);
|
|
}
|
|
} else if (task_id == NUM_WEIGHT_TASKS * 2 + num_down_tasks) {
|
|
fp32_to_bf16(w13_scale_dst + offset_in_gpu_scale, gate_bb_[expert_id]->d, cpu_tp_scale_elem_count);
|
|
} else {
|
|
fp32_to_bf16(w13_scale_dst + offset_in_gpu_scale + gpu_tp_scale_elem_count, up_bb_[expert_id]->d,
|
|
cpu_tp_scale_elem_count);
|
|
}
|
|
},
|
|
nullptr);
|
|
} else {
|
|
int gpu_tps_per_cpu_tp = gpu_tp_count / cpu_tp_count;
|
|
int start_gpu_tp = tp_part_idx * gpu_tps_per_cpu_tp;
|
|
size_t data_per_gpu_tp_weight = cpu_tp_weight_bytes / gpu_tps_per_cpu_tp;
|
|
size_t data_per_gpu_tp_scale = cpu_tp_scale_elem_count / gpu_tps_per_cpu_tp;
|
|
constexpr int NUM_WEIGHT_TASKS = 8;
|
|
constexpr int MIN_COLS_PER_TASK = 128;
|
|
int num_down_tasks = std::min(std::max(1, config_.hidden_size / MIN_COLS_PER_TASK), 32);
|
|
int tasks_per_gpu_tp = NUM_WEIGHT_TASKS * 2 + num_down_tasks + 2;
|
|
int total_tasks = tasks_per_gpu_tp * gpu_tps_per_cpu_tp;
|
|
size_t weight_chunk_size = (data_per_gpu_tp_weight + NUM_WEIGHT_TASKS - 1) / NUM_WEIGHT_TASKS;
|
|
weight_chunk_size = (weight_chunk_size + 63) & ~63ULL;
|
|
|
|
pool->do_work_stealing_job(
|
|
total_tasks, nullptr,
|
|
[=, this, &w13_weight_ptrs, &w13_scale_ptrs, &w2_weight_ptrs, &w2_scale_ptrs](int task_id) {
|
|
int local_gpu_idx = task_id / tasks_per_gpu_tp;
|
|
int task_type = task_id % tasks_per_gpu_tp;
|
|
int gpu_tp_idx = start_gpu_tp + local_gpu_idx;
|
|
uint8_t* w13_weight_dst = (uint8_t*)w13_weight_ptrs[gpu_tp_idx];
|
|
ggml_bf16_t* w13_scale_dst = (ggml_bf16_t*)w13_scale_ptrs[gpu_tp_idx];
|
|
uint8_t* w2_weight_dst = (uint8_t*)w2_weight_ptrs[gpu_tp_idx];
|
|
ggml_bf16_t* w2_scale_dst = (ggml_bf16_t*)w2_scale_ptrs[gpu_tp_idx];
|
|
size_t cpu_offset_weight = (size_t)local_gpu_idx * data_per_gpu_tp_weight;
|
|
size_t cpu_offset_scale = (size_t)local_gpu_idx * data_per_gpu_tp_scale;
|
|
if (task_type < NUM_WEIGHT_TASKS) {
|
|
size_t start = (size_t)task_type * weight_chunk_size;
|
|
size_t end = std::min(start + weight_chunk_size, data_per_gpu_tp_weight);
|
|
if (start < end)
|
|
std::memcpy(w13_weight_dst + start, gate_bb_[expert_id]->b + cpu_offset_weight + start, end - start);
|
|
} else if (task_type < NUM_WEIGHT_TASKS * 2) {
|
|
int chunk_idx = task_type - NUM_WEIGHT_TASKS;
|
|
size_t start = (size_t)chunk_idx * weight_chunk_size;
|
|
size_t end = std::min(start + weight_chunk_size, data_per_gpu_tp_weight);
|
|
if (start < end)
|
|
std::memcpy(w13_weight_dst + gpu_tp_weight_bytes + start,
|
|
up_bb_[expert_id]->b + cpu_offset_weight + start, end - start);
|
|
} else if (task_type < NUM_WEIGHT_TASKS * 2 + num_down_tasks) {
|
|
int chunk_idx = task_type - NUM_WEIGHT_TASKS * 2;
|
|
size_t cols_per_chunk = (config_.hidden_size + num_down_tasks - 1) / num_down_tasks;
|
|
size_t col_start = (size_t)chunk_idx * cols_per_chunk;
|
|
size_t col_end = std::min(col_start + cols_per_chunk, (size_t)config_.hidden_size);
|
|
size_t weight_per_gpu_col = (config_.intermediate_size / gpu_tps_per_cpu_tp) >> 1;
|
|
size_t scale_per_gpu_col = (config_.intermediate_size / gpu_tps_per_cpu_tp) / group_size;
|
|
for (size_t col = col_start; col < col_end; col++) {
|
|
size_t col_offset_weight = (col * config_.intermediate_size / 2) +
|
|
(local_gpu_idx * data_per_gpu_tp_weight / config_.hidden_size);
|
|
size_t col_offset_scale = (col * (config_.intermediate_size / group_size)) +
|
|
(local_gpu_idx * data_per_gpu_tp_scale / config_.hidden_size);
|
|
std::memcpy(w2_weight_dst + col * weight_per_gpu_col, down_bb_[expert_id]->b + col_offset_weight,
|
|
weight_per_gpu_col);
|
|
fp32_to_bf16(w2_scale_dst + col * scale_per_gpu_col, down_bb_[expert_id]->d + col_offset_scale,
|
|
scale_per_gpu_col);
|
|
}
|
|
} else if (task_type == NUM_WEIGHT_TASKS * 2 + num_down_tasks) {
|
|
fp32_to_bf16(w13_scale_dst, gate_bb_[expert_id]->d + cpu_offset_scale, data_per_gpu_tp_scale);
|
|
} else {
|
|
fp32_to_bf16(w13_scale_dst + gpu_tp_scale_elem_count, up_bb_[expert_id]->d + cpu_offset_scale,
|
|
data_per_gpu_tp_scale);
|
|
}
|
|
},
|
|
nullptr);
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename K>
|
|
class TP_MOE<AVX2_RAW_INT4_MOE_TP<K>> : public TP_MOE<AVX2_MOE_BASE<K, AVX2_RAW_INT4_MOE_TP<K>>> {
|
|
public:
|
|
using Base = TP_MOE<AVX2_MOE_BASE<K, AVX2_RAW_INT4_MOE_TP<K>>>;
|
|
using Base::Base;
|
|
|
|
void load_weights() override {
|
|
auto& config = this->config;
|
|
auto& tps = this->tps;
|
|
auto pool = config.pool;
|
|
const uint64_t* physical_to_logical_map = (const uint64_t*)config.physical_to_logical_map;
|
|
bool use_per_expert_ptrs = !config.gate_projs.empty();
|
|
if (config.gate_projs.empty() && config.gate_scale == nullptr) {
|
|
throw std::runtime_error("RAWINT4 AVX2 MoE only supports packed INT4 with KGroup scales");
|
|
}
|
|
|
|
if (use_per_expert_ptrs) {
|
|
// Direct-pointer mode: inner load_weights() sets BufferB.b directly from mmap'd data.
|
|
// The down projection column-gather required for tp_count > 1 is NOT supported in
|
|
// this mode; enforce single NUMA pool (kt_threadpool_count=1).
|
|
if (this->tp_count > 1) {
|
|
throw std::runtime_error(
|
|
"RAWINT4 per-expert pointer mode requires kt_threadpool_count=1 "
|
|
"(down projection TP column-gather is unsupported with direct pointers)");
|
|
}
|
|
DO_TPS_LOAD_WEIGHTS(pool);
|
|
} else {
|
|
// Flat-buffer mode: build a TP-sliced contiguous buffer per TP partition, then
|
|
// call inner load_weights() which copies into BufferB.b from the flat buffer.
|
|
int group_size = config.quant_config.group_size;
|
|
pool->dispense_backend()->do_numa_job([&, this](int i) {
|
|
auto& tpc = tps[i]->config_;
|
|
size_t weight_elem_count = (size_t)tpc.intermediate_size * tpc.hidden_size;
|
|
size_t scales_elem_count = ((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];
|
|
|
|
pool->get_subpool(i)->do_work_stealing_job(
|
|
tpc.expert_num, nullptr,
|
|
[&, i](int expert_id_) {
|
|
size_t expert_id = expert_map(physical_to_logical_map, expert_id_);
|
|
uint8_t* src_gate = (uint8_t*)config.gate_proj +
|
|
((expert_id * (size_t)config.intermediate_size * config.hidden_size) >> 1);
|
|
uint8_t* src_up =
|
|
(uint8_t*)config.up_proj + ((expert_id * (size_t)config.intermediate_size * config.hidden_size) >> 1);
|
|
uint8_t* src_down = (uint8_t*)config.down_proj +
|
|
((expert_id * (size_t)config.intermediate_size * config.hidden_size) >> 1);
|
|
ggml_bf16_t* src_gate_scale =
|
|
(ggml_bf16_t*)config.gate_scale +
|
|
expert_id * ((size_t)config.hidden_size / group_size) * config.intermediate_size;
|
|
ggml_bf16_t* src_up_scale =
|
|
(ggml_bf16_t*)config.up_scale +
|
|
expert_id * ((size_t)config.hidden_size / group_size) * config.intermediate_size;
|
|
ggml_bf16_t* src_down_scale =
|
|
(ggml_bf16_t*)config.down_scale +
|
|
expert_id * ((size_t)config.intermediate_size / group_size) * config.hidden_size;
|
|
|
|
std::memcpy((uint8_t*)tpc.gate_proj + ((expert_id * weight_elem_count) >> 1),
|
|
src_gate + ((i * weight_elem_count) >> 1), weight_elem_count >> 1);
|
|
std::memcpy((uint8_t*)tpc.up_proj + ((expert_id * weight_elem_count) >> 1),
|
|
src_up + ((i * weight_elem_count) >> 1), weight_elem_count >> 1);
|
|
std::memcpy((ggml_bf16_t*)tpc.gate_scale + expert_id * scales_elem_count,
|
|
src_gate_scale + i * scales_elem_count, sizeof(ggml_bf16_t) * scales_elem_count);
|
|
std::memcpy((ggml_bf16_t*)tpc.up_scale + expert_id * scales_elem_count,
|
|
src_up_scale + i * scales_elem_count, sizeof(ggml_bf16_t) * scales_elem_count);
|
|
|
|
for (size_t col = 0; col < (size_t)config.hidden_size; col++) {
|
|
std::memcpy(
|
|
(uint8_t*)tpc.down_proj + ((expert_id * weight_elem_count + col * tpc.intermediate_size) >> 1),
|
|
src_down + ((col * config.intermediate_size + i * tpc.intermediate_size) >> 1),
|
|
tpc.intermediate_size >> 1);
|
|
std::memcpy((ggml_bf16_t*)tpc.down_scale +
|
|
(expert_id * scales_elem_count + col * (tpc.intermediate_size / group_size)),
|
|
src_down_scale + (col * (config.intermediate_size / group_size) +
|
|
i * (tpc.intermediate_size / group_size)),
|
|
sizeof(ggml_bf16_t) * (tpc.intermediate_size / group_size));
|
|
}
|
|
},
|
|
nullptr);
|
|
printf("AVX2 RAWINT4 TP %d load weight done.\n", i);
|
|
});
|
|
|
|
DO_TPS_LOAD_WEIGHTS(pool);
|
|
|
|
pool->dispense_backend()->do_numa_job([&, this](int i) {
|
|
auto& tpc = tps[i]->config_;
|
|
delete[] (uint8_t*)tpc.gate_proj;
|
|
delete[] (uint8_t*)tpc.up_proj;
|
|
delete[] (uint8_t*)tpc.down_proj;
|
|
delete[] (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 == false) 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_RAW_INT4_MOE_H
|