mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-17 05:11:19 +08:00
Add blocked RAWINT4 prefill backend
This commit is contained in:
yyj
@@ -791,6 +791,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) {
|
||||
bind_moe_module<AMX_MOE_TP<amx::GemmKernel224Int4_1>>(moe_module, "AMXInt4_1_MOE");
|
||||
bind_moe_module<AMX_AWQ_MOE_TP<amx::GemmKernel224Int4_1_LowKGroup>>(moe_module, "AMXInt4_1KGroup_MOE");
|
||||
bind_moe_module<AMX_K2_MOE_TP<amx::GemmKernel224Int4SmallKGroup>>(moe_module, "AMXInt4_KGroup_MOE");
|
||||
bind_moe_module<AMX_K2_MOE_TP<amx::GemmKernel224Int4SmallKGroupBlocked>>(moe_module, "AMXInt4_KGroupBlocked_MOE");
|
||||
#if defined(__AVX512F__)
|
||||
bind_moe_module<AMX_BF16_MOE_TP<amx::GemmKernel224BF16>>(moe_module, "AMXBF16_MOE");
|
||||
bind_moe_module<AMX_FP8_MOE_TP<amx::GemmKernel224FP8>>(moe_module, "AMXFP8_MOE");
|
||||
|
||||
@@ -241,6 +241,91 @@ class AMX_K2_MOE_TP : public AMX_MOE_BASE<T, AMX_K2_MOE_TP<T>> {
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void write_weights_to_buffer_blocked(int gpu_tp_count, [[maybe_unused]] int cpu_tp_count, int expert_id,
|
||||
const GeneralMOEConfig& full_config,
|
||||
const std::vector<uintptr_t>& w13_weight_ptrs,
|
||||
const std::vector<uintptr_t>& w13_scale_ptrs,
|
||||
const std::vector<uintptr_t>& w2_weight_ptrs,
|
||||
const std::vector<uintptr_t>& w2_scale_ptrs) const {
|
||||
const int group_size = config_.quant_config.group_size;
|
||||
auto pool = config_.pool->get_subpool(tp_part_idx);
|
||||
|
||||
constexpr int NUM_W13_TASKS = 32;
|
||||
constexpr int NUM_W2_TASKS = 32;
|
||||
const int total_tasks = NUM_W13_TASKS + NUM_W2_TASKS;
|
||||
|
||||
const int cpu_n_w13 = config_.intermediate_size;
|
||||
const int cpu_k_w13 = config_.hidden_size;
|
||||
const int gpu_n_w13 = full_config.intermediate_size / gpu_tp_count;
|
||||
const int gpu_k_w13 = full_config.hidden_size;
|
||||
const int global_n_offset_w13 = tp_part_idx * cpu_n_w13;
|
||||
const size_t gpu_w13_weight_per_mat = static_cast<size_t>(gpu_n_w13) * gpu_k_w13 / 2;
|
||||
const size_t gpu_w13_scale_per_mat = static_cast<size_t>(gpu_n_w13) * (gpu_k_w13 / group_size);
|
||||
|
||||
const int cpu_n_w2 = config_.hidden_size;
|
||||
const int cpu_k_w2 = config_.intermediate_size;
|
||||
const int gpu_k_w2 = full_config.intermediate_size / gpu_tp_count;
|
||||
const int global_k_offset_w2 = tp_part_idx * cpu_k_w2;
|
||||
|
||||
pool->do_work_stealing_job(
|
||||
total_tasks, nullptr,
|
||||
[=, &w13_weight_ptrs, &w13_scale_ptrs, &w2_weight_ptrs, &w2_scale_ptrs, this](int task_id) {
|
||||
if (task_id < NUM_W13_TASKS) {
|
||||
const int rows_per_task = (cpu_n_w13 + NUM_W13_TASKS - 1) / NUM_W13_TASKS;
|
||||
const int row_start = task_id * rows_per_task;
|
||||
const int row_end = std::min(row_start + rows_per_task, cpu_n_w13);
|
||||
for (int local_n = row_start; local_n < row_end; local_n++) {
|
||||
const int global_n = global_n_offset_w13 + local_n;
|
||||
const int target_gpu = global_n / gpu_n_w13;
|
||||
const int n_in_gpu = global_n % gpu_n_w13;
|
||||
|
||||
uint8_t* w13_weight_base = reinterpret_cast<uint8_t*>(w13_weight_ptrs[target_gpu]);
|
||||
ggml_bf16_t* w13_scale_base = reinterpret_cast<ggml_bf16_t*>(w13_scale_ptrs[target_gpu]);
|
||||
const size_t weight_row_offset = static_cast<size_t>(n_in_gpu) * gpu_k_w13 / 2;
|
||||
const size_t scale_row_offset = static_cast<size_t>(n_in_gpu) * (gpu_k_w13 / group_size);
|
||||
|
||||
gate_bb_[expert_id]->copy_weight_rows_to(w13_weight_base + weight_row_offset, local_n, 1, 0, cpu_k_w13,
|
||||
gpu_k_w13 / 2);
|
||||
up_bb_[expert_id]->copy_weight_rows_to(w13_weight_base + gpu_w13_weight_per_mat + weight_row_offset,
|
||||
local_n, 1, 0, cpu_k_w13, gpu_k_w13 / 2);
|
||||
|
||||
gate_bb_[expert_id]->copy_scale_rows_to(w13_scale_base + scale_row_offset, local_n, 1, 0,
|
||||
cpu_k_w13 / group_size, gpu_k_w13 / group_size);
|
||||
up_bb_[expert_id]->copy_scale_rows_to(w13_scale_base + gpu_w13_scale_per_mat + scale_row_offset, local_n,
|
||||
1, 0, cpu_k_w13 / group_size, gpu_k_w13 / group_size);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const int w2_task_id = task_id - NUM_W13_TASKS;
|
||||
const int rows_per_task = (cpu_n_w2 + NUM_W2_TASKS - 1) / NUM_W2_TASKS;
|
||||
const int row_start = w2_task_id * rows_per_task;
|
||||
const int row_end = std::min(row_start + rows_per_task, cpu_n_w2);
|
||||
for (int row = row_start; row < row_end; row++) {
|
||||
int k_local = 0;
|
||||
while (k_local < cpu_k_w2) {
|
||||
const int global_k = global_k_offset_w2 + k_local;
|
||||
const int target_gpu = global_k / gpu_k_w2;
|
||||
const int k_in_gpu = global_k % gpu_k_w2;
|
||||
const int k_count = std::min(cpu_k_w2 - k_local, gpu_k_w2 - k_in_gpu);
|
||||
|
||||
uint8_t* w2_weight_base = reinterpret_cast<uint8_t*>(w2_weight_ptrs[target_gpu]);
|
||||
ggml_bf16_t* w2_scale_base = reinterpret_cast<ggml_bf16_t*>(w2_scale_ptrs[target_gpu]);
|
||||
uint8_t* weight_dst = w2_weight_base + static_cast<size_t>(row) * gpu_k_w2 / 2 + k_in_gpu / 2;
|
||||
ggml_bf16_t* scale_dst =
|
||||
w2_scale_base + static_cast<size_t>(row) * (gpu_k_w2 / group_size) + k_in_gpu / group_size;
|
||||
|
||||
down_bb_[expert_id]->copy_weight_rows_to(weight_dst, row, 1, k_local, k_count, gpu_k_w2 / 2);
|
||||
down_bb_[expert_id]->copy_scale_rows_to(scale_dst, row, 1, k_local / group_size, k_count / group_size,
|
||||
gpu_k_w2 / group_size);
|
||||
k_local += k_count;
|
||||
}
|
||||
}
|
||||
},
|
||||
nullptr);
|
||||
}
|
||||
|
||||
// Write a single expert's weights to the output buffers
|
||||
// The caller provides pointers that already point to the target expert's location (no offset needed)
|
||||
// expert_id: the index of the expert to write
|
||||
@@ -250,6 +335,12 @@ class AMX_K2_MOE_TP : public AMX_MOE_BASE<T, AMX_K2_MOE_TP<T>> {
|
||||
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 constexpr (T::BLOCKED_B_LAYOUT) {
|
||||
write_weights_to_buffer_blocked(gpu_tp_count, cpu_tp_count, expert_id, full_config, w13_weight_ptrs,
|
||||
w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs);
|
||||
return;
|
||||
}
|
||||
|
||||
const int group_size = config_.quant_config.group_size;
|
||||
auto pool = config_.pool->get_subpool(tp_part_idx);
|
||||
|
||||
|
||||
@@ -1141,6 +1141,121 @@ struct BufferBInt4KGroupImpl {
|
||||
}
|
||||
};
|
||||
|
||||
// Block-major BufferB for signed RAWINT4 with KGroup scale. The external
|
||||
// format remains row-major packed int4; internally weights are stored by
|
||||
// N_BLOCK and 64-K-element tiles for prefill-friendly access.
|
||||
template <typename K>
|
||||
struct BufferBInt4KGroupBlockedImpl {
|
||||
using dt = typename K::dt;
|
||||
dt* b;
|
||||
float* d;
|
||||
int n, k, k_group_size, k_group_count;
|
||||
|
||||
static constexpr int N_STEP = K::N_STEP;
|
||||
static constexpr int K_STEP = K::K_STEP;
|
||||
static constexpr int N_BLOCK = K::N_BLOCK;
|
||||
static constexpr int K_TILE = 64;
|
||||
static constexpr int K_TILE_BYTES = K_TILE / 2;
|
||||
static constexpr bool SCALE = true;
|
||||
|
||||
static size_t required_size(int n, int k, int k_group_size) {
|
||||
return sizeof(int8_t) * n * k / 2 + sizeof(float) * n * (k / k_group_size);
|
||||
}
|
||||
|
||||
BufferBInt4KGroupBlockedImpl(int n, int k, int k_group_size, void* ptr) : n(n), k(k), k_group_size(k_group_size) {
|
||||
assert(reinterpret_cast<intptr_t>(ptr) % 64 == 0);
|
||||
if (n % N_STEP || k % K_TILE || k % k_group_size) {
|
||||
printf("BufferBInt4KGroupBlockedImpl: n: %d, k: %d, N_STEP: %d, K_TILE: %d, k_group_size: %d\n", n, k,
|
||||
N_STEP, K_TILE, k_group_size);
|
||||
throw std::runtime_error("n or k is not aligned to blocked RAWINT4 layout");
|
||||
}
|
||||
k_group_count = k / k_group_size;
|
||||
b = reinterpret_cast<dt*>(ptr);
|
||||
d = reinterpret_cast<float*>(offset_pointer(b, n * k / 2));
|
||||
}
|
||||
|
||||
size_t block_offset_bytes(int n_begin, int k_begin) const {
|
||||
const int n_block_begin = n_begin / N_BLOCK * N_BLOCK;
|
||||
const int n_in_block = n_begin - n_block_begin;
|
||||
const int n_block_size = std::min(N_BLOCK, n - n_block_begin);
|
||||
const int k_tile = k_begin / K_TILE;
|
||||
return static_cast<size_t>(n_block_begin) * k / 2 + static_cast<size_t>(k_tile) * n_block_size * K_TILE_BYTES +
|
||||
static_cast<size_t>(n_in_block) * K_TILE_BYTES;
|
||||
}
|
||||
|
||||
void from_raw_mat(uint8_t* proj, int ith, int nth) {
|
||||
auto [n_start, n_end] = K::split_range_n(n, ith, nth);
|
||||
if (n_start >= n_end) return;
|
||||
const size_t src_row_bytes = static_cast<size_t>(k) / 2;
|
||||
for (int row = n_start; row < n_end; row++) {
|
||||
const uint8_t* src_row = proj + static_cast<size_t>(row) * src_row_bytes;
|
||||
for (int k_begin = 0; k_begin < k; k_begin += K_TILE) {
|
||||
uint8_t* dst = reinterpret_cast<uint8_t*>(b) + block_offset_bytes(row, k_begin);
|
||||
std::memcpy(dst, src_row + k_begin / 2, K_TILE_BYTES);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
dt* get_submat(int n_, int k_, int n_begin, int k_begin) {
|
||||
(void)n_;
|
||||
(void)k_;
|
||||
return reinterpret_cast<dt*>(reinterpret_cast<uint8_t*>(b) + block_offset_bytes(n_begin, k_begin));
|
||||
}
|
||||
|
||||
const uint8_t* get_kblock(int n_begin, int k_begin) const {
|
||||
return reinterpret_cast<const uint8_t*>(b) + block_offset_bytes(n_begin, k_begin);
|
||||
}
|
||||
|
||||
uint8_t* get_kblock(int n_begin, int k_begin) {
|
||||
return reinterpret_cast<uint8_t*>(b) + block_offset_bytes(n_begin, k_begin);
|
||||
}
|
||||
|
||||
float* get_scale(int n_, int n_begin, int k_, int k_begin) {
|
||||
(void)n_;
|
||||
int k_group_idx = k_begin / k_group_size;
|
||||
return d + n_begin * (k_ / k_group_size) + k_group_idx;
|
||||
}
|
||||
|
||||
const float* get_scale(int n_, int n_begin, int k_, int k_begin) const {
|
||||
(void)n_;
|
||||
int k_group_idx = k_begin / k_group_size;
|
||||
return d + n_begin * (k_ / k_group_size) + k_group_idx;
|
||||
}
|
||||
|
||||
void copy_weight_rows_to(uint8_t* dst, int row_start, int row_count, int k_start, int k_count,
|
||||
size_t dst_row_stride_bytes) const {
|
||||
assert(k_start % K_TILE == 0);
|
||||
assert(k_count % K_TILE == 0);
|
||||
for (int r = 0; r < row_count; r++) {
|
||||
uint8_t* dst_row = dst + static_cast<size_t>(r) * dst_row_stride_bytes;
|
||||
const int src_row = row_start + r;
|
||||
for (int kk = 0; kk < k_count; kk += K_TILE) {
|
||||
std::memcpy(dst_row + kk / 2, get_kblock(src_row, k_start + kk), K_TILE_BYTES);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void copy_scale_rows_to(ggml_bf16_t* dst, int row_start, int row_count, int kg_start, int kg_count,
|
||||
size_t dst_row_stride_elems) const {
|
||||
for (int r = 0; r < row_count; r++) {
|
||||
const float* src = d + static_cast<size_t>(row_start + r) * k_group_count + kg_start;
|
||||
ggml_bf16_t* dst_row = dst + static_cast<size_t>(r) * dst_row_stride_elems;
|
||||
for (int kg = 0; kg < kg_count; kg++) {
|
||||
dst_row[kg] = ggml_fp32_to_bf16(src[kg]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static std::pair<int, int> split_range_n(int n, int ith, int nth) {
|
||||
int n_per_thread = (n + nth - 1) / nth;
|
||||
n_per_thread = (n_per_thread + N_STEP - 1) / N_STEP * N_STEP;
|
||||
int n_start = std::min(ith * n_per_thread, n);
|
||||
int n_end = std::min(n_start + n_per_thread, n);
|
||||
return {n_start, n_end};
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
template <typename K>
|
||||
struct BufferBInt4WithZeroKGroupImpl {
|
||||
using dt = typename K::dt;
|
||||
|
||||
@@ -3223,6 +3223,7 @@ struct GemmKernel224Int4SmallKGroup {
|
||||
using output_t = int32_t;
|
||||
static constexpr double ELEMENT_SIZE = 0.5;
|
||||
static constexpr int VNNI_BLK = 4;
|
||||
static constexpr bool BLOCKED_B_LAYOUT = false;
|
||||
|
||||
static constexpr int M_STEP = 1;
|
||||
static constexpr int N_STEP = 32;
|
||||
@@ -3470,6 +3471,173 @@ struct GemmKernel224Int4SmallKGroup {
|
||||
}
|
||||
};
|
||||
|
||||
struct GemmKernel224Int4SmallKGroupBlocked : public GemmKernel224Int4SmallKGroup {
|
||||
static constexpr bool BLOCKED_B_LAYOUT = true;
|
||||
static std::string name() { return "K2_INT4_KGROUP_BLOCKED"; }
|
||||
|
||||
using BufferA = BufferASmallKGroupImpl<GemmKernel224Int4SmallKGroupBlocked>;
|
||||
using BufferB = BufferBInt4KGroupBlockedImpl<GemmKernel224Int4SmallKGroupBlocked>;
|
||||
using BufferC = BufferCReduceImpl<GemmKernel224Int4SmallKGroupBlocked>;
|
||||
|
||||
static inline __m256i load_packed_kblock(BufferB* bb, int n_begin, int k_begin) {
|
||||
return _mm256_loadu_si256(reinterpret_cast<const __m256i*>(bb->get_kblock(n_begin, k_begin)));
|
||||
}
|
||||
|
||||
static inline void integer_mat_vec_kgroup(int m, int n, int k, int k_group_size, BufferA* ba, BufferB* bb,
|
||||
BufferC* bc, int ith, int nth) {
|
||||
auto [n_start, n_end] = split_range_n(n, ith, nth);
|
||||
for (int m_begin = 0; m_begin < m; m_begin++) {
|
||||
float* c = bc->get_submat(m, n, m_begin, n_start);
|
||||
__m512i* a512 = (__m512i*)ba->get_submat(m, k, m_begin, 0);
|
||||
float* as = (float*)ba->get_scale(m, m_begin, k, 0);
|
||||
|
||||
for (int n_block_begin = n_start; n_block_begin < n_end; n_block_begin++) {
|
||||
float* bs = (float*)bb->get_scale(n, n_block_begin, k, 0);
|
||||
|
||||
__m512 sum = _mm512_setzero_ps();
|
||||
for (int k_block = 0; k_block < k / 64; k_block++) {
|
||||
__m256i b256 = load_packed_kblock(bb, n_block_begin, k_block * 64);
|
||||
sum = _mm512_add_ps(sum, dot_scaled_kblock(a512[k_block], b256, as[k_block * 2] * bs[k_block * 2],
|
||||
as[k_block * 2 + 1] * bs[k_block * 2 + 1]));
|
||||
}
|
||||
|
||||
c[n_block_begin - n_start] = _mm512_reduce_add_ps(sum) / 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static inline void integer_mat_mat_kgroup(int m, int n, int k, int k_group_size, BufferA* ba, BufferB* bb,
|
||||
BufferC* bc, int ith, int nth) {
|
||||
auto [n_start, n_end] = split_range_n(n, ith, nth);
|
||||
if (n_start >= n_end) return;
|
||||
|
||||
constexpr int MB = 4;
|
||||
constexpr int NB = 4;
|
||||
const int k_blocks = k / 64;
|
||||
|
||||
int m_pos = 0;
|
||||
for (; m_pos + MB <= m; m_pos += MB) {
|
||||
__m512i* a_rows[MB] = {
|
||||
(__m512i*)ba->get_submat(m, k, m_pos + 0, 0),
|
||||
(__m512i*)ba->get_submat(m, k, m_pos + 1, 0),
|
||||
(__m512i*)ba->get_submat(m, k, m_pos + 2, 0),
|
||||
(__m512i*)ba->get_submat(m, k, m_pos + 3, 0),
|
||||
};
|
||||
float* as[MB] = {
|
||||
(float*)ba->get_scale(m, m_pos + 0, k, 0),
|
||||
(float*)ba->get_scale(m, m_pos + 1, k, 0),
|
||||
(float*)ba->get_scale(m, m_pos + 2, k, 0),
|
||||
(float*)ba->get_scale(m, m_pos + 3, k, 0),
|
||||
};
|
||||
|
||||
int n_pos = n_start;
|
||||
for (; n_pos + NB <= n_end; n_pos += NB) {
|
||||
float* bs[NB] = {
|
||||
(float*)bb->get_scale(n, n_pos + 0, k, 0),
|
||||
(float*)bb->get_scale(n, n_pos + 1, k, 0),
|
||||
(float*)bb->get_scale(n, n_pos + 2, k, 0),
|
||||
(float*)bb->get_scale(n, n_pos + 3, k, 0),
|
||||
};
|
||||
|
||||
__m512 acc[MB][NB];
|
||||
for (int i = 0; i < MB; i++) {
|
||||
for (int j = 0; j < NB; j++) acc[i][j] = _mm512_setzero_ps();
|
||||
}
|
||||
|
||||
for (int k_block = 0; k_block < k_blocks; k_block++) {
|
||||
const int k_begin = k_block * 64;
|
||||
__m512i w[NB] = {
|
||||
compressed_int4_to_int8_avx512(load_packed_kblock(bb, n_pos + 0, k_begin)),
|
||||
compressed_int4_to_int8_avx512(load_packed_kblock(bb, n_pos + 1, k_begin)),
|
||||
compressed_int4_to_int8_avx512(load_packed_kblock(bb, n_pos + 2, k_begin)),
|
||||
compressed_int4_to_int8_avx512(load_packed_kblock(bb, n_pos + 3, k_begin)),
|
||||
};
|
||||
|
||||
#define K2_INT4_BLOCKED_ACCUM_ROW4(M_I) \
|
||||
do { \
|
||||
const __m512 ab0 = make_scale_pair(as[M_I][k_block * 2] * bs[0][k_block * 2], \
|
||||
as[M_I][k_block * 2 + 1] * bs[0][k_block * 2 + 1]); \
|
||||
const __m512 ab1 = make_scale_pair(as[M_I][k_block * 2] * bs[1][k_block * 2], \
|
||||
as[M_I][k_block * 2 + 1] * bs[1][k_block * 2 + 1]); \
|
||||
const __m512 ab2 = make_scale_pair(as[M_I][k_block * 2] * bs[2][k_block * 2], \
|
||||
as[M_I][k_block * 2 + 1] * bs[2][k_block * 2 + 1]); \
|
||||
const __m512 ab3 = make_scale_pair(as[M_I][k_block * 2] * bs[3][k_block * 2], \
|
||||
as[M_I][k_block * 2 + 1] * bs[3][k_block * 2 + 1]); \
|
||||
accumulate_row4(acc[M_I], a_rows[M_I][k_block], w[0], w[1], w[2], w[3], ab0, ab1, ab2, ab3); \
|
||||
} while (0)
|
||||
K2_INT4_BLOCKED_ACCUM_ROW4(0);
|
||||
K2_INT4_BLOCKED_ACCUM_ROW4(1);
|
||||
K2_INT4_BLOCKED_ACCUM_ROW4(2);
|
||||
K2_INT4_BLOCKED_ACCUM_ROW4(3);
|
||||
#undef K2_INT4_BLOCKED_ACCUM_ROW4
|
||||
}
|
||||
|
||||
for (int i = 0; i < MB; i++) {
|
||||
float* c = bc->get_submat(m, n, m_pos + i, n_start);
|
||||
store4_reduce_div16(acc[i][0], acc[i][1], acc[i][2], acc[i][3], c + (n_pos - n_start));
|
||||
}
|
||||
}
|
||||
|
||||
for (; n_pos < n_end; n_pos++) {
|
||||
float* bs = (float*)bb->get_scale(n, n_pos, k, 0);
|
||||
for (int i = 0; i < MB; i++) {
|
||||
float* c = bc->get_submat(m, n, m_pos + i, n_start);
|
||||
__m512 sum = _mm512_setzero_ps();
|
||||
for (int k_block = 0; k_block < k_blocks; k_block++) {
|
||||
__m256i b256 = load_packed_kblock(bb, n_pos, k_block * 64);
|
||||
sum = _mm512_add_ps(
|
||||
sum, dot_scaled_kblock(a_rows[i][k_block], b256, as[i][k_block * 2] * bs[k_block * 2],
|
||||
as[i][k_block * 2 + 1] * bs[k_block * 2 + 1]));
|
||||
}
|
||||
c[n_pos - n_start] = _mm512_reduce_add_ps(sum) / 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (int mi = m_pos; mi < m; mi++) {
|
||||
float* c = bc->get_submat(m, n, mi, n_start);
|
||||
__m512i* a512 = (__m512i*)ba->get_submat(m, k, mi, 0);
|
||||
float* as = (float*)ba->get_scale(m, mi, k, 0);
|
||||
int n_pos = n_start;
|
||||
for (; n_pos + NB <= n_end; n_pos += NB) {
|
||||
float* bs[NB] = {
|
||||
(float*)bb->get_scale(n, n_pos + 0, k, 0),
|
||||
(float*)bb->get_scale(n, n_pos + 1, k, 0),
|
||||
(float*)bb->get_scale(n, n_pos + 2, k, 0),
|
||||
(float*)bb->get_scale(n, n_pos + 3, k, 0),
|
||||
};
|
||||
__m512 acc[NB] = {_mm512_setzero_ps(), _mm512_setzero_ps(), _mm512_setzero_ps(), _mm512_setzero_ps()};
|
||||
for (int k_block = 0; k_block < k_blocks; k_block++) {
|
||||
const int k_begin = k_block * 64;
|
||||
__m512i w[NB] = {
|
||||
compressed_int4_to_int8_avx512(load_packed_kblock(bb, n_pos + 0, k_begin)),
|
||||
compressed_int4_to_int8_avx512(load_packed_kblock(bb, n_pos + 1, k_begin)),
|
||||
compressed_int4_to_int8_avx512(load_packed_kblock(bb, n_pos + 2, k_begin)),
|
||||
compressed_int4_to_int8_avx512(load_packed_kblock(bb, n_pos + 3, k_begin)),
|
||||
};
|
||||
for (int j = 0; j < NB; j++) {
|
||||
__m512 abscale = make_scale_pair(as[k_block * 2] * bs[j][k_block * 2],
|
||||
as[k_block * 2 + 1] * bs[j][k_block * 2 + 1]);
|
||||
acc[j] = _mm512_add_ps(acc[j], dot_scaled_decoded_kblock(a512[k_block], w[j], abscale));
|
||||
}
|
||||
}
|
||||
store4_reduce_div16(acc[0], acc[1], acc[2], acc[3], c + (n_pos - n_start));
|
||||
}
|
||||
for (; n_pos < n_end; n_pos++) {
|
||||
float* bs = (float*)bb->get_scale(n, n_pos, k, 0);
|
||||
__m512 sum = _mm512_setzero_ps();
|
||||
for (int k_block = 0; k_block < k_blocks; k_block++) {
|
||||
__m256i b256 = load_packed_kblock(bb, n_pos, k_block * 64);
|
||||
sum = _mm512_add_ps(sum, dot_scaled_kblock(a512[k_block], b256, as[k_block * 2] * bs[k_block * 2],
|
||||
as[k_block * 2 + 1] * bs[k_block * 2 + 1]));
|
||||
}
|
||||
c[n_pos - n_start] = _mm512_reduce_add_ps(sum) / 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
inline void vec_mul_kgroup(int m, int n, int k, int k_group_size,
|
||||
std::shared_ptr<GemmKernel224Int4SmallKGroup::BufferA> ba,
|
||||
std::shared_ptr<GemmKernel224Int4SmallKGroup::BufferB> bb,
|
||||
@@ -3484,6 +3652,23 @@ inline void mat_mul_kgroup(int m, int n, int k, int k_group_size,
|
||||
GemmKernel224Int4SmallKGroup::integer_mat_mat_kgroup(m, n, k, k_group_size, ba.get(), bb.get(), bc.get(), ith, nth);
|
||||
}
|
||||
|
||||
|
||||
inline void vec_mul_kgroup(int m, int n, int k, int k_group_size,
|
||||
std::shared_ptr<GemmKernel224Int4SmallKGroupBlocked::BufferA> ba,
|
||||
std::shared_ptr<GemmKernel224Int4SmallKGroupBlocked::BufferB> bb,
|
||||
std::shared_ptr<GemmKernel224Int4SmallKGroupBlocked::BufferC> bc, int ith, int nth) {
|
||||
GemmKernel224Int4SmallKGroupBlocked::integer_mat_vec_kgroup(m, n, k, k_group_size, ba.get(), bb.get(), bc.get(), ith,
|
||||
nth);
|
||||
}
|
||||
|
||||
inline void mat_mul_kgroup(int m, int n, int k, int k_group_size,
|
||||
std::shared_ptr<GemmKernel224Int4SmallKGroupBlocked::BufferA> ba,
|
||||
std::shared_ptr<GemmKernel224Int4SmallKGroupBlocked::BufferB> bb,
|
||||
std::shared_ptr<GemmKernel224Int4SmallKGroupBlocked::BufferC> bc, int ith, int nth) {
|
||||
GemmKernel224Int4SmallKGroupBlocked::integer_mat_mat_kgroup(m, n, k, k_group_size, ba.get(), bb.get(), bc.get(), ith,
|
||||
nth);
|
||||
}
|
||||
|
||||
// New k-group aware matrix multiplication function
|
||||
template <typename K, bool amx_or_avx = true>
|
||||
void integer_mat_mul_kgroup(int m, int n, int k, int k_group_size, typename K::BufferA* ba, typename K::BufferB* bb,
|
||||
|
||||
@@ -24,6 +24,7 @@ import kt_kernel_ext.moe as _moe_mod
|
||||
AMXInt4_MOE = getattr(_moe_mod, "AMXInt4_MOE", None)
|
||||
AMXInt8_MOE = getattr(_moe_mod, "AMXInt8_MOE", None)
|
||||
AMXInt4_KGroup_MOE = getattr(_moe_mod, "AMXInt4_KGroup_MOE", None)
|
||||
AMXInt4_KGroupBlocked_MOE = getattr(_moe_mod, "AMXInt4_KGroupBlocked_MOE", None)
|
||||
AMXFP4_KGroup_MOE = getattr(_moe_mod, "AMXFP4_KGroup_MOE", None)
|
||||
AMXMXFP8_KGroup_MOE = getattr(_moe_mod, "AMXMXFP8_KGroup_MOE", None)
|
||||
AMXFP8_MOE = getattr(_moe_mod, "AMXFP8_MOE", None)
|
||||
@@ -41,6 +42,7 @@ AVXVNNI256RawInt4_MOE = getattr(_moe_mod, "AVXVNNI256RawInt4_MOE", None)
|
||||
_HAS_AMXINT4_SUPPORT = AMXInt4_MOE is not None
|
||||
_HAS_AMXINT8_SUPPORT = AMXInt8_MOE is not None
|
||||
_HAS_RAWINT4_SUPPORT = AMXInt4_KGroup_MOE is not None
|
||||
_HAS_RAWINT4_BLOCKED_SUPPORT = AMXInt4_KGroupBlocked_MOE is not None
|
||||
_HAS_MXFP4_SUPPORT = AMXFP4_KGroup_MOE is not None
|
||||
_HAS_MXFP8_SUPPORT = AMXMXFP8_KGroup_MOE is not None
|
||||
_HAS_FP8_SUPPORT = AMXFP8_MOE is not None
|
||||
@@ -123,6 +125,13 @@ def _select_rawint4_backend(group_size: Optional[int] = None):
|
||||
raise RuntimeError("KT_RAWINT4_BACKEND=amx requested, but AMXInt4_KGroup_MOE is not compiled in.")
|
||||
return AMXInt4_KGroup_MOE
|
||||
|
||||
if forced in {"amx_blocked", "blocked"}:
|
||||
if not _HAS_RAWINT4_BLOCKED_SUPPORT:
|
||||
raise RuntimeError(
|
||||
"KT_RAWINT4_BACKEND=amx_blocked requested, but AMXInt4_KGroupBlocked_MOE is not compiled in."
|
||||
)
|
||||
return AMXInt4_KGroupBlocked_MOE
|
||||
|
||||
if forced in {"avxvnni", "avxvnni256"}:
|
||||
if not _HAS_AVXVNNI256_RAW_INT4_SUPPORT:
|
||||
raise RuntimeError("KT_RAWINT4_BACKEND=avxvnni requested, but AVXVNNI256RawInt4_MOE is not compiled in.")
|
||||
|
||||
@@ -25,7 +25,7 @@ expert_num = 8
|
||||
hidden_size = 256
|
||||
intermediate_size = 512
|
||||
num_experts_per_tok = 2
|
||||
max_len = 128
|
||||
max_len = 512
|
||||
group_size = 128
|
||||
validation_iter = 3
|
||||
CPUINFER_PARAM = 16
|
||||
@@ -226,7 +226,7 @@ def run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen, quant_
|
||||
)
|
||||
|
||||
config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size, 0)
|
||||
config.max_len = max_len
|
||||
config.max_len = max(max_len, qlen)
|
||||
config.gate_proj = gate_qw.data_ptr()
|
||||
config.up_proj = up_qw.data_ptr()
|
||||
config.down_proj = down_qw.data_ptr()
|
||||
@@ -296,6 +296,113 @@ def test_amxint4_kgroup_accuracy():
|
||||
run_backend_accuracy_test("AMXInt4_KGroup_MOE", backend_cls, 0.20, qlen=128, quant_group_size=32)
|
||||
|
||||
|
||||
def test_amxint4_kgroup_blocked_accuracy():
|
||||
if not hasattr(kt_kernel_ext.moe, "AMXInt4_KGroupBlocked_MOE"):
|
||||
pytest.skip("AMXInt4_KGroupBlocked_MOE is not available")
|
||||
|
||||
backend_cls = kt_kernel_ext.moe.AMXInt4_KGroupBlocked_MOE
|
||||
run_backend_accuracy_test("AMXInt4_KGroupBlocked_MOE", backend_cls, 0.20, qlen=1, quant_group_size=32)
|
||||
run_backend_accuracy_test("AMXInt4_KGroupBlocked_MOE", backend_cls, 0.20, qlen=32, quant_group_size=32)
|
||||
run_backend_accuracy_test("AMXInt4_KGroupBlocked_MOE", backend_cls, 0.20, qlen=128, quant_group_size=32)
|
||||
run_backend_accuracy_test("AMXInt4_KGroupBlocked_MOE", backend_cls, 0.20, qlen=512, quant_group_size=32)
|
||||
|
||||
|
||||
def _make_writer_inputs(quant_group_size=32):
|
||||
torch.manual_seed(1234)
|
||||
gate_qw = torch.randint(0, 256, (expert_num, intermediate_size, hidden_size // 2), dtype=torch.uint8).contiguous()
|
||||
up_qw = torch.randint(0, 256, (expert_num, intermediate_size, hidden_size // 2), dtype=torch.uint8).contiguous()
|
||||
down_qw = torch.randint(0, 256, (expert_num, hidden_size, intermediate_size // 2), dtype=torch.uint8).contiguous()
|
||||
gate_scales = torch.rand((expert_num, intermediate_size, hidden_size // quant_group_size), dtype=torch.float32).to(
|
||||
torch.bfloat16
|
||||
).contiguous()
|
||||
up_scales = torch.rand((expert_num, intermediate_size, hidden_size // quant_group_size), dtype=torch.float32).to(
|
||||
torch.bfloat16
|
||||
).contiguous()
|
||||
down_scales = torch.rand((expert_num, hidden_size, intermediate_size // quant_group_size), dtype=torch.float32).to(
|
||||
torch.bfloat16
|
||||
).contiguous()
|
||||
return gate_qw, up_qw, down_qw, gate_scales, up_scales, down_scales
|
||||
|
||||
|
||||
def _load_writer_backend(backend_cls, tensors, quant_group_size=32):
|
||||
gate_qw, up_qw, down_qw, gate_scales, up_scales, down_scales = tensors
|
||||
cpu_infer = kt_kernel_ext.CPUInfer(CPUINFER_PARAM)
|
||||
config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size, 0)
|
||||
config.max_len = max_len
|
||||
config.gate_proj = gate_qw.data_ptr()
|
||||
config.up_proj = up_qw.data_ptr()
|
||||
config.down_proj = down_qw.data_ptr()
|
||||
config.gate_scale = gate_scales.data_ptr()
|
||||
config.up_scale = up_scales.data_ptr()
|
||||
config.down_scale = down_scales.data_ptr()
|
||||
config.quant_config.bits = 4
|
||||
config.quant_config.group_size = quant_group_size
|
||||
config.quant_config.zero_point = False
|
||||
config.pool = cpu_infer.backend_
|
||||
|
||||
moe = backend_cls(config)
|
||||
physical_to_logical_map = torch.tensor(range(expert_num), dtype=torch.int64).contiguous()
|
||||
cpu_infer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
|
||||
cpu_infer.sync()
|
||||
return cpu_infer, moe
|
||||
|
||||
|
||||
def _export_writer_buffers(cpu_infer, moe, gpu_tp_count, quant_group_size=32, gpu_experts=2):
|
||||
per_mat_weight_bytes = intermediate_size * hidden_size // 2
|
||||
per_mat_scale_elems = intermediate_size * (hidden_size // quant_group_size)
|
||||
weight_bytes_per_expert_per_tp = per_mat_weight_bytes // gpu_tp_count
|
||||
scale_elems_per_expert_per_tp = per_mat_scale_elems // gpu_tp_count
|
||||
|
||||
w13_weight_bufs = [torch.empty(2 * gpu_experts * weight_bytes_per_expert_per_tp, dtype=torch.uint8) for _ in range(gpu_tp_count)]
|
||||
w13_scale_bufs = [torch.empty(2 * gpu_experts * scale_elems_per_expert_per_tp, dtype=torch.bfloat16) for _ in range(gpu_tp_count)]
|
||||
w2_weight_bufs = [torch.empty(gpu_experts * weight_bytes_per_expert_per_tp, dtype=torch.uint8) for _ in range(gpu_tp_count)]
|
||||
w2_scale_bufs = [torch.empty(gpu_experts * scale_elems_per_expert_per_tp, dtype=torch.bfloat16) for _ in range(gpu_tp_count)]
|
||||
|
||||
for expert_id in range(gpu_experts):
|
||||
w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = [], [], [], []
|
||||
for tp_idx in range(gpu_tp_count):
|
||||
w13_weight_ptrs.append(w13_weight_bufs[tp_idx].data_ptr() + expert_id * 2 * weight_bytes_per_expert_per_tp)
|
||||
w13_scale_ptrs.append(
|
||||
w13_scale_bufs[tp_idx].data_ptr() + expert_id * 2 * scale_elems_per_expert_per_tp * 2
|
||||
)
|
||||
w2_weight_ptrs.append(w2_weight_bufs[tp_idx].data_ptr() + expert_id * weight_bytes_per_expert_per_tp)
|
||||
w2_scale_ptrs.append(w2_scale_bufs[tp_idx].data_ptr() + expert_id * scale_elems_per_expert_per_tp * 2)
|
||||
|
||||
cpu_infer.submit(
|
||||
moe.write_weight_scale_to_buffer_task(
|
||||
gpu_tp_count=gpu_tp_count,
|
||||
expert_id=expert_id,
|
||||
w13_weight_ptrs=w13_weight_ptrs,
|
||||
w13_scale_ptrs=w13_scale_ptrs,
|
||||
w2_weight_ptrs=w2_weight_ptrs,
|
||||
w2_scale_ptrs=w2_scale_ptrs,
|
||||
)
|
||||
)
|
||||
cpu_infer.sync()
|
||||
|
||||
return w13_weight_bufs, w13_scale_bufs, w2_weight_bufs, w2_scale_bufs
|
||||
|
||||
|
||||
def test_amxint4_kgroup_blocked_write_buffer_matches_rowmajor_backend():
|
||||
if not hasattr(kt_kernel_ext.moe, "AMXInt4_KGroup_MOE") or not hasattr(
|
||||
kt_kernel_ext.moe, "AMXInt4_KGroupBlocked_MOE"
|
||||
):
|
||||
pytest.skip("AMX RAWINT4 row-major and blocked backends are both required")
|
||||
|
||||
tensors = _make_writer_inputs(quant_group_size=32)
|
||||
old_cpu, old_moe = _load_writer_backend(kt_kernel_ext.moe.AMXInt4_KGroup_MOE, tensors, quant_group_size=32)
|
||||
blocked_cpu, blocked_moe = _load_writer_backend(
|
||||
kt_kernel_ext.moe.AMXInt4_KGroupBlocked_MOE, tensors, quant_group_size=32
|
||||
)
|
||||
|
||||
for gpu_tp_count in (1, 2, 4):
|
||||
old_buffers = _export_writer_buffers(old_cpu, old_moe, gpu_tp_count, quant_group_size=32)
|
||||
blocked_buffers = _export_writer_buffers(blocked_cpu, blocked_moe, gpu_tp_count, quant_group_size=32)
|
||||
for old_group, blocked_group in zip(old_buffers, blocked_buffers):
|
||||
for old_buf, blocked_buf in zip(old_group, blocked_group):
|
||||
assert torch.equal(blocked_buf, old_buf)
|
||||
|
||||
|
||||
def test_compressed_loader_normalizes_int32_pack_quantized_weights():
|
||||
load_amx_utils()
|
||||
loader_mod = sys.modules["kt_kernel.utils.loader"]
|
||||
@@ -386,6 +493,17 @@ def test_rawint4_backend_selection_falls_back_to_avx2_for_large_group_size(monke
|
||||
assert amx_utils._select_rawint4_backend(128) is fake_avxvnni_backend
|
||||
|
||||
|
||||
def test_rawint4_backend_selection_accepts_forced_blocked_amx(monkeypatch):
|
||||
amx_utils = load_amx_utils()
|
||||
fake_blocked_backend = object()
|
||||
|
||||
monkeypatch.setattr(amx_utils, "AMXInt4_KGroupBlocked_MOE", fake_blocked_backend)
|
||||
monkeypatch.setattr(amx_utils, "_HAS_RAWINT4_BLOCKED_SUPPORT", True)
|
||||
monkeypatch.setenv("KT_RAWINT4_BACKEND", "amx_blocked")
|
||||
|
||||
assert amx_utils._select_rawint4_backend(32) is fake_blocked_backend
|
||||
|
||||
|
||||
def test_rawint4_backend_selection_rejects_forced_avxvnni_with_large_group_size(monkeypatch):
|
||||
amx_utils = load_amx_utils()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user