| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu |
| 3 | |
| 4 | // clang-format off |
| 5 | |
| 6 | #include "common.h" |
| 7 | #include "gemm.h" |
| 8 | #include "moe.h" |
| 9 | |
| 10 | template <int64_t N> |
| 11 | inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m, int64_t ldn) { |
| 12 | using Vec = at::vec::Vectorized<float>; |
| 13 | constexpr int kVecSize = Vec::size(); |
| 14 | static_assert(N % kVecSize == 0, "copy_bias requires N to be a multiple of Vectorized<float>::size()"); |
| 15 | const bool has_bias = bias_ptr != nullptr; |
| 16 | const Vec zero_vec(0.f); |
| 17 | for (int i = 0; i < m; ++i) { |
| 18 | #pragma GCC unroll 2 |
| 19 | for (int j = 0; j < N; j += kVecSize) { |
| 20 | Vec vec = has_bias ? Vec::loadu(bias_ptr + j) : zero_vec; |
| 21 | vec.store(y_buf + i * ldn + j); |
| 22 | } |
| 23 | } |
| 24 | } |
| 25 | |
| 26 | template <typename scalar_t> |
| 27 | void fused_experts_int4_w4a8_kernel_impl( |
| 28 | scalar_t* __restrict__ output, |
| 29 | scalar_t* __restrict__ ic0, |
| 30 | scalar_t* __restrict__ ic1, |
| 31 | scalar_t* __restrict__ ic2, |
| 32 | uint8_t* __restrict__ A_tmp, |
| 33 | uint8_t* __restrict__ Aq_tmp, |
| 34 | float* __restrict__ As_tmp, |
| 35 | int32_t* __restrict__ Azp_tmp, |
| 36 | float* __restrict__ C_tmp, |
| 37 | int8_t* __restrict__ dqB_tmp, |
| 38 | const scalar_t* __restrict__ input, |
| 39 | const uint8_t* __restrict__ packed_w1, |
| 40 | const uint8_t* __restrict__ packed_w2, |
| 41 | const int8_t* __restrict__ w1z, |
| 42 | const int8_t* __restrict__ w2z, |
| 43 | const float* __restrict__ w1s, |
| 44 | const float* __restrict__ w2s, |
| 45 | int group_size, |
| 46 | const float* __restrict__ topk_weights, |
| 47 | const int32_t* __restrict__ sorted_ids, |
| 48 | const int32_t* __restrict__ expert_ids, |
| 49 | const int32_t* __restrict__ offsets, |
| 50 | int64_t M, |
| 51 | int64_t N, |
| 52 | int64_t K, |
| 53 | int64_t E, |
| 54 | int64_t topk, |
| 55 | int64_t num_tokens_post_pad) { |
| 56 | constexpr int64_t BLOCK_M = block_size_m(); |
| 57 | constexpr int64_t BLOCK_N = block_size_n(); |
| 58 | int num_threads = at::get_num_threads(); |
| 59 | // int64_t buffer_size_nbytes = M * topk * N * 2 |
| 60 | // M * topk * K * 2 + |
| 61 | // num_threads * BLOCK_M * K + |
| 62 | // num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float) + |
| 63 | // M * topk * 2 * N * 2 + |
| 64 | // max(M * K, M * topk * N) + |
| 65 | // M * topk * sizeof(float); |
| 66 | |
| 67 | // intermediate_cache1 (scalar_t): START + M * topk * N |
| 68 | // intermediate_cache2 (scalar_t): + M * topk * K |
| 69 | // A_tmp (uint8_t): + num_threads * BLOCK_M * K |
| 70 | // C_tmp (float): + num_threads * 2 * BLOCK_M * BLOCK_N |
| 71 | // intermediate_cache0 (scalar_t): + M * topk * 2 * N |
| 72 | // Aq_tmp (uint8_t): + max(M * K, M * topk * N) |
| 73 | // As_tmp (float): + M * topk |
| 74 | // dqB_tmp (int8_t) + num_threads * _block_k * BlOCK_N |
| 75 | |
| 76 | // stage 0: quantize input to uint8, [M, K] |
| 77 | at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) { |
| 78 | for (int64_t m = begin; m < end; ++m) { |
| 79 | quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K); |
| 80 | } |
| 81 | }); |
| 82 | int64_t _block_k = get_4bit_block_k_size(group_size); |
| 83 | auto Azp = at::ones({M * topk}).to(at::kInt).mul(128); |
| 84 | auto Azp_ptr = Azp.data_ptr<int32_t>(); |
| 85 | // stage 1: intermediate_cache0 = hidden_states @ w1 |
| 86 | const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M); |
| 87 | const int64_t NB = div_up(N, BLOCK_N); |
| 88 | |
| 89 | int64_t block_per_group = group_size / _block_k; |
| 90 | int64_t Kc = K / _block_k; |
| 91 | int64_t num_groups = K / group_size; |
| 92 | |
| 93 | const int64_t stride_e = 2 * NB * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))); |
| 94 | const bool sym_quant_act = false; |
| 95 | // weight + compensation shape = [E, Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)] |
| 96 | // scales/qzeros shape = [E, Nc, G, block_n] |
| 97 | |
| 98 | // here we only parallel on half of 2N to fuse silu_and_mul with gemm |
| 99 | at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) { |
| 100 | // get local pointers |
| 101 | int tid = at::get_thread_num(); |
| 102 | int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N; |
| 103 | int8_t* dqB_tmp2 = dqB_tmp1 + _block_k * BLOCK_N; |
| 104 | alignas(64) float As[BLOCK_M]; |
| 105 | uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K; |
| 106 | float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N; |
| 107 | float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N; |
| 108 | bool is_brgemm_used = false; |
| 109 | for (int64_t i = begin; i < end; ++i) { |
| 110 | int64_t mb = i / NB; |
| 111 | int64_t nb = i % NB; |
| 112 | int64_t nb1 = nb + NB; |
| 113 | int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N); |
| 114 | // B shape [K, n_size] in vnni format |
| 115 | int32_t expert_id = expert_ids[mb]; |
| 116 | const uint8_t* __restrict__ B = packed_w1 + expert_id * stride_e; |
| 117 | // Bz and Bs: [E, K/gs, 2N] |
| 118 | const int8_t* __restrict__ Bz = w1z + expert_id * (num_groups) * (2 * N); |
| 119 | const float* __restrict__ Bs = w1s + expert_id * (num_groups) * (2 * N); |
| 120 | |
| 121 | // 1.a load A |
| 122 | const int32_t* A_ids = sorted_ids + mb * BLOCK_M; |
| 123 | int64_t m_size = offsets[mb + 1] - offsets[mb]; |
| 124 | const bool use_brgemm = can_use_brgemm<int8_t>(m_size); |
| 125 | is_brgemm_used = is_brgemm_used || use_brgemm; |
| 126 | // copy to A [BLOCK_M, K] |
| 127 | for (int64_t m = 0; m < m_size; ++m) { |
| 128 | int32_t index = A_ids[m] / topk; |
| 129 | copy_stub(A + m * K, Aq_tmp + index * K, K); |
| 130 | As[m] = As_tmp[index]; |
| 131 | } |
| 132 | const int64_t offset = offsets[mb]; |
| 133 | copy_bias<BLOCK_N>(nullptr, C0, m_size, BLOCK_N); |
| 134 | copy_bias<BLOCK_N>(nullptr, C1, m_size, BLOCK_N); |
| 135 | for (int kci = 0; kci < Kc; ++kci) { |
| 136 | int32_t* compensation_ptr = |
| 137 | sym_quant_act ? nullptr |
| 138 | : (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) + |
| 139 | _block_k * BLOCK_N / 2) /*Bcomp*/; |
| 140 | tinygemm_kernel<scalar_t>( |
| 141 | ic0 + offset * 2 * N + nb * BLOCK_N, |
| 142 | C0, |
| 143 | A + kci * _block_k, |
| 144 | As, |
| 145 | Azp_ptr, |
| 146 | B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/, |
| 147 | Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/, |
| 148 | Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/, |
| 149 | compensation_ptr, |
| 150 | dqB_tmp1, |
| 151 | m_size, |
| 152 | _block_k, |
| 153 | K, |
| 154 | BLOCK_N, |
| 155 | 2 * N, |
| 156 | kci == Kc - 1, |
| 157 | use_brgemm); |
| 158 | } |
| 159 | |
| 160 | for (int kci = 0; kci < Kc; ++kci) { |
| 161 | int32_t* compensation_ptr = |
| 162 | sym_quant_act ? nullptr |
| 163 | : (int32_t*)(void*)(B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) + |
| 164 | _block_k * BLOCK_N / 2) /*Bcomp*/; |
| 165 | tinygemm_kernel<scalar_t>( |
| 166 | ic0 + offset * 2 * N + nb1 * BLOCK_N, |
| 167 | C1, |
| 168 | A + kci * _block_k, |
| 169 | As, |
| 170 | Azp_ptr, |
| 171 | B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/, |
| 172 | Bs + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/, |
| 173 | Bz + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/, |
| 174 | compensation_ptr, |
| 175 | dqB_tmp2, |
| 176 | m_size, |
| 177 | _block_k, |
| 178 | K, |
| 179 | BLOCK_N, |
| 180 | 2 * N, |
| 181 | kci == Kc - 1, |
| 182 | use_brgemm); |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | if (is_brgemm_used) { |
| 187 | at::native::cpublas::brgemm_release(); |
| 188 | } |
| 189 | }); |
| 190 | |
| 191 | // stage 1.5: intermediate_cache1 = silu(intermediate_cache0) |
| 192 | at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { |
| 193 | for (int64_t m = begin; m < end; ++m) { |
| 194 | silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N); |
| 195 | } |
| 196 | }); |
| 197 | |
| 198 | // stage 1.5: quantize ic1 to uint8, [M * topk, N] |
| 199 | at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { |
| 200 | for (int64_t m = begin; m < end; ++m) { |
| 201 | quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N); |
| 202 | } |
| 203 | }); |
| 204 | // stage 2: intermediate_cache2 = intermediate_cache1 @ w2 |
| 205 | // w2 : [E, K, N] as [E, OC, IC] |
| 206 | const int64_t OC = K; // rename K as OC |
| 207 | const int64_t IC = N; // rename N as IC |
| 208 | const int64_t MB2 = MB; |
| 209 | const int64_t NB2 = div_up(OC, BLOCK_N); |
| 210 | const int64_t stride_oc = IC; |
| 211 | num_groups = IC / group_size; |
| 212 | Kc = IC / _block_k; |
| 213 | const int64_t stride_e2 = NB2 * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))); |
| 214 | // parallel on [MB2, NB2] |
| 215 | at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) { |
| 216 | int tid = at::get_thread_num(); |
| 217 | int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N; |
| 218 | float* __restrict__ C2 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N; |
| 219 | bool is_brgemm_used = false; |
| 220 | for (int64_t i = begin; i < end; ++i) { |
| 221 | int64_t mb = i / NB2; |
| 222 | int64_t nb = i % NB2; |
| 223 | |
| 224 | int64_t m_size = offsets[mb + 1] - offsets[mb]; |
| 225 | int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N); |
| 226 | const bool use_brgemm = can_use_brgemm<int8_t>(m_size); |
| 227 | is_brgemm_used = is_brgemm_used || use_brgemm; |
| 228 | const int32_t* A_ids = sorted_ids + mb * BLOCK_M; |
| 229 | |
| 230 | // B shape [IC, n_size] in vnni format |
| 231 | int32_t expert_id = expert_ids[mb]; |
| 232 | const uint8_t* __restrict__ B = packed_w2 + expert_id * stride_e2; |
| 233 | |
| 234 | // Bz and Bs: [E, IC/gs, OC] |
| 235 | const int8_t* __restrict__ Bz = w2z + expert_id * (num_groups)*OC; |
| 236 | const float* __restrict__ Bs = w2s + expert_id * (num_groups)*OC; |
| 237 | |
| 238 | // A ptr from ic1 of [M * topk, N] in sorted order |
| 239 | // so as to avoid copy A to tmp buffer again |
| 240 | const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * IC; |
| 241 | const float* __restrict__ As = As_tmp + offsets[mb]; |
| 242 | copy_bias<BLOCK_N>(nullptr, C2, m_size, BLOCK_N); |
| 243 | for (int kci = 0; kci < Kc; ++kci) { |
| 244 | int32_t* compensation_ptr = |
| 245 | sym_quant_act ? nullptr |
| 246 | : (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) + |
| 247 | _block_k * BLOCK_N / 2) /*Bcomp*/; |
| 248 | tinygemm_kernel<scalar_t>( |
| 249 | nullptr, /*store_out is false*/ |
| 250 | C2, |
| 251 | A + kci * _block_k, |
| 252 | As, |
| 253 | Azp_ptr, |
| 254 | B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))), |
| 255 | Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/, |
| 256 | Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*zeros_b*/, |
| 257 | compensation_ptr, |
| 258 | dqB_tmp1, |
| 259 | m_size, |
| 260 | _block_k, |
| 261 | IC, |
| 262 | BLOCK_N, |
| 263 | BLOCK_N, |
| 264 | false, |
| 265 | use_brgemm); |
| 266 | } |
| 267 | |
| 268 | // 2.b copy from C to ic2 in original order |
| 269 | // and also mul topk_weights in float32 |
| 270 | for (int64_t m = 0; m < m_size; ++m) { |
| 271 | int32_t index = A_ids[m]; |
| 272 | float weight = topk_weights[index]; |
| 273 | copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C2 + m * BLOCK_N, weight, n_size); |
| 274 | } |
| 275 | } |
| 276 | |
| 277 | if (is_brgemm_used) { |
| 278 | at::native::cpublas::brgemm_release(); |
| 279 | } |
| 280 | }); |
| 281 | |
| 282 | // stage 3: out = intermediate_cache2.sum(dim=1) |
| 283 | // from [M, topk, K] to [M, K] |
| 284 | at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) { |
| 285 | for (int64_t m = begin; m < end; ++m) { |
| 286 | sum_stub(output + m * K, ic2 + m * topk * K, topk, K); |
| 287 | } |
| 288 | }); |
| 289 | } |
| 290 | |
| 291 | #define INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(TYPE) \ |
| 292 | template void fused_experts_int4_w4a8_kernel_impl<TYPE>( \ |
| 293 | TYPE* __restrict__ output, \ |
| 294 | TYPE* __restrict__ ic0, \ |
| 295 | TYPE* __restrict__ ic1, \ |
| 296 | TYPE* __restrict__ ic2, \ |
| 297 | uint8_t* __restrict__ A_tmp, \ |
| 298 | uint8_t* __restrict__ Aq_tmp, \ |
| 299 | float* __restrict__ As_tmp, \ |
| 300 | int32_t* __restrict__ Azp_tmp, \ |
| 301 | float* __restrict__ C_tmp, \ |
| 302 | int8_t* __restrict__ dqB_tmp, \ |
| 303 | const TYPE* __restrict__ input, \ |
| 304 | const uint8_t* __restrict__ packed_w1, \ |
| 305 | const uint8_t* __restrict__ packed_w2, \ |
| 306 | const int8_t* __restrict__ w1z, \ |
| 307 | const int8_t* __restrict__ w2z, \ |
| 308 | const float* __restrict__ w1s, \ |
| 309 | const float* __restrict__ w2s, \ |
| 310 | int group_size, \ |
| 311 | const float* __restrict__ topk_weights, \ |
| 312 | const int32_t* __restrict__ sorted_ids, \ |
| 313 | const int32_t* __restrict__ expert_ids, \ |
| 314 | const int32_t* __restrict__ offsets, \ |
| 315 | int64_t M, \ |
| 316 | int64_t N, \ |
| 317 | int64_t K, \ |
| 318 | int64_t E, \ |
| 319 | int64_t topk, \ |
| 320 | int64_t num_tokens_post_pad) |
| 321 | |
| 322 | INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::BFloat16); |
| 323 | INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::Half); |