| 1 | |
| 2 | #ifndef MARLIN_NAMESPACE_NAME |
| 3 | #define MARLIN_NAMESPACE_NAME marlin |
| 4 | #endif |
| 5 | |
| 6 | #include "marlin.cuh" |
| 7 | #include "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, int prob_m, int prob_n, int prob_k, \ |
| 18 | int lda, int *locks, bool has_bias, bool use_atomic_add, \ |
| 19 | bool use_fp32_reduce, int max_shared_mem |
| 20 | |
| 21 | namespace MARLIN_NAMESPACE_NAME { |
| 22 | template <const vllm::ScalarTypeId a_type_id, // A ScalarType id |
| 23 | const vllm::ScalarTypeId b_type_id, // B ScalarType id |
| 24 | const vllm::ScalarTypeId c_type_id, // C ScalarType id |
| 25 | const vllm::ScalarTypeId s_type_id, // B_SCALE ScalarType id |
| 26 | const int threads, // number of threads in a threadblock |
| 27 | const int thread_m_blocks, // number of 16x16 blocks in the m |
| 28 | // dimension (batchsize) of the |
| 29 | // threadblock |
| 30 | const int thread_n_blocks, // same for n dimension (output) |
| 31 | const int thread_k_blocks, // same for k dimension (reduction) |
| 32 | const bool m_block_size_8, // whether m_block_size == 8 |
| 33 | // only works when thread_m_blocks == 1 |
| 34 | const int stages, // number of stages for the async global->shared |
| 35 | // fetch pipeline |
| 36 | const int group_blocks, // number of consecutive 16x16 blocks |
| 37 | // with a separate quantization scale |
| 38 | const bool is_zp_float // is zero point of float16 type? |
| 39 | > |
| 40 | __global__ void Marlin(MARLIN_KERNEL_PARAMS); |
| 41 | |
| 42 | } |