| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu |
| 3 | |
| 4 | // clang-format off |
| 5 | |
| 6 | #include <c10/util/Unroll.h> |
| 7 | |
| 8 | #include "common.h" |
| 9 | #include "gemm.h" |
| 10 | #include "moe.h" |
| 11 | |
| 12 | namespace { |
| 13 | |
| 14 | // Byte-wise copy for FP8 elements (no remainder assumption) |
| 15 | inline void copy_stub2( |
| 16 | at::Float8_e4m3fn* __restrict__ out, |
| 17 | const at::Float8_e4m3fn* __restrict__ input, |
| 18 | int64_t size) { |
| 19 | for (int64_t d = 0; d < size; d++) { |
| 20 | out[d] = input[d]; |
| 21 | } |
| 22 | } |
| 23 | |
| 24 | #define PER_TENSOR 1 |
| 25 | #define PER_ROW 2 |
| 26 | #define PER_GROUP 3 |
| 27 | |
| 28 | // Store float32 GEMM result to output buffer with dequantization and dtype |
| 29 | // conversion. Activations are PER_ROW quantized; weights are PER_GROUP quantized. |
| 30 | // Uses AVX512 on capable platforms. |
| 31 | #if defined(CPU_CAPABILITY_AVX512) |
| 32 | |
| 33 | template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode> |
| 34 | inline void store_out( |
| 35 | const float* y_buf, |
| 36 | out_dtype* c_ptr, |
| 37 | int64_t M, |
| 38 | int64_t lda, |
| 39 | const float* scales_a, |
| 40 | const float* scales_b, |
| 41 | const float* bias) { |
| 42 | float a_scale = 1.0, b_scale = 1.0; |
| 43 | __m512 va_scale, vb_scale; |
| 44 | if constexpr (act_quant_mode == PER_TENSOR) { |
| 45 | a_scale = *scales_a; |
| 46 | } |
| 47 | if constexpr (wei_quant_mode == PER_TENSOR) { |
| 48 | b_scale = *scales_b; |
| 49 | vb_scale = _mm512_set1_ps(b_scale); |
| 50 | } |
| 51 | for (int i = 0; i < M; ++i) { |
| 52 | if constexpr (act_quant_mode == PER_ROW) { |
| 53 | a_scale = *(scales_a + i); |
| 54 | } |
| 55 | if constexpr (act_quant_mode != PER_GROUP) { |
| 56 | va_scale = _mm512_set1_ps(a_scale); |
| 57 | } |
| 58 | constexpr int N_UNROLL = N / 16; |
| 59 | c10::ForcedUnroll<N_UNROLL>{}([&](auto idx) { |
| 60 | constexpr int j = idx * 16; |
| 61 | __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j); |
| 62 | __m512 bias_vec = bias ? _mm512_loadu_ps(bias + j) : _mm512_setzero_ps(); |
| 63 | if constexpr (act_quant_mode != PER_GROUP) { |
| 64 | y_vec = _mm512_mul_ps(y_vec, va_scale); |
| 65 | } |
| 66 | if constexpr (wei_quant_mode == PER_ROW) { |
| 67 | vb_scale = _mm512_loadu_ps(scales_b + j); |
| 68 | } |
| 69 | if constexpr (wei_quant_mode != PER_GROUP) { |
| 70 | y_vec = _mm512_mul_ps(y_vec, vb_scale); |
| 71 | } |
| 72 | y_vec = _mm512_add_ps(y_vec, bias_vec); |
| 73 | if constexpr (std::is_same<out_dtype, float>::value) { |
| 74 | _mm512_storeu_ps(c_ptr + i * lda + j, y_vec); |
| 75 | } else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) { |
| 76 | __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec); |
| 77 | _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec); |
| 78 | } else if constexpr (std::is_same<out_dtype, at::Half>::value) { |
| 79 | __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec); |
| 80 | _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec); |
| 81 | } else { |
| 82 | TORCH_CHECK(false, "Unsupported output dtype"); |
| 83 | } |
| 84 | }); |
| 85 | constexpr int tail_start = N / 16 * 16; |
| 86 | for (int j = tail_start; j < N; ++j) { |
| 87 | if constexpr (wei_quant_mode == PER_ROW) { |
| 88 | b_scale = scales_b[j]; |
| 89 | } |
| 90 | c_ptr[i * lda + j] = static_cast<out_dtype>(y_buf[i * N + j] * a_scale * b_scale); |
| 91 | } |
| 92 | } // for M |
| 93 | } |
| 94 | |
| 95 | #else // no AVX512 |
| 96 | |
| 97 | template <typename out_dtype, int64_t N, int act_quant_mode, int wei_quant_mode> |
| 98 | inline void store_out( |
| 99 | const float* y_buf, |
| 100 | out_dtype* c_ptr, |
| 101 | int64_t M, |
| 102 | int64_t lda, |
| 103 | const float* scales_a, |
| 104 | const float* scales_b, |
| 105 | const float* bias) { |
| 106 | float a_scale = 1.0, b_scale = 1.0; |
| 107 | if constexpr (act_quant_mode == PER_TENSOR) { |
| 108 | a_scale = *scales_a; |
| 109 | } |
| 110 | if constexpr (wei_quant_mode == PER_TENSOR) { |
| 111 | b_scale = *scales_b; |
| 112 | } |
| 113 | for (int i = 0; i < M; ++i) { |
| 114 | if constexpr (act_quant_mode == PER_ROW) { |
| 115 | a_scale = *(scales_a + i); |
| 116 | } |
| 117 | for (int j = 0; j < N; ++j) { |
| 118 | if constexpr (wei_quant_mode == PER_ROW) { |
| 119 | b_scale = scales_b[j]; |
| 120 | } |
| 121 | float val = y_buf[i * N + j] * a_scale * b_scale; |
| 122 | if (bias) val += bias[j]; |
| 123 | c_ptr[i * lda + j] = static_cast<out_dtype>(val); |
| 124 | } |
| 125 | } |
| 126 | } |
| 127 | |
| 128 | #endif // CPU_CAPABILITY_AVX512 |
| 129 | |
| 130 | } // anonymous namespace |
| 131 | |
| 132 | // FP8 W8A8 fused MoE kernel. |
| 133 | // |
| 134 | // hidden_states: [M, K] FP8 E4M3, pre-quantized activation |
| 135 | // packed_w1: [E, 2*Nc, Kc, BLOCK_K, BLOCK_N] FP8 W8A8 prepack (float8_linear_prepack_cpu) |
| 136 | // packed_w2: [E, K, N_packed] FP8 W8A16 VNNI pack (convert_weight_packed) |
| 137 | // As: [M] float32 per-token activation scales |
| 138 | // w1s: [E, num_groups, 2N] float32 per-group weight scales for w1 |
| 139 | // w2s: [E, ...] float32 per-group weight scales for w2 |
| 140 | // |
| 141 | // Stage 1: hidden_states(FP8) @ w1(FP8) → ic0(BF16) [M*topk, 2N] |
| 142 | // Uses FP8×FP8 GEMM (tinygemm_kernel with FP8 inputs) |
| 143 | // Stage 1.5: ic1 = silu(ic0[:,:N]) * ic0[:,N:] [M*topk, N] |
| 144 | // Stage 2: ic1(BF16) @ w2(FP8) → ic2(BF16) [M*topk, K] |
| 145 | // Uses BF16×FP8 GEMM (W8A16 path, tinygemm_kernel<scalar_t>) |
| 146 | // Stage 3: out = weighted_sum(ic2, topk_weights) [M, K] |
| 147 | |
| 148 | template <typename scalar_t> |
| 149 | void fused_experts_fp8_a8_kernel_impl( |
| 150 | scalar_t* __restrict__ output, |
| 151 | scalar_t* __restrict__ ic0, |
| 152 | scalar_t* __restrict__ ic1, |
| 153 | scalar_t* __restrict__ ic2, |
| 154 | at::Float8_e4m3fn* __restrict__ A_tmp, |
| 155 | scalar_t* __restrict__ B_tmp, |
| 156 | float* __restrict__ C_tmp, |
| 157 | float* __restrict__ Ukernel_tmp, |
| 158 | const at::Float8_e4m3fn* __restrict__ input, |
| 159 | const at::Float8_e4m3fn* __restrict__ packed_w1, |
| 160 | const at::Float8_e4m3fn* __restrict__ packed_w2, |
| 161 | const float* __restrict__ As, |
| 162 | const float* __restrict__ w1s, |
| 163 | const float* __restrict__ w2s, |
| 164 | int64_t block_size_N, |
| 165 | int64_t block_size_K, |
| 166 | const float* __restrict__ topk_weights, |
| 167 | const int32_t* __restrict__ sorted_ids, |
| 168 | const int32_t* __restrict__ expert_ids, |
| 169 | const int32_t* __restrict__ offsets, |
| 170 | int64_t M, |
| 171 | int64_t N, |
| 172 | int64_t K, |
| 173 | int64_t E, |
| 174 | int64_t topk, |
| 175 | int64_t num_tokens_post_pad) { |
| 176 | // 1. intermediate_cache1 : [M * topk, N] |
| 177 | // 2. intermediate_cache2 : [M * topk, K] |
| 178 | // 3. A_tmp : [T, BLOCK_M * K] FP8 |
| 179 | // 4. C_tmp : [T, 2 * BLOCK_M * BLOCK_N] float32 |
| 180 | // 5. intermediate_cache0 : [M * topk, 2N] |
| 181 | // 6. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, max(K, N)] scalar_t |
| 182 | // 7. Ukernel_tmp: [T, 2 * BLOCK_M * BLOCK_N] float32 |
| 183 | // 8. (non-FP8-brgemm only) dqA [T, BLOCK_M * BLOCK_K] + dqB [T, BLOCK_K * BLOCK_N] in BF16 |
| 184 | // appended after Ukernel_tmp; see moe.cpp FP8_W8A8 buffer allocation |
| 185 | |
| 186 | constexpr int64_t BLOCK_M = block_size_m(); |
| 187 | constexpr int64_t BLOCK_N_VAL = block_size_n(); // = 32 |
| 188 | const int64_t KB = K / BLOCK_K; // K blocks |
| 189 | |
| 190 | int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N_VAL * std::max(K, N); |
| 191 | |
| 192 | // stage 1: intermediate_cache0 = hidden_states @ w1 |
| 193 | const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M); |
| 194 | const int64_t NB = div_up(N, BLOCK_N_VAL); |
| 195 | |
| 196 | int64_t scale_size_N = div_up(2 * N, block_size_N); |
| 197 | int64_t scale_size_K = div_up(K, block_size_K); |
| 198 | int64_t blocks_n_per_group = block_size_N / BLOCK_N_VAL; |
| 199 | int64_t num_groups = div_up(K, block_size_K); // G |
| 200 | int64_t blocks_k_per_group = block_size_K / BLOCK_K; |
| 201 | const int64_t stride_e = 2 * N * K; // per-expert stride for w1 (total elements) |
| 202 | bool use_brgemm = true; |
| 203 | int64_t num_thread = at::get_num_threads(); |
| 204 | |
| 205 | // For the BF16 fallback path (no native FP8 brgemm), tinygemm_kernel needs temporary |
| 206 | // BF16 buffers to dequantize FP8 inputs before running AMX-BF16 brgemm. |
| 207 | // These are placed after the Ukernel_tmp region (see moe.cpp FP8_W8A8 buffer allocation). |
| 208 | #ifndef CPUBLAS_BRGEMM_F8F8F32 |
| 209 | // dqA: [T, BLOCK_M, BLOCK_K] BF16 |
| 210 | // dqB: [T, BLOCK_K, BLOCK_N] BF16 |
| 211 | // Layout: immediately after Ukernel_tmp[num_thread * 2 * BLOCK_M * BLOCK_N_VAL] |
| 212 | at::BFloat16* __restrict__ dq_base = reinterpret_cast<at::BFloat16*>( |
| 213 | Ukernel_tmp + num_thread * 2 * BLOCK_M * BLOCK_N_VAL); |
| 214 | const int64_t dq_per_thread = BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N_VAL; |
| 215 | #endif |
| 216 | |
| 217 | at::parallel_for(0, MB * NB, 1, [&](int64_t begin, int64_t end) { |
| 218 | int tid = get_thread_num(); |
| 219 | // A_tmp: FP8 activation buffer for this thread [BLOCK_M, K] |
| 220 | at::Float8_e4m3fn* __restrict__ A = A_tmp + tid * BLOCK_M * K; |
| 221 | float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL; |
| 222 | float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N_VAL; |
| 223 | alignas(64) float As_[BLOCK_M]; |
| 224 | float* __restrict__ ukernel_buf_1 = Ukernel_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL; |
| 225 | #ifndef CPUBLAS_BRGEMM_F8F8F32 |
| 226 | at::BFloat16* __restrict__ dqA_buf = dq_base + tid * dq_per_thread; |
| 227 | at::BFloat16* __restrict__ dqB_buf = dqA_buf + BLOCK_M * BLOCK_K; |
| 228 | #else |
| 229 | at::BFloat16* __restrict__ dqA_buf = nullptr; |
| 230 | at::BFloat16* __restrict__ dqB_buf = nullptr; |
| 231 | #endif |
| 232 | float* __restrict__ ukernel_buf_2 = ukernel_buf_1 + BLOCK_M * BLOCK_N_VAL; |
| 233 | auto ldsa = 1; // activation per-row quantized |
| 234 | |
| 235 | for (int64_t i = begin; i < end; ++i) { |
| 236 | int64_t mb = i / NB; |
| 237 | int64_t nb = i % NB; |
| 238 | int64_t nb1 = nb + NB; // gate is first N columns, up is second N columns |
| 239 | |
| 240 | int64_t n_size = std::min(N - nb * BLOCK_N_VAL, BLOCK_N_VAL); |
| 241 | |
| 242 | // B (weight) for expert expert_id: gate+up fused [2N, K] prepackaged as FP8 |
| 243 | int32_t expert_id = expert_ids[mb]; |
| 244 | const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e; |
| 245 | const float* __restrict__ Bs = w1s + expert_id * num_groups * (2 * N); |
| 246 | |
| 247 | // Load activation tile for this M-block |
| 248 | const int32_t* A_ids = sorted_ids + mb * BLOCK_M; |
| 249 | int64_t m_size = offsets[mb + 1] - offsets[mb]; |
| 250 | |
| 251 | for (int64_t m = 0; m < m_size; ++m) { |
| 252 | int32_t index = A_ids[m] / topk; |
| 253 | copy_stub2(A + m * K, input + index * K, K); |
| 254 | As_[m] = As[index]; |
| 255 | } |
| 256 | |
| 257 | const int64_t offset = offsets[mb]; |
| 258 | |
| 259 | // Gate path (nb): w1[:N, :] |
| 260 | zero_buffer(C0, BLOCK_M * BLOCK_N_VAL); |
| 261 | for (int kci = 0; kci < KB; ++kci) { |
| 262 | // FP8×FP8 tinygemm: act_quant=PER_ROW(2), wei_quant=PER_GROUP(3) |
| 263 | tinygemm_kernel( |
| 264 | /* C */ C0, |
| 265 | /* A */ A + kci * BLOCK_K, |
| 266 | /* scales_a*/ As_ + kci / blocks_k_per_group, |
| 267 | /* B */ B + (nb * KB + kci) * BLOCK_K * BLOCK_N_VAL, |
| 268 | /* scales_b*/ Bs + nb * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL, |
| 269 | /* M */ m_size, |
| 270 | /* K */ BLOCK_K, |
| 271 | /* lda */ K, |
| 272 | /* ldc */ BLOCK_N_VAL, |
| 273 | /* ldsa */ ldsa, |
| 274 | /* ukernel */ ukernel_buf_1, |
| 275 | /* dqA_buf */ dqA_buf, |
| 276 | /* dqB_buf */ dqB_buf); |
| 277 | } |
| 278 | store_out<scalar_t, BLOCK_N_VAL, PER_ROW, PER_GROUP>( |
| 279 | C0, |
| 280 | ic0 + offset * 2 * N + nb * BLOCK_N_VAL, |
| 281 | m_size, |
| 282 | 2 * N, // lda |
| 283 | As_, |
| 284 | nullptr, |
| 285 | nullptr); |
| 286 | |
| 287 | // Up path (nb1): w1[N:, :] |
| 288 | zero_buffer(C1, BLOCK_M * BLOCK_N_VAL); |
| 289 | for (int kci = 0; kci < KB; ++kci) { |
| 290 | tinygemm_kernel( |
| 291 | /* C */ C1, |
| 292 | /* A */ A + kci * BLOCK_K, |
| 293 | /* scales_a*/ As_ + kci / blocks_k_per_group, |
| 294 | /* B */ B + (nb1 * KB + kci) * BLOCK_K * BLOCK_N_VAL, |
| 295 | /* scales_b*/ Bs + nb1 * BLOCK_N_VAL * num_groups + kci / blocks_k_per_group * BLOCK_N_VAL, |
| 296 | /* M */ m_size, |
| 297 | /* K */ BLOCK_K, |
| 298 | /* lda */ K, |
| 299 | /* ldc */ BLOCK_N_VAL, |
| 300 | /* ldsa */ ldsa, |
| 301 | /* ukernel */ ukernel_buf_2, |
| 302 | /* dqA_buf */ dqA_buf, |
| 303 | /* dqB_buf */ dqB_buf); |
| 304 | } |
| 305 | store_out<scalar_t, BLOCK_N_VAL, PER_ROW, PER_GROUP>( |
| 306 | C1, |
| 307 | ic0 + offset * 2 * N + nb1 * BLOCK_N_VAL, |
| 308 | m_size, |
| 309 | 2 * N, // lda |
| 310 | As_, |
| 311 | nullptr, |
| 312 | nullptr); |
| 313 | } |
| 314 | if (use_brgemm) { |
| 315 | at::native::cpublas::brgemm_release(); |
| 316 | } |
| 317 | }); |
| 318 | |
| 319 | // stage 1.5: intermediate_cache1 = silu(ic0[:, :N]) * ic0[:, N:] |
| 320 | at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { |
| 321 | for (int64_t m = begin; m < end; ++m) { |
| 322 | silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N); |
| 323 | } |
| 324 | }); |
| 325 | |
| 326 | // stage 2: intermediate_cache2 = intermediate_cache1 @ w2 |
| 327 | // w2: [E, K, N] FP8, VNNI packed; treated as W8A16 (BF16 act × FP8 weight) |
| 328 | const int64_t OC = K; // output channels (hidden dim) |
| 329 | const int64_t IC = N; // input channels (expert intermediate dim) |
| 330 | const int64_t MB2 = MB; |
| 331 | const int64_t NB2 = div_up(OC, BLOCK_N_VAL); |
| 332 | scale_size_N = div_up(K, block_size_N); |
| 333 | scale_size_K = div_up(N, block_size_K); |
| 334 | const int64_t stride_e2 = OC * IC; // per-expert stride for w2 |
| 335 | const int64_t stride_oc = IC; |
| 336 | |
| 337 | int64_t avg_M = std::max(int64_t(1), M * topk / E); |
| 338 | bool use_brgemm2 = can_use_brgemm<at::Float8_e4m3fn>(avg_M); |
| 339 | |
| 340 | parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { |
| 341 | int tid = get_thread_num(); |
| 342 | alignas(64) scalar_t C[BLOCK_M * BLOCK_K]; |
| 343 | |
| 344 | loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N_VAL * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) { |
| 345 | int64_t m_size = offsets[mb + 1] - offsets[mb]; |
| 346 | int64_t n_size = std::min(OC - nb * BLOCK_N_VAL, BLOCK_N_VAL); |
| 347 | |
| 348 | // A from ic1: [M*topk, N] in sorted order (contiguous after stage 1.5) |
| 349 | const scalar_t* __restrict__ A2 = ic1 + offsets[mb] * N; |
| 350 | const int32_t* A_ids = sorted_ids + mb * BLOCK_M; |
| 351 | |
| 352 | // B from w2: [E, K, N] VNNI packed, offset to expert + output block |
| 353 | int32_t expert_id = expert_ids[mb]; |
| 354 | const at::Float8_e4m3fn* __restrict__ B2 = |
| 355 | packed_w2 + expert_id * stride_e2 + nb * BLOCK_N_VAL * stride_oc; |
| 356 | const float* __restrict__ Bs2 = |
| 357 | w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K; |
| 358 | |
| 359 | // Only unpack B on first visit or new expert |
| 360 | int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1]; |
| 361 | bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id); |
| 362 | |
| 363 | tinygemm_kernel<scalar_t>( |
| 364 | /* A */ A2, |
| 365 | /* B */ B2, |
| 366 | /* C */ C, |
| 367 | /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N_VAL * IC, |
| 368 | /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N_VAL, |
| 369 | /* Bbias */ nullptr, |
| 370 | /* scale */ Bs2, |
| 371 | /* M */ m_size, |
| 372 | /* N */ n_size, |
| 373 | /* K */ IC, |
| 374 | /* lda */ IC, |
| 375 | /* ldb */ n_size, |
| 376 | /* ldc */ BLOCK_N_VAL, |
| 377 | /* brg */ use_brgemm2, |
| 378 | /* block_sz_K */ block_size_K, |
| 379 | /* do_unpack */ do_unpack); |
| 380 | |
| 381 | // Scatter ic2 in original token order, weighting by topk_weights |
| 382 | for (int64_t m = 0; m < m_size; ++m) { |
| 383 | int32_t index = A_ids[m]; |
| 384 | float weight = topk_weights[index]; |
| 385 | copy_mul_stub(ic2 + index * K + nb * BLOCK_N_VAL, C + m * BLOCK_N_VAL, weight, n_size); |
| 386 | } |
| 387 | }); |
| 388 | |
| 389 | if (use_brgemm2) { |
| 390 | at::native::cpublas::brgemm_release(); |
| 391 | } |
| 392 | }); |
| 393 | |
| 394 | // stage 3: output = sum(ic2, dim=topk) [M, K] |
| 395 | at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) { |
| 396 | for (int64_t m = begin; m < end; ++m) { |
| 397 | sum_stub(output + m * K, ic2 + m * topk * K, topk, K); |
| 398 | } |
| 399 | }); |
| 400 | } |
| 401 | |
| 402 | #define INSTANTIATE_MOE_FP8_A8_TEMPLATE(TYPE) \ |
| 403 | template void fused_experts_fp8_a8_kernel_impl<TYPE>( \ |
| 404 | TYPE* __restrict__ output, \ |
| 405 | TYPE* __restrict__ ic0, \ |
| 406 | TYPE* __restrict__ ic1, \ |
| 407 | TYPE* __restrict__ ic2, \ |
| 408 | at::Float8_e4m3fn* __restrict__ A_tmp, \ |
| 409 | TYPE* __restrict__ B_tmp, \ |
| 410 | float* __restrict__ C_tmp, \ |
| 411 | float* __restrict__ Ukernel_tmp, \ |
| 412 | const at::Float8_e4m3fn* __restrict__ input, \ |
| 413 | const at::Float8_e4m3fn* __restrict__ packed_w1, \ |
| 414 | const at::Float8_e4m3fn* __restrict__ packed_w2, \ |
| 415 | const float* __restrict__ As, \ |
| 416 | const float* __restrict__ w1s, \ |
| 417 | const float* __restrict__ w2s, \ |
| 418 | int64_t block_size_N, \ |
| 419 | int64_t block_size_K, \ |
| 420 | const float* __restrict__ topk_weights, \ |
| 421 | const int32_t* __restrict__ sorted_ids, \ |
| 422 | const int32_t* __restrict__ expert_ids, \ |
| 423 | const int32_t* __restrict__ offsets, \ |
| 424 | int64_t M, \ |
| 425 | int64_t N, \ |
| 426 | int64_t K, \ |
| 427 | int64_t E, \ |
| 428 | int64_t topk, \ |
| 429 | int64_t num_tokens_post_pad) |
| 430 | |
| 431 | INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::BFloat16); |
| 432 | INSTANTIATE_MOE_FP8_A8_TEMPLATE(at::Half); |