[perf]: add K2 RAWINT4 prefill mat-mat dispatch#2080
Open
yyj6666667 wants to merge 5 commits into
Open
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a new register-blocked matrix-matrix multiplication kernel integer_mat_mat_kgroup for AMX Int4 quantization, refactors integer_mat_vec_kgroup using helper functions, and adds corresponding accuracy tests. The reviewer pointed out potential alignment issues in both kernels where submatrix pointers are directly cast to SIMD pointer types (__m512i* and __m256i*), which could lead to undefined behavior or segmentation faults. They recommended using unaligned load intrinsics (_mm512_loadu_si512 and _mm256_loadu_si256) to safely load the data.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
3291
to
3312
| 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; | ||
| } | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
Casting ba->get_submat and bb->get_submat directly to __m512i* and __m256i* and dereferencing them assumes strict 64-byte and 32-byte alignment. If the alignment is not met at runtime (e.g., when k is a multiple of 32 but not 64, or on arbitrary column offsets), this can lead to undefined behavior or segmentation faults due to alignment violations (e.g., when the compiler emits aligned load instructions like vmovdqa64).\n\nUsing _mm512_loadu_si512 and _mm256_loadu_si256 is the standard, safe, and idiomatic way to load SIMD registers from potentially unaligned memory with zero overhead on modern CPUs.
static inline void integer_mat_vec_kgroup(int m, int n, int k, int k_group_size, BufferA* ba, BufferB* bb,\n BufferC* bc, int ith, int nth) {\n auto [n_start, n_end] = split_range_n(n, ith, nth);\n for (int m_begin = 0; m_begin < m; m_begin++) {\n float* c = bc->get_submat(m, n, m_begin, n_start);\n int8_t* a_ptr = ba->get_submat(m, k, m_begin, 0);\n float* as = (float*)ba->get_scale(m, m_begin, k, 0);\n\n for (int n_block_begin = n_start; n_block_begin < n_end; n_block_begin++) {\n uint8_t* b_ptr = bb->get_submat(n, k, n_block_begin, 0);\n float* bs = (float*)bb->get_scale(n, n_block_begin, k, 0);\n\n __m512 sum = _mm512_setzero_ps();\n for (int k_block = 0; k_block < k / 64; k_block++) {\n __m512i a_val = _mm512_loadu_si512((const __m512i*)(a_ptr + k_block * 64));\n __m256i b_val = _mm256_loadu_si256((const __m256i*)(b_ptr + k_block * 32));\n sum = _mm512_add_ps(sum, dot_scaled_kblock(a_val, b_val, as[k_block * 2] * bs[k_block * 2],\n as[k_block * 2 + 1] * bs[k_block * 2 + 1]));\n }\n\n c[n_block_begin - n_start] = _mm512_reduce_add_ps(sum) / 16;\n }\n }\n }
Comment on lines
+3314
to
+3448
| 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; | ||
| } | ||
| } | ||
| } |
Contributor
There was a problem hiding this comment.
Casting the submatrix pointers directly to __m512i* and __m256i* and dereferencing them (a_rows[i][k_block] and b_rows[j][k_block]) assumes strict 64-byte and 32-byte alignment. If the alignment is not met at runtime, this can lead to undefined behavior or segmentation faults due to alignment violations.\n\nUsing _mm512_loadu_si512 and _mm256_loadu_si256 is the standard, safe, and idiomatic way to load SIMD registers from potentially unaligned memory with zero overhead on modern CPUs.
static inline void integer_mat_mat_kgroup(int m, int n, int k, int k_group_size, BufferA* ba, BufferB* bb,\n BufferC* bc, int ith, int nth) {\n auto [n_start, n_end] = split_range_n(n, ith, nth);\n if (n_start >= n_end) return;\n\n constexpr int MB = 4;\n constexpr int NB = 4;\n const int k_blocks = k / 64;\n\n int m_pos = 0;\n for (; m_pos + MB <= m; m_pos += MB) {\n int8_t* a_rows[MB] = {\n ba->get_submat(m, k, m_pos + 0, 0),\n ba->get_submat(m, k, m_pos + 1, 0),\n ba->get_submat(m, k, m_pos + 2, 0),\n ba->get_submat(m, k, m_pos + 3, 0),\n };\n float* as[MB] = {\n (float*)ba->get_scale(m, m_pos + 0, k, 0),\n (float*)ba->get_scale(m, m_pos + 1, k, 0),\n (float*)ba->get_scale(m, m_pos + 2, k, 0),\n (float*)ba->get_scale(m, m_pos + 3, k, 0),\n };\n\n int n_pos = n_start;\n for (; n_pos + NB <= n_end; n_pos += NB) {\n uint8_t* b_rows[NB] = {\n bb->get_submat(n, k, n_pos + 0, 0),\n bb->get_submat(n, k, n_pos + 1, 0),\n bb->get_submat(n, k, n_pos + 2, 0),\n bb->get_submat(n, k, n_pos + 3, 0),\n };\n float* bs[NB] = {\n (float*)bb->get_scale(n, n_pos + 0, k, 0),\n (float*)bb->get_scale(n, n_pos + 1, k, 0),\n (float*)bb->get_scale(n, n_pos + 2, k, 0),\n (float*)bb->get_scale(n, n_pos + 3, k, 0),\n };\n\n __m512 acc[MB][NB];\n for (int i = 0; i < MB; i++) {\n for (int j = 0; j < NB; j++) acc[i][j] = _mm512_setzero_ps();\n }\n\n for (int k_block = 0; k_block < k_blocks; k_block++) {\n __m512i w[NB] = {\n compressed_int4_to_int8_avx512(_mm256_loadu_si256((const __m256i*)(b_rows[0] + k_block * 32))),\n compressed_int4_to_int8_avx512(_mm256_loadu_si256((const __m256i*)(b_rows[1] + k_block * 32))),\n compressed_int4_to_int8_avx512(_mm256_loadu_si256((const __m256i*)(b_rows[2] + k_block * 32))),\n compressed_int4_to_int8_avx512(_mm256_loadu_si256((const __m256i*)(b_rows[3] + k_block * 32))),\n };\n\n for (int i = 0; i < MB; i++) {\n __m512i a_val = _mm512_loadu_si512((const __m512i*)(a_rows[i] + k_block * 64));\n for (int j = 0; j < NB; j++) {\n __m256 abscale0 = _mm256_set1_ps(as[i][k_block * 2] * bs[j][k_block * 2]);\n __m256 abscale1 = _mm256_set1_ps(as[i][k_block * 2 + 1] * bs[j][k_block * 2 + 1]);\n __m512 abscale = _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1);\n __m512i mul = _mm512_setzero_si512();\n mul = _mm512_dpbssd_epi32(mul, a_val, w[j]);\n acc[i][j] = _mm512_add_ps(acc[i][j], _mm512_mul_ps(abscale, _mm512_cvtepi32_ps(mul)));\n }\n }\n }\n\n for (int i = 0; i < MB; i++) {\n float* c = bc->get_submat(m, n, m_pos + i, n_start);\n store4_reduce_div16(acc[i][0], acc[i][1], acc[i][2], acc[i][3], c + (n_pos - n_start));\n }\n }\n\n for (; n_pos < n_end; n_pos++) {\n uint8_t* b_ptr = bb->get_submat(n, k, n_pos, 0);\n float* bs = (float*)bb->get_scale(n, n_pos, k, 0);\n for (int i = 0; i < MB; i++) {\n float* c = bc->get_submat(m, n, m_pos + i, n_start);\n __m512 sum = _mm512_setzero_ps();\n for (int k_block = 0; k_block < k_blocks; k_block++) {\n __m512i a_val = _mm512_loadu_si512((const __m512i*)(a_rows[i] + k_block * 64));\n __m256i b_val = _mm256_loadu_si256((const __m256i*)(b_ptr + k_block * 32));\n sum = _mm512_add_ps(\n sum, dot_scaled_kblock(a_val, b_val, as[i][k_block * 2] * bs[k_block * 2],\n as[i][k_block * 2 + 1] * bs[k_block * 2 + 1]));\n }\n c[n_pos - n_start] = _mm512_reduce_add_ps(sum) / 16;\n }\n }\n }\n\n for (int mi = m_pos; mi < m; mi++) {\n float* c = bc->get_submat(m, n, mi, n_start);\n int8_t* a_ptr = ba->get_submat(m, k, mi, 0);\n float* as = (float*)ba->get_scale(m, mi, k, 0);\n int n_pos = n_start;\n for (; n_pos + NB <= n_end; n_pos += NB) {\n uint8_t* b_rows[NB] = {\n bb->get_submat(n, k, n_pos + 0, 0),\n bb->get_submat(n, k, n_pos + 1, 0),\n bb->get_submat(n, k, n_pos + 2, 0),\n bb->get_submat(n, k, n_pos + 3, 0),\n };\n float* bs[NB] = {\n (float*)bb->get_scale(n, n_pos + 0, k, 0),\n (float*)bb->get_scale(n, n_pos + 1, k, 0),\n (float*)bb->get_scale(n, n_pos + 2, k, 0),\n (float*)bb->get_scale(n, n_pos + 3, k, 0),\n };\n __m512 acc[NB] = {_mm512_setzero_ps(), _mm512_setzero_ps(), _mm512_setzero_ps(), _mm512_setzero_ps()};\n for (int k_block = 0; k_block < k_blocks; k_block++) {\n __m512i w[NB] = {\n compressed_int4_to_int8_avx512(_mm256_loadu_si256((const __m256i*)(b_rows[0] + k_block * 32))),\n compressed_int4_to_int8_avx512(_mm256_loadu_si256((const __m256i*)(b_rows[1] + k_block * 32))),\n compressed_int4_to_int8_avx512(_mm256_loadu_si256((const __m256i*)(b_rows[2] + k_block * 32))),\n compressed_int4_to_int8_avx512(_mm256_loadu_si256((const __m256i*)(b_rows[3] + k_block * 32))),\n };\n __m512i a_val = _mm512_loadu_si512((const __m512i*)(a_ptr + k_block * 64));\n for (int j = 0; j < NB; j++) {\n __m256 abscale0 = _mm256_set1_ps(as[k_block * 2] * bs[j][k_block * 2]);\n __m256 abscale1 = _mm256_set1_ps(as[k_block * 2 + 1] * bs[j][k_block * 2 + 1]);\n __m512 abscale = _mm512_insertf32x8(_mm512_castps256_ps512(abscale0), abscale1, 1);\n __m512i mul = _mm512_setzero_si512();\n mul = _mm512_dpbssd_epi32(mul, a_val, w[j]);\n acc[j] = _mm512_add_ps(acc[j], _mm512_mul_ps(abscale, _mm512_cvtepi32_ps(mul)));\n }\n }\n store4_reduce_div16(acc[0], acc[1], acc[2], acc[3], c + (n_pos - n_start));\n }\n for (; n_pos < n_end; n_pos++) {\n uint8_t* b_ptr = bb->get_submat(n, k, n_pos, 0);\n float* bs = (float*)bb->get_scale(n, n_pos, k, 0);\n __m512 sum = _mm512_setzero_ps();\n for (int k_block = 0; k_block < k_blocks; k_block++) {\n __m512i a_val = _mm512_loadu_si512((const __m512i*)(a_ptr + k_block * 64));\n __m256i b_val = _mm256_loadu_si256((const __m256i*)(b_ptr + k_block * 32));\n sum = _mm512_add_ps(sum, dot_scaled_kblock(a_val, b_val, as[k_block * 2] * bs[k_block * 2],\n as[k_block * 2 + 1] * bs[k_block * 2 + 1]));\n }\n c[n_pos - n_start] = _mm512_reduce_add_ps(sum) / 16;\n }\n }\n }…om/kvcache-ai/ktransformers into feat/k2-rawint4-prefill-dispatch-8e46e58 # Conflicts: # kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds a prefill-oriented mat-mat path for Kimi-K2 RAWINT4 CPU MoE.
The previous K2 RAWINT4 KGroup kernel used the mat-vec implementation for both decode and prefill. During prefill, when an expert receives multiple tokens, this repeatedly loads and unpacks the same int4 weight blocks for each token.
This change adds a small blocked mat-mat kernel for
GemmKernel224Int4SmallKGroup:qlento a new mat-mat path4 tokens x 4 output columnsper tileDetails
The main change is in:
New helper/path:
dot_scaled_kblockstore4_reduce_div16integer_mat_mat_kgroupThe K2 MoE dispatch already selects mat-mat for larger prefill chunks:
qlen > 4 * expert_num / num_experts_per_tokFor Kimi-K2.6, this means prefill chunks larger than 192 tokens use mat-mat, while decode remains mat-vec.
Why
For prefill, multiple tokens routed to the same expert reuse the same expert weights. The new mat-mat path loads and unpacks each B/int4 block once per small tile, then applies it to multiple token rows. This reduces repeated B-side memory traffic and int4 unpack overhead compared with the old per-token mat-vec loop.
Tests
Added RAWINT4 KGroup coverage for both decode-like and prefill-like shapes:
Local validation:
kt-kernelbuilds and installs successfullyqlen=1