From 9f3fb6c8bbf0b877b1cfc2de4eed57119d5de8a3 Mon Sep 17 00:00:00 2001 From: yyj6666667 Date: Tue, 7 Jul 2026 15:43:14 +0000 Subject: [PATCH 1/4] [perf]: add K2 RAWINT4 prefill mat-mat dispatch --- kt-kernel/operators/amx/la/amx_kernels.hpp | 220 ++++++++++++++---- .../per_commit/test_moe_rawint4_accuracy.py | 45 ++-- 2 files changed, 203 insertions(+), 62 deletions(-) diff --git a/kt-kernel/operators/amx/la/amx_kernels.hpp b/kt-kernel/operators/amx/la/amx_kernels.hpp index cb0363ee3..25d7e587c 100644 --- a/kt-kernel/operators/amx/la/amx_kernels.hpp +++ b/kt-kernel/operators/amx/la/amx_kernels.hpp @@ -810,8 +810,8 @@ struct GemmKernel224BF { for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { int k_block_size = std::min(K_BLOCK, k - k_block_begin); for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { - ggml_bf16_t* tile_src = b + n_block_begin * k + k_block_begin * n_block_size + - n_begin * k_block_size + k_begin * N_STEP; + ggml_bf16_t* tile_src = + b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP; // Copy tile and reverse VNNI transpose (self-inverse) memcpy(tile_copy, tile_src, N_STEP * K_STEP * sizeof(ggml_bf16_t)); @@ -845,15 +845,14 @@ struct GemmKernel224BF { if (dst_nb_size <= 0) return; // Helper: compute tile pointer in a packed BF16 BB - auto tile_ptr = [](ggml_bf16_t* base, int total_n, int total_k, - int abs_n, int abs_k) -> ggml_bf16_t* { + auto tile_ptr = [](ggml_bf16_t* base, int total_n, int total_k, int abs_n, int abs_k) -> ggml_bf16_t* { int nb_begin = abs_n / N_BLOCK * N_BLOCK; int n_within = abs_n - nb_begin; int nb_size = std::min(N_BLOCK, total_n - nb_begin); int kb_begin = abs_k / K_BLOCK * K_BLOCK; int k_within = abs_k - kb_begin; - return base + nb_begin * total_k + kb_begin * nb_size + - n_within * std::min(K_BLOCK, total_k - kb_begin) + k_within * N_STEP; + return base + nb_begin * total_k + kb_begin * nb_size + n_within * std::min(K_BLOCK, total_k - kb_begin) + + k_within * N_STEP; }; alignas(64) ggml_bf16_t src_tile[N_STEP * K_STEP]; @@ -1198,8 +1197,8 @@ struct GemmKernel224Int8 { for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { int k_block_size = std::min(K_BLOCK, k - k_block_begin); for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { - int8_t* tile_src = b + n_block_begin * k + k_block_begin * n_block_size + - n_begin * k_block_size + k_begin * N_STEP; + int8_t* tile_src = + b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP; // Copy tile and reverse VNNI transpose (transpose_16x16_32bit is self-inverse) memcpy(tile_copy, tile_src, N_STEP * K_STEP); @@ -1245,15 +1244,14 @@ struct GemmKernel224Int8 { int dst_nb_size = n_end - dst_nb_begin; if (dst_nb_size <= 0) return; - auto tile_ptr = [](int8_t* base, int total_n, int total_k, - int abs_n, int abs_k) -> int8_t* { + auto tile_ptr = [](int8_t* base, int total_n, int total_k, int abs_n, int abs_k) -> int8_t* { int nb_begin = abs_n / N_BLOCK * N_BLOCK; int n_within = abs_n - nb_begin; int nb_size = std::min(N_BLOCK, total_n - nb_begin); int kb_begin = abs_k / K_BLOCK * K_BLOCK; int k_within = abs_k - kb_begin; - return base + nb_begin * total_k + kb_begin * nb_size + - n_within * std::min(K_BLOCK, total_k - kb_begin) + k_within * N_STEP; + return base + nb_begin * total_k + kb_begin * nb_size + n_within * std::min(K_BLOCK, total_k - kb_begin) + + k_within * N_STEP; }; alignas(64) int8_t tile_copy[N_STEP * K_STEP]; // 2KB un-VNNI workspace @@ -1273,8 +1271,7 @@ struct GemmKernel224Int8 { int nchunks = ncols / 16; __m512 amax[4]; - for (int c = 0; c < nchunks; c++) - amax[c] = _mm512_setzero_ps(); + for (int c = 0; c < nchunks; c++) amax[c] = _mm512_setzero_ps(); for (int src_r = 0; src_r < src.n; src_r += N_STEP) { int8_t* sp = tile_ptr(src.b, src.n, src.k, src_r, src_c); @@ -1291,18 +1288,15 @@ struct GemmKernel224Int8 { for (int c = 0; c < nchunks; c++) { __m128i i8_16 = _mm_load_si128((__m128i*)(row + c * 16)); __m512i abs_i32 = _mm512_abs_epi32(_mm512_cvtepi8_epi32(i8_16)); - amax[c] = _mm512_max_ps(amax[c], - _mm512_mul_ps(_mm512_cvtepi32_ps(abs_i32), vs)); + amax[c] = _mm512_max_ps(amax[c], _mm512_mul_ps(_mm512_cvtepi32_ps(abs_i32), vs)); } } } - for (int c = 0; c < nchunks; c++) - _mm512_store_ps(absmax_arr + buf_offset + c * 16, amax[c]); + for (int c = 0; c < nchunks; c++) _mm512_store_ps(absmax_arr + buf_offset + c * 16, amax[c]); } - for (int j = 0; j < dst_nb_size; j++) - d[dst_nb_begin + j] = absmax_arr[j] / 127.0f; + for (int j = 0; j < dst_nb_size; j++) d[dst_nb_begin + j] = absmax_arr[j] / 127.0f; // === Pass 2: register-based 16×16 sub-block transpose === alignas(64) int8_t quant_tile[N_STEP * K_STEP]; // 2KB @@ -1331,8 +1325,7 @@ struct GemmKernel224Int8 { int8_t* addr = tile_copy + (src_rb + i) * K_STEP + c_offset + src_cb; float scale = src.d[src_r + src_rb + i]; __m512i i32 = _mm512_cvtepi8_epi32(_mm_load_si128((__m128i*)addr)); - regs[i] = _mm512_castps_si512( - _mm512_mul_ps(_mm512_cvtepi32_ps(i32), _mm512_set1_ps(scale))); + regs[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_cvtepi32_ps(i32), _mm512_set1_ps(scale))); } // Transpose 16×16 in registers (32-bit element shuffle) @@ -1344,11 +1337,9 @@ struct GemmKernel224Int8 { for (int i = 0; i < 16; i++) { float sv = d[abs_dn + dest_rb + i]; float id = sv ? 1.0f / sv : 0.0f; - __m512i q = _mm512_cvtps_epi32( - _mm512_mul_ps(_mm512_castsi512_ps(regs[i]), _mm512_set1_ps(id))); - _mm_store_si128( - (__m128i*)(quant_tile + (dest_rb + i) * K_STEP + dest_cb), - _mm512_cvtsepi32_epi8(q)); + __m512i q = _mm512_cvtps_epi32(_mm512_mul_ps(_mm512_castsi512_ps(regs[i]), _mm512_set1_ps(id))); + _mm_store_si128((__m128i*)(quant_tile + (dest_rb + i) * K_STEP + dest_cb), + _mm512_cvtsepi32_epi8(q)); } } } @@ -1359,8 +1350,7 @@ struct GemmKernel224Int8 { transpose_16x16_32bit((__m512i*)(quant_tile + TILE_N * K_STEP)); // Write to dest BB - int8_t* dp = b + dst_nb_begin * k + dk_block * dst_nb_size + - dn * dk_block_size + dk * N_STEP; + int8_t* dp = b + dst_nb_begin * k + dk_block * dst_nb_size + dn * dk_block_size + dk * N_STEP; memcpy(dp, quant_tile, N_STEP * K_STEP); } } @@ -3282,38 +3272,180 @@ struct GemmKernel224Int4SmallKGroup { const __m512i lane_shuffle = _mm512_set_epi64(7, 6, 3, 2, 5, 4, 1, 0); return _mm512_permutexvar_epi64(lane_shuffle, result); } + static inline __m512 dot_scaled_kblock(__m512i a512, __m256i b256, float scale0, float scale1) { + __m256 abscale0 = _mm256_set1_ps(scale0); + __m256 abscale1 = _mm256_set1_ps(scale1); + __m512 abscale = _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1); + __m512i mul = _mm512_setzero_si512(); + mul = _mm512_dpbssd_epi32(mul, a512, compressed_int4_to_int8_avx512(b256)); + return _mm512_mul_ps(abscale, _mm512_cvtepi32_ps(mul)); + } + + static inline void store4_reduce_div16(__m512 s0, __m512 s1, __m512 s2, __m512 s3, float* dst) { + dst[0] = _mm512_reduce_add_ps(s0) / 16; + dst[1] = _mm512_reduce_add_ps(s1) / 16; + dst[2] = _mm512_reduce_add_ps(s2) / 16; + dst[3] = _mm512_reduce_add_ps(s3) / 16; + } + 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++) { __m256i* b256 = (__m256i*)bb->get_submat(n, k, n_block_begin, 0); - float* as = (float*)ba->get_scale(m, m_begin, k, 0); float* bs = (float*)bb->get_scale(n, n_block_begin, k, 0); __m512 sum = _mm512_setzero_ps(); -#define WORK_K_BLOCK(k_block) \ - { \ - __m256 abscale0 = _mm256_set1_ps(as[(k_block) * 2] * bs[(k_block) * 2]); \ - __m256 abscale1 = _mm256_set1_ps(as[(k_block) * 2 + 1] * bs[(k_block) * 2 + 1]); \ - __m512 abscale = _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1); \ - __m512i mul = _mm512_setzero_si512(); \ - mul = _mm512_dpbssd_epi32(mul, a512[k_block], compressed_int4_to_int8_avx512(b256[k_block])); \ - sum = _mm512_add_ps(sum, _mm512_mul_ps(abscale, _mm512_cvtepi32_ps(mul))); \ - } - - for (int k_block = 0; k_block < k / 64; k_block += 2) { - WORK_K_BLOCK(k_block); - WORK_K_BLOCK(k_block + 1); + for (int k_block = 0; k_block < k / 64; k_block++) { + sum = _mm512_add_ps(sum, dot_scaled_kblock(a512[k_block], b256[k_block], 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) { + __m256i* b_rows[NB] = { + (__m256i*)bb->get_submat(n, k, n_pos + 0, 0), + (__m256i*)bb->get_submat(n, k, n_pos + 1, 0), + (__m256i*)bb->get_submat(n, k, n_pos + 2, 0), + (__m256i*)bb->get_submat(n, k, n_pos + 3, 0), + }; + 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++) { + __m512i w[NB] = { + compressed_int4_to_int8_avx512(b_rows[0][k_block]), + compressed_int4_to_int8_avx512(b_rows[1][k_block]), + compressed_int4_to_int8_avx512(b_rows[2][k_block]), + compressed_int4_to_int8_avx512(b_rows[3][k_block]), + }; + + for (int i = 0; i < MB; i++) { + for (int j = 0; j < NB; j++) { + __m256 abscale0 = _mm256_set1_ps(as[i][k_block * 2] * bs[j][k_block * 2]); + __m256 abscale1 = _mm256_set1_ps(as[i][k_block * 2 + 1] * bs[j][k_block * 2 + 1]); + __m512 abscale = _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1); + __m512i mul = _mm512_setzero_si512(); + mul = _mm512_dpbssd_epi32(mul, a_rows[i][k_block], w[j]); + acc[i][j] = _mm512_add_ps(acc[i][j], _mm512_mul_ps(abscale, _mm512_cvtepi32_ps(mul))); + } + } + } + + 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++) { + __m256i* b256 = (__m256i*)bb->get_submat(n, k, n_pos, 0); + 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++) { + sum = _mm512_add_ps( + sum, dot_scaled_kblock(a_rows[i][k_block], b256[k_block], 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) { + __m256i* b_rows[NB] = { + (__m256i*)bb->get_submat(n, k, n_pos + 0, 0), + (__m256i*)bb->get_submat(n, k, n_pos + 1, 0), + (__m256i*)bb->get_submat(n, k, n_pos + 2, 0), + (__m256i*)bb->get_submat(n, k, n_pos + 3, 0), + }; + 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++) { + __m512i w[NB] = { + compressed_int4_to_int8_avx512(b_rows[0][k_block]), + compressed_int4_to_int8_avx512(b_rows[1][k_block]), + compressed_int4_to_int8_avx512(b_rows[2][k_block]), + compressed_int4_to_int8_avx512(b_rows[3][k_block]), + }; + for (int j = 0; j < NB; j++) { + __m256 abscale0 = _mm256_set1_ps(as[k_block * 2] * bs[j][k_block * 2]); + __m256 abscale1 = _mm256_set1_ps(as[k_block * 2 + 1] * bs[j][k_block * 2 + 1]); + __m512 abscale = _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1); + __m512i mul = _mm512_setzero_si512(); + mul = _mm512_dpbssd_epi32(mul, a512[k_block], w[j]); + acc[j] = _mm512_add_ps(acc[j], _mm512_mul_ps(abscale, _mm512_cvtepi32_ps(mul))); + } + } + store4_reduce_div16(acc[0], acc[1], acc[2], acc[3], c + (n_pos - n_start)); + } + for (; n_pos < n_end; n_pos++) { + __m256i* b256 = (__m256i*)bb->get_submat(n, k, n_pos, 0); + 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++) { + sum = _mm512_add_ps(sum, dot_scaled_kblock(a512[k_block], b256[k_block], 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, @@ -3327,7 +3459,7 @@ inline void mat_mul_kgroup(int m, int n, int k, int k_group_size, std::shared_ptr ba, std::shared_ptr bb, std::shared_ptr bc, int ith, int nth) { - GemmKernel224Int4SmallKGroup::integer_mat_vec_kgroup(m, n, k, k_group_size, ba.get(), bb.get(), bc.get(), ith, nth); + GemmKernel224Int4SmallKGroup::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 diff --git a/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py b/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py index 36054670a..8e461b772 100644 --- a/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py +++ b/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py @@ -66,20 +66,20 @@ def load_amx_utils(): return sys.modules["kt_kernel.utils.amx"] -def rawint4_quantize(weight_bf16): +def rawint4_quantize(weight_bf16, quant_group_size=group_size): """Quantize [N, K] BF16 weight to RAWINT4 layout.""" n, k = weight_bf16.shape assert k % 2 == 0 - assert k % group_size == 0 + assert k % quant_group_size == 0 weight_fp32 = weight_bf16.float() qweight = torch.zeros((n, k // 2), dtype=torch.uint8) - scales = torch.zeros((n, k // group_size), dtype=torch.bfloat16) + scales = torch.zeros((n, k // quant_group_size), dtype=torch.bfloat16) for ni in range(n): - for g in range(k // group_size): - k_start = g * group_size - k_end = k_start + group_size + for g in range(k // quant_group_size): + k_start = g * quant_group_size + k_end = k_start + quant_group_size block = weight_fp32[ni, k_start:k_end] amax = block.abs().max().item() scale = amax / 7.0 if amax > 0 else 1.0 @@ -95,14 +95,14 @@ def rawint4_quantize(weight_bf16): return qweight, scales -def rawint4_dequantize(qweight, scales, out_features, in_features): +def rawint4_dequantize(qweight, scales, out_features, in_features, quant_group_size=group_size): """Dequantize RAWINT4 qweight/scales back to fp32 [N, K].""" result = torch.zeros((out_features, in_features), dtype=torch.float32) for ni in range(out_features): - for g in range(in_features // group_size): + for g in range(in_features // quant_group_size): scale = scales[ni, g].float().item() - k_start = g * group_size - k_end = k_start + group_size + k_start = g * quant_group_size + k_end = k_start + quant_group_size for kk in range(k_start, k_end, 2): packed = int(qweight[ni, kk // 2].item()) result[ni, kk] = ((packed & 0x0F) - 8) * scale @@ -160,7 +160,7 @@ def available_backends(): return backends -def run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen): +def run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen, quant_group_size=group_size): physical_to_logical_map = torch.tensor(range(expert_num), dtype=torch.int64).contiguous() cpu_infer = kt_kernel_ext.CPUInfer(CPUINFER_PARAM) @@ -180,15 +180,15 @@ def run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen): down_qw_list, down_scale_list = [], [] for e in range(expert_num): - qw, sc = rawint4_quantize(gate_bf16[e]) + qw, sc = rawint4_quantize(gate_bf16[e], quant_group_size) gate_qw_list.append(qw) gate_scale_list.append(sc) - qw, sc = rawint4_quantize(up_bf16[e]) + qw, sc = rawint4_quantize(up_bf16[e], quant_group_size) up_qw_list.append(qw) up_scale_list.append(sc) - qw, sc = rawint4_quantize(down_bf16[e]) + qw, sc = rawint4_quantize(down_bf16[e], quant_group_size) down_qw_list.append(qw) down_scale_list.append(sc) @@ -201,19 +201,19 @@ def run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen): gate_deq = torch.stack( [ - rawint4_dequantize(gate_qw_list[e], gate_scale_list[e], intermediate_size, hidden_size) + rawint4_dequantize(gate_qw_list[e], gate_scale_list[e], intermediate_size, hidden_size, quant_group_size) for e in range(expert_num) ] ) up_deq = torch.stack( [ - rawint4_dequantize(up_qw_list[e], up_scale_list[e], intermediate_size, hidden_size) + rawint4_dequantize(up_qw_list[e], up_scale_list[e], intermediate_size, hidden_size, quant_group_size) for e in range(expert_num) ] ) down_deq = torch.stack( [ - rawint4_dequantize(down_qw_list[e], down_scale_list[e], hidden_size, intermediate_size) + rawint4_dequantize(down_qw_list[e], down_scale_list[e], hidden_size, intermediate_size, quant_group_size) for e in range(expert_num) ] ) @@ -227,7 +227,7 @@ def run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen): config.up_scale = up_scales.data_ptr() config.down_scale = down_scales.data_ptr() config.quant_config.bits = 4 - config.quant_config.group_size = group_size + config.quant_config.group_size = quant_group_size config.quant_config.zero_point = False config.pool = cpu_infer.backend_ @@ -279,6 +279,15 @@ def test_rawint4_accuracy(): run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen=16) +def test_amxint4_kgroup_accuracy(): + if not hasattr(kt_kernel_ext.moe, "AMXInt4_KGroup_MOE"): + pytest.skip("AMXInt4_KGroup_MOE is not available") + + backend_cls = kt_kernel_ext.moe.AMXInt4_KGroup_MOE + run_backend_accuracy_test("AMXInt4_KGroup_MOE", backend_cls, 0.20, qlen=1, quant_group_size=32) + run_backend_accuracy_test("AMXInt4_KGroup_MOE", backend_cls, 0.20, qlen=32, quant_group_size=32) + + def test_rawint4_backend_selection_falls_back_to_avx2_for_large_group_size(monkeypatch): amx_utils = load_amx_utils() fake_amx_backend = object() From 9b0ff46f4cbcf5505006d367299377b1a9d5cac5 Mon Sep 17 00:00:00 2001 From: yyj Date: Wed, 8 Jul 2026 11:42:28 +0800 Subject: [PATCH 2/4] [perf] Refactor K2 RAWINT4 prefill accumulation --- kt-kernel/operators/amx/la/amx_kernels.hpp | 48 ++++++++++++++----- .../per_commit/test_moe_rawint4_accuracy.py | 1 + 2 files changed, 36 insertions(+), 13 deletions(-) diff --git a/kt-kernel/operators/amx/la/amx_kernels.hpp b/kt-kernel/operators/amx/la/amx_kernels.hpp index 25d7e587c..b681c8805 100644 --- a/kt-kernel/operators/amx/la/amx_kernels.hpp +++ b/kt-kernel/operators/amx/la/amx_kernels.hpp @@ -3272,15 +3272,22 @@ struct GemmKernel224Int4SmallKGroup { const __m512i lane_shuffle = _mm512_set_epi64(7, 6, 3, 2, 5, 4, 1, 0); return _mm512_permutexvar_epi64(lane_shuffle, result); } - static inline __m512 dot_scaled_kblock(__m512i a512, __m256i b256, float scale0, float scale1) { + static inline __m512 make_scale_pair(float scale0, float scale1) { __m256 abscale0 = _mm256_set1_ps(scale0); __m256 abscale1 = _mm256_set1_ps(scale1); - __m512 abscale = _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1); + return _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1); + } + + static inline __m512 dot_scaled_decoded_kblock(__m512i a512, __m512i w512, __m512 abscale) { __m512i mul = _mm512_setzero_si512(); - mul = _mm512_dpbssd_epi32(mul, a512, compressed_int4_to_int8_avx512(b256)); + mul = _mm512_dpbssd_epi32(mul, a512, w512); return _mm512_mul_ps(abscale, _mm512_cvtepi32_ps(mul)); } + static inline __m512 dot_scaled_kblock(__m512i a512, __m256i b256, float scale0, float scale1) { + return dot_scaled_decoded_kblock(a512, compressed_int4_to_int8_avx512(b256), make_scale_pair(scale0, scale1)); + } + static inline void store4_reduce_div16(__m512 s0, __m512 s1, __m512 s2, __m512 s3, float* dst) { dst[0] = _mm512_reduce_add_ps(s0) / 16; dst[1] = _mm512_reduce_add_ps(s1) / 16; @@ -3288,6 +3295,14 @@ struct GemmKernel224Int4SmallKGroup { dst[3] = _mm512_reduce_add_ps(s3) / 16; } + static inline void accumulate_row4(__m512* acc, __m512i a512, __m512i w0, __m512i w1, __m512i w2, __m512i w3, + __m512 abscale0, __m512 abscale1, __m512 abscale2, __m512 abscale3) { + acc[0] = _mm512_add_ps(acc[0], dot_scaled_decoded_kblock(a512, w0, abscale0)); + acc[1] = _mm512_add_ps(acc[1], dot_scaled_decoded_kblock(a512, w1, abscale1)); + acc[2] = _mm512_add_ps(acc[2], dot_scaled_decoded_kblock(a512, w2, abscale2)); + acc[3] = _mm512_add_ps(acc[3], dot_scaled_decoded_kblock(a512, w3, abscale3)); + } + 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); @@ -3363,16 +3378,23 @@ struct GemmKernel224Int4SmallKGroup { compressed_int4_to_int8_avx512(b_rows[3][k_block]), }; - for (int i = 0; i < MB; i++) { - for (int j = 0; j < NB; j++) { - __m256 abscale0 = _mm256_set1_ps(as[i][k_block * 2] * bs[j][k_block * 2]); - __m256 abscale1 = _mm256_set1_ps(as[i][k_block * 2 + 1] * bs[j][k_block * 2 + 1]); - __m512 abscale = _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1); - __m512i mul = _mm512_setzero_si512(); - mul = _mm512_dpbssd_epi32(mul, a_rows[i][k_block], w[j]); - acc[i][j] = _mm512_add_ps(acc[i][j], _mm512_mul_ps(abscale, _mm512_cvtepi32_ps(mul))); - } - } +#define K2_INT4_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_ACCUM_ROW4(0); + K2_INT4_ACCUM_ROW4(1); + K2_INT4_ACCUM_ROW4(2); + K2_INT4_ACCUM_ROW4(3); +#undef K2_INT4_ACCUM_ROW4 } for (int i = 0; i < MB; i++) { diff --git a/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py b/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py index 0c63f3b31..03ae6e2dd 100644 --- a/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py +++ b/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py @@ -293,6 +293,7 @@ def test_amxint4_kgroup_accuracy(): backend_cls = kt_kernel_ext.moe.AMXInt4_KGroup_MOE run_backend_accuracy_test("AMXInt4_KGroup_MOE", backend_cls, 0.20, qlen=1, quant_group_size=32) run_backend_accuracy_test("AMXInt4_KGroup_MOE", backend_cls, 0.20, qlen=32, quant_group_size=32) + run_backend_accuracy_test("AMXInt4_KGroup_MOE", backend_cls, 0.20, qlen=128, quant_group_size=32) def test_compressed_loader_normalizes_int32_pack_quantized_weights(): From be40f8e25d8bb7a8ef71aafdddf9495a95109296 Mon Sep 17 00:00:00 2001 From: yyj Date: Wed, 8 Jul 2026 17:15:42 +0800 Subject: [PATCH 3/4] Add blocked RAWINT4 prefill backend --- kt-kernel/ext_bindings.cpp | 1 + kt-kernel/operators/amx/k2-moe.hpp | 91 +++++++++ kt-kernel/operators/amx/la/amx_buffers.hpp | 115 +++++++++++ kt-kernel/operators/amx/la/amx_kernels.hpp | 185 ++++++++++++++++++ kt-kernel/python/utils/amx.py | 9 + .../per_commit/test_moe_rawint4_accuracy.py | 122 +++++++++++- 6 files changed, 521 insertions(+), 2 deletions(-) diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 87074e417..ca35d4262 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -791,6 +791,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { bind_moe_module>(moe_module, "AMXInt4_1_MOE"); bind_moe_module>(moe_module, "AMXInt4_1KGroup_MOE"); bind_moe_module>(moe_module, "AMXInt4_KGroup_MOE"); + bind_moe_module>(moe_module, "AMXInt4_KGroupBlocked_MOE"); #if defined(__AVX512F__) bind_moe_module>(moe_module, "AMXBF16_MOE"); bind_moe_module>(moe_module, "AMXFP8_MOE"); diff --git a/kt-kernel/operators/amx/k2-moe.hpp b/kt-kernel/operators/amx/k2-moe.hpp index b8b8a05a1..63f9f8d12 100644 --- a/kt-kernel/operators/amx/k2-moe.hpp +++ b/kt-kernel/operators/amx/k2-moe.hpp @@ -241,6 +241,91 @@ class AMX_K2_MOE_TP : public AMX_MOE_BASE> { } } + + 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& w13_weight_ptrs, + const std::vector& w13_scale_ptrs, + const std::vector& w2_weight_ptrs, + const std::vector& 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(gpu_n_w13) * gpu_k_w13 / 2; + const size_t gpu_w13_scale_per_mat = static_cast(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(w13_weight_ptrs[target_gpu]); + ggml_bf16_t* w13_scale_base = reinterpret_cast(w13_scale_ptrs[target_gpu]); + const size_t weight_row_offset = static_cast(n_in_gpu) * gpu_k_w13 / 2; + const size_t scale_row_offset = static_cast(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(w2_weight_ptrs[target_gpu]); + ggml_bf16_t* w2_scale_base = reinterpret_cast(w2_scale_ptrs[target_gpu]); + uint8_t* weight_dst = w2_weight_base + static_cast(row) * gpu_k_w2 / 2 + k_in_gpu / 2; + ggml_bf16_t* scale_dst = + w2_scale_base + static_cast(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> { const std::vector& w13_scale_ptrs, const std::vector& w2_weight_ptrs, const std::vector& 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); diff --git a/kt-kernel/operators/amx/la/amx_buffers.hpp b/kt-kernel/operators/amx/la/amx_buffers.hpp index 1d0b8ac64..a3f214e50 100644 --- a/kt-kernel/operators/amx/la/amx_buffers.hpp +++ b/kt-kernel/operators/amx/la/amx_buffers.hpp @@ -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 +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(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(ptr); + d = reinterpret_cast(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(n_block_begin) * k / 2 + static_cast(k_tile) * n_block_size * K_TILE_BYTES + + static_cast(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(k) / 2; + for (int row = n_start; row < n_end; row++) { + const uint8_t* src_row = proj + static_cast(row) * src_row_bytes; + for (int k_begin = 0; k_begin < k; k_begin += K_TILE) { + uint8_t* dst = reinterpret_cast(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(reinterpret_cast(b) + block_offset_bytes(n_begin, k_begin)); + } + + const uint8_t* get_kblock(int n_begin, int k_begin) const { + return reinterpret_cast(b) + block_offset_bytes(n_begin, k_begin); + } + + uint8_t* get_kblock(int n_begin, int k_begin) { + return reinterpret_cast(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(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(row_start + r) * k_group_count + kg_start; + ggml_bf16_t* dst_row = dst + static_cast(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 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 struct BufferBInt4WithZeroKGroupImpl { using dt = typename K::dt; diff --git a/kt-kernel/operators/amx/la/amx_kernels.hpp b/kt-kernel/operators/amx/la/amx_kernels.hpp index b681c8805..0da7891ae 100644 --- a/kt-kernel/operators/amx/la/amx_kernels.hpp +++ b/kt-kernel/operators/amx/la/amx_kernels.hpp @@ -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; + using BufferB = BufferBInt4KGroupBlockedImpl; + using BufferC = BufferCReduceImpl; + + static inline __m256i load_packed_kblock(BufferB* bb, int n_begin, int k_begin) { + return _mm256_loadu_si256(reinterpret_cast(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 ba, std::shared_ptr 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 ba, + std::shared_ptr bb, + std::shared_ptr 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 ba, + std::shared_ptr bb, + std::shared_ptr 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 void integer_mat_mul_kgroup(int m, int n, int k, int k_group_size, typename K::BufferA* ba, typename K::BufferB* bb, diff --git a/kt-kernel/python/utils/amx.py b/kt-kernel/python/utils/amx.py index 1b04e4def..23efbff2d 100644 --- a/kt-kernel/python/utils/amx.py +++ b/kt-kernel/python/utils/amx.py @@ -24,6 +24,7 @@ AMXInt4_MOE = getattr(_moe_mod, "AMXInt4_MOE", None) AMXInt8_MOE = getattr(_moe_mod, "AMXInt8_MOE", None) AMXInt4_KGroup_MOE = getattr(_moe_mod, "AMXInt4_KGroup_MOE", None) +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 @@ _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.") diff --git a/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py b/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py index 03ae6e2dd..775815531 100644 --- a/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py +++ b/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py @@ -25,7 +25,7 @@ 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() From 385f5b1db5c82b9492cb9fe8d1452a938ff3deba Mon Sep 17 00:00:00 2001 From: yyj Date: Wed, 8 Jul 2026 17:45:48 +0800 Subject: [PATCH 4/4] Optimize blocked RAWINT4 decode vecmul --- kt-kernel/operators/amx/la/amx_kernels.hpp | 43 +++++++++++++++++++--- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/kt-kernel/operators/amx/la/amx_kernels.hpp b/kt-kernel/operators/amx/la/amx_kernels.hpp index 0da7891ae..4bc55ddd2 100644 --- a/kt-kernel/operators/amx/la/amx_kernels.hpp +++ b/kt-kernel/operators/amx/la/amx_kernels.hpp @@ -3486,22 +3486,55 @@ struct GemmKernel224Int4SmallKGroupBlocked : public GemmKernel224Int4SmallKGroup 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); + if (n_start >= n_end) return; + + constexpr int NB = 4; + const int k_blocks = k / 64; 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); + 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 / 64; k_block++) { - __m256i b256 = load_packed_kblock(bb, n_block_begin, k_block * 64); + 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_block_begin - n_start] = _mm512_reduce_add_ps(sum) / 16; + c[n_pos - n_start] = _mm512_reduce_add_ps(sum) / 16; } } }