| 1 | |
| 2 | #ifndef MARLIN_NAMESPACE_NAME |
| 3 | #define MARLIN_NAMESPACE_NAME marlin_moe_wna16 |
| 4 | #endif |
| 5 | |
| 6 | #include "libtorch_stable/quantization/marlin/marlin.cuh" |
| 7 | #include "libtorch_stable/quantization/marlin/marlin_dtypes.cuh" |
| 8 | #include "core/scalar_type.hpp" |
| 9 | |
| 10 | #define MARLIN_KERNEL_PARAMS \ |
| 11 | const int4 *__restrict__ A, const int4 *__restrict__ B, \ |
| 12 | int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ |
| 13 | const int4 *__restrict__ b_bias_ptr, \ |
| 14 | const float *__restrict__ a_scales_ptr, \ |
| 15 | const int4 *__restrict__ scales_ptr, \ |
| 16 | const float *__restrict__ global_scale_ptr, \ |
| 17 | const int4 *__restrict__ zp_ptr, \ |
| 18 | const int32_t *__restrict__ sorted_token_ids_ptr, \ |
| 19 | const int32_t *__restrict__ expert_ids_ptr, \ |
| 20 | const int32_t *__restrict__ num_tokens_past_padded_ptr, \ |
| 21 | const float *__restrict__ topk_weights_ptr, int top_k, \ |
| 22 | bool mul_topk_weights, int prob_m, int prob_n, int prob_k, int *locks, \ |
| 23 | bool has_bias, bool use_atomic_add, bool use_fp32_reduce |
| 24 | |
| 25 | namespace MARLIN_NAMESPACE_NAME { |
| 26 | template <const vllm::ScalarTypeId a_type_id, // A ScalarType id |
| 27 | const vllm::ScalarTypeId b_type_id, // B ScalarType id |
| 28 | const vllm::ScalarTypeId c_type_id, // C ScalarType id |
| 29 | const vllm::ScalarTypeId s_type_id, // B_SCALE ScalarType id |
| 30 | const int threads, // number of threads in a threadblock |
| 31 | const int thread_m_blocks, // number of 16x16 blocks in the m |
| 32 | // dimension (batchsize) of the |
| 33 | // threadblock |
| 34 | const int thread_n_blocks, // same for n dimension (output) |
| 35 | const int thread_k_blocks, // same for k dimension (reduction) |
| 36 | const bool m_block_size_8, // whether m_block_size == 8 |
| 37 | // only works when thread_m_blocks == 1 |
| 38 | const int stages, // number of stages for the async global->shared |
| 39 | // fetch pipeline |
| 40 | const int group_blocks, // number of consecutive 16x16 blocks |
| 41 | // with a separate quantization scale |
| 42 | const bool is_zp_float // is zero point of float16 type? |
| 43 | > |
| 44 | __global__ void Marlin(MARLIN_KERNEL_PARAMS); |
| 45 | |
| 46 | } |