From 9dff9331648a25017733244bb17a78eea6d4645d Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Wed, 25 Feb 2026 10:08:22 +0800 Subject: [PATCH] [Kernel Slimming] Remove sgl-kernel AOT marlin kernels (#19241) --- .../jit_kernel/tests/test_gptq_marlin.py | 24 - .../tests/test_gptq_marlin_repack.py | 9 - sgl-kernel/CMakeLists.txt | 3 - sgl-kernel/csrc/common_extension.cc | 14 - .../csrc/gemm/marlin/awq_marlin_repack.cu | 253 ---- sgl-kernel/csrc/gemm/marlin/gptq_marlin.cu | 1120 ----------------- .../csrc/gemm/marlin/gptq_marlin_repack.cu | 329 ----- sgl-kernel/include/sgl_kernel_ops.h | 24 - sgl-kernel/python/sgl_kernel/__init__.py | 6 - sgl-kernel/python/sgl_kernel/gemm.py | 41 - sgl-kernel/python/sgl_kernel/marlin.py | 44 - sgl-kernel/tests/test_marlin_gemm.py | 121 -- sgl-kernel/tests/test_marlin_repack.py | 148 --- 13 files changed, 2136 deletions(-) delete mode 100644 sgl-kernel/csrc/gemm/marlin/awq_marlin_repack.cu delete mode 100644 sgl-kernel/csrc/gemm/marlin/gptq_marlin.cu delete mode 100644 sgl-kernel/csrc/gemm/marlin/gptq_marlin_repack.cu delete mode 100644 sgl-kernel/python/sgl_kernel/marlin.py delete mode 100644 sgl-kernel/tests/test_marlin_gemm.py delete mode 100644 sgl-kernel/tests/test_marlin_repack.py diff --git a/python/sglang/jit_kernel/tests/test_gptq_marlin.py b/python/sglang/jit_kernel/tests/test_gptq_marlin.py index 4b4a37f4d..b023b5956 100644 --- a/python/sglang/jit_kernel/tests/test_gptq_marlin.py +++ b/python/sglang/jit_kernel/tests/test_gptq_marlin.py @@ -1,6 +1,5 @@ import pytest import torch -from sgl_kernel import gptq_marlin_gemm as aot_gptq_marlin_gemm from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm @@ -86,26 +85,6 @@ def test_gptq_marlin_gemm( is_zp_float=False, ) - aot_output = aot_gptq_marlin_gemm( - a_input, - None, - marlin_q_w, - marlin_s, - marlin_s2, - marlin_zp, - g_idx, - sort_indices, - workspace, - quant_type, - a_input.shape[0], - b_weight.shape[1], - a_input.shape[1], - is_k_full=True, - use_atomic_add=False, - use_fp32_reduce=False, - is_zp_float=False, - ) - output_ref = torch.matmul(a_input, w_ref) torch.cuda.synchronize() @@ -115,9 +94,6 @@ def test_gptq_marlin_gemm( ) assert max_diff < 0.04 - # JIT kernel should produce bitwise identical results to AOT kernel - torch.testing.assert_close(output, aot_output, rtol=0, atol=0) - if __name__ == "__main__": import subprocess diff --git a/python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py b/python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py index dadfd0ecf..ef02be3e8 100644 --- a/python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py +++ b/python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py @@ -1,6 +1,5 @@ import pytest import torch -from sgl_kernel import gptq_marlin_repack as aot_gptq_marlin_repack from sgl_kernel.scalar_type import scalar_types from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack @@ -81,19 +80,11 @@ def test_gptq_marlin_repack( q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits ) - # Run AOT repack kernel - aot_output = aot_gptq_marlin_repack( - q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits - ) - torch.cuda.synchronize() # JIT should match the reference (computed from CPU marlin_weights) torch.testing.assert_close(jit_output, q_w_marlin_ref) - # JIT should produce bitwise identical results to AOT - torch.testing.assert_close(jit_output, aot_output, rtol=0, atol=0) - if __name__ == "__main__": import subprocess diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index f5ad30872..680f41438 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -297,9 +297,6 @@ set(SOURCES "csrc/gemm/per_token_quant_fp8.cu" "csrc/gemm/qserve_w4a8_per_chn_gemm.cu" "csrc/gemm/qserve_w4a8_per_group_gemm.cu" - "csrc/gemm/marlin/gptq_marlin.cu" - "csrc/gemm/marlin/gptq_marlin_repack.cu" - "csrc/gemm/marlin/awq_marlin_repack.cu" "csrc/gemm/gptq/gptq_kernel.cu" "csrc/grammar/apply_token_bitmask_inplace_cuda.cu" diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 004b85304..087383f1d 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -196,14 +196,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { /* * From csrc/gemm/gptq */ - m.def( - "gptq_marlin_gemm(Tensor! a, Tensor? c_or_none," - "Tensor! b_q_weight, Tensor! b_scales, Tensor? global_scale_or_none," - "Tensor? b_zeros_or_none, Tensor? g_idx_or_none, Tensor? perm_or_none," - "Tensor! workspace, int b_q_type_id, int size_m, int size_n, int size_k," - "bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor"); - m.impl("gptq_marlin_gemm", torch::kCUDA, &gptq_marlin_gemm); - m.def( "gptq_gemm(Tensor a, Tensor b_q_weight, Tensor b_gptq_qzeros, Tensor b_gptq_scales, Tensor b_g_idx, bool " "use_shuffle, int bit) -> Tensor"); @@ -212,12 +204,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("gptq_shuffle(Tensor! q_weight, Tensor q_perm, int bit) -> ()"); m.impl("gptq_shuffle", torch::kCUDA, &gptq_shuffle); - m.def("gptq_marlin_repack(Tensor! b_q_weight, Tensor! perm, int size_k, int size_n, int num_bits) -> Tensor"); - m.impl("gptq_marlin_repack", torch::kCUDA, &gptq_marlin_repack); - - m.def("awq_marlin_repack(Tensor! b_q_weight, int size_k, int size_n, int num_bits) -> Tensor"); - m.impl("awq_marlin_repack", torch::kCUDA, &awq_marlin_repack); - /* * From csrc/moe */ diff --git a/sgl-kernel/csrc/gemm/marlin/awq_marlin_repack.cu b/sgl-kernel/csrc/gemm/marlin/awq_marlin_repack.cu deleted file mode 100644 index 1c8de396d..000000000 --- a/sgl-kernel/csrc/gemm/marlin/awq_marlin_repack.cu +++ /dev/null @@ -1,253 +0,0 @@ -#include "marlin.cuh" - -namespace marlin { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 -template -__global__ void awq_marlin_repack_kernel( - uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { - return; -} -#else - -template -__global__ void awq_marlin_repack_kernel( - uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { - constexpr int pack_factor = 32 / num_bits; - - int k_tiles = size_k / tile_k_size; - int n_tiles = size_n / tile_n_size; - int block_k_tiles = div_ceil(k_tiles, gridDim.x); - - auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) { - return; - } - - int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - extern __shared__ int4 sh[]; - - constexpr int tile_n_ints = tile_n_size / pack_factor; - - constexpr int stage_n_threads = tile_n_ints / 4; - constexpr int stage_k_threads = tile_k_size; - constexpr int stage_size = stage_k_threads * stage_n_threads; - - auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - cp_async_fence(); - return; - } - - int first_n = n_tile_id * tile_n_size; - int first_n_packed = first_n / pack_factor; - - int4* sh_ptr = sh + stage_size * pipe; - - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - int first_k = k_tile_id * tile_k_size; - - cp_async4( - &sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast( - &(b_q_weight_ptr[(first_k + k_id) * (size_n / pack_factor) + first_n_packed + (n_id * 4)]))); - } - - cp_async_fence(); - }; - - auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - return; - } - - auto warp_id = threadIdx.x / 32; - auto th_id = threadIdx.x % 32; - - if (warp_id >= 4) { - return; - } - - int tc_col = th_id / 4; - int tc_row = (th_id % 4) * 2; - - constexpr int tc_offsets[4] = {0, 1, 8, 9}; - - int cur_n = warp_id * 16 + tc_col; - int cur_n_packed = cur_n / pack_factor; - int cur_n_pos = cur_n % pack_factor; - - constexpr int sh_stride = tile_n_ints; - constexpr uint32_t mask = (1 << num_bits) - 1; - - int4* sh_stage_ptr = sh + stage_size * pipe; - uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - - // Undo interleaving - int cur_n_pos_unpacked; - if constexpr (num_bits == 4) { - constexpr int undo_pack[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - cur_n_pos_unpacked = undo_pack[cur_n_pos]; - } else { - constexpr int undo_pack[4] = {0, 2, 1, 3}; - cur_n_pos_unpacked = undo_pack[cur_n_pos]; - } - - uint32_t vals[8]; -#pragma unroll - for (int i = 0; i < 4; i++) { - int cur_elem = tc_row + tc_offsets[i]; - - int packed_src_0 = sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem]; - int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) + sh_stride * cur_elem]; - - vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask; - vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask; - } - - constexpr int tile_size = tile_k_size * tile_n_size / pack_factor; - int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; - - // Result of: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (num_bits == 4) { - constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - - } else { - constexpr int pack_idx[4] = {0, 2, 1, 3}; - - uint32_t res1 = 0; - uint32_t res2 = 0; -#pragma unroll - for (int i = 0; i < 4; i++) { - res1 |= vals[pack_idx[i]] << (i * 8); - res2 |= vals[4 + pack_idx[i]] << (i * 8); - } - - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; - } - }; - - auto start_pipes = [&](int k_tile_id, int n_tile_id) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages - 1; pipe++) { - fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); - } - - wait_for_stage(); - }; -#pragma unroll - for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { - int n_tile_id = 0; - - start_pipes(k_tile_id, n_tile_id); - - while (n_tile_id < n_tiles) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages; pipe++) { - fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += repack_stages; - } - } -} -#endif -} // namespace marlin - -#define CALL_IF(NUM_BITS) \ - else if (num_bits == NUM_BITS) { \ - cudaFuncSetAttribute( \ - marlin::awq_marlin_repack_kernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, \ - max_shared_mem); \ - marlin::awq_marlin_repack_kernel \ - <<>>(b_q_weight_ptr, out_ptr, size_k, size_n); \ - } - -torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k, int64_t size_n, int64_t num_bits) { - // Verify compatibility with marlin tile of 16x64 - TORCH_CHECK( - size_k % marlin::tile_k_size == 0, - "size_k = ", - size_k, - " is not divisible by tile_k_size = ", - marlin::tile_k_size); - TORCH_CHECK( - size_n % marlin::tile_n_size == 0, - "size_n = ", - size_n, - " is not divisible by tile_n_size = ", - marlin::tile_n_size); - - TORCH_CHECK(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits); - int const pack_factor = 32 / num_bits; - - // Verify B - TORCH_CHECK(b_q_weight.size(0) == size_k, "b_q_weight.size(0) = ", b_q_weight.size(0), " is not size_k = ", size_k); - TORCH_CHECK( - (size_n / pack_factor) == b_q_weight.size(1), - "Shape mismatch: b_q_weight.size(1) = ", - b_q_weight.size(1), - ", size_n = ", - size_n, - ", pack_factor = ", - pack_factor); - - // Verify device and strides - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); - - // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); - auto options = torch::TensorOptions().dtype(b_q_weight.dtype()).device(b_q_weight.device()); - torch::Tensor out = torch::empty({size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, options); - - // Get ptrs - uint32_t const* b_q_weight_ptr = reinterpret_cast(b_q_weight.data_ptr()); - uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); - - // Get dev info - int dev = b_q_weight.get_device(); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); - int blocks; - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); - - if (false) { - } - CALL_IF(4) - CALL_IF(8) - else { - TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits); - } - - return out; -} diff --git a/sgl-kernel/csrc/gemm/marlin/gptq_marlin.cu b/sgl-kernel/csrc/gemm/marlin/gptq_marlin.cu deleted file mode 100644 index b3437fc4b..000000000 --- a/sgl-kernel/csrc/gemm/marlin/gptq_marlin.cu +++ /dev/null @@ -1,1120 +0,0 @@ -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#include "kernel.h" -#include "marlin_template.h" - -#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ - static_assert( \ - std::is_same::value || std::is_same::value, \ - "only float16 and bfloat16 is supported"); - -namespace marlin { - -__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; - -using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__global__ void permute_cols_kernel( - int4 const* __restrict__ a_int4_ptr, - int const* __restrict__ perm_int_ptr, - int4* __restrict__ out_int4_ptr, - int size_m, - int size_k, - int lda, - int block_rows) {} - -} // namespace marlin - -torch::Tensor gptq_marlin_gemm( - torch::Tensor& a, - std::optional c_or_none, - torch::Tensor& b_q_weight, - torch::Tensor& b_scales, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, - torch::Tensor& workspace, - sglang::ScalarTypeId const& b_q_type_id, - int64_t size_m, - int64_t size_n, - int64_t size_k, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float) { - TORCH_CHECK_NOT_IMPLEMENTED(false, "marlin_gemm(..) requires CUDA_ARCH >= 8.0"); - return torch::empty({1, 1}); -} - -#else - -// For a given "a" of size [M,K] performs a permutation of the K columns based -// on the given "perm" indices. -__global__ void permute_cols_kernel( - int4 const* __restrict__ a_int4_ptr, - int const* __restrict__ perm_int_ptr, - int4* __restrict__ out_int4_ptr, - int size_m, - int size_k, - int lda, - int block_rows) { - auto start_row = block_rows * blockIdx.x; - int finish_row = start_row + block_rows; - if (finish_row > size_m) { - finish_row = size_m; - } - int cur_block_rows = finish_row - start_row; - - int input_row_stride = lda * sizeof(half) / 16; - int output_row_stride = size_k * sizeof(half) / 16; - - auto permute_row = [&](int row) { - int iters = size_k / default_threads; - int rest = size_k % default_threads; - - int input_offset = row * input_row_stride; - int output_offset = row * output_row_stride; - - half const* a_row_half = reinterpret_cast(a_int4_ptr + input_offset); - half* out_half = reinterpret_cast(out_int4_ptr + output_offset); - - int base_k = 0; - - for (int i = 0; i < iters; i++) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - - base_k += default_threads; - } - - if (rest) { - if (threadIdx.x < rest) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - } - } - }; - - for (int i = 0; i < cur_block_rows; i++) { - int cur_row = start_row + i; - if (cur_row < size_m) { - permute_row(cur_row); - } - } -} - -typedef struct { - int thread_k; - int thread_n; - int num_threads; -} thread_config_t; - -thread_config_t small_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {128, 128, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -thread_config_t large_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {64, 256, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -typedef struct { - int blocks_per_sm; - thread_config_t tb_cfg; -} exec_config_t; - -int get_scales_cache_size( - thread_config_t const& th_config, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full) { - bool cache_scales_chunk = has_act_order && !is_k_full; - - int tb_n = th_config.thread_n; - int tb_k = th_config.thread_k; - - // Get max scale groups per thread-block - int tb_groups; - if (group_size == -1) { - tb_groups = 1; - } else if (group_size == 0) { - tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size - } else { - tb_groups = div_ceil(tb_k, group_size); - } - - if (cache_scales_chunk) { - int load_groups = tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K - load_groups = max(load_groups, 32); // We load at least 32 scale groups - return load_groups * tb_n * 2; - } else { - int tb_scales = tb_groups * tb_n * 2; - - return tb_scales * pipe_stages; - } -} - -int get_kernel_cache_size( - thread_config_t const& th_config, - int thread_m_blocks, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - int has_zp, - int is_zp_float) { - int pack_factor = 32 / num_bits; - - // Get B size - int tb_k = th_config.thread_k; - int tb_n = th_config.thread_n; - int tb_m = thread_m_blocks * 16; - int sh_a_size = pipe_stages * (tb_m * tb_k) * 2; - int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4; - int sh_red_size = tb_m * (tb_n + 8); - int sh_s_size = - get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full); - int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0; - int sh_zp_size = 0; - if (has_zp) { - if (is_zp_float) - sh_zp_size = sh_s_size; - else if (num_bits == 4) - sh_zp_size = sh_s_size / 4; - else if (num_bits == 8) - sh_zp_size = sh_s_size / 2; - } - - int total_size = max(sh_b_size, sh_red_size) + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size; - - return total_size; -} - -bool is_valid_config( - thread_config_t const& th_config, - int thread_m_blocks, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - int has_zp, - int is_zp_float, - int max_shared_mem) { - // Sanity - if (th_config.thread_k == -1 || th_config.thread_n == -1 || th_config.num_threads == -1) { - return false; - } - - // Verify K/N are divisible by thread K/N - if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { - return false; - } - - // Verify min for thread K/N - if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { - return false; - } - - // num_threads must be at least 128 (= 4 warps) - if (th_config.num_threads < 128) { - return false; - } - - // Check that pipeline fits into cache - int cache_size = get_kernel_cache_size( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float); - return cache_size <= max_shared_mem; -} - -#define _GET_IF( \ - W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \ - else if ( \ - q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && thread_n_blocks == THREAD_N_BLOCKS && \ - thread_k_blocks == THREAD_K_BLOCKS && m_block_size_8 == M_BLOCK_SIZE_8 && group_blocks == GROUP_BLOCKS && \ - num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) { \ - kernel = Marlin< \ - scalar_t, \ - W_TYPE.id(), \ - NUM_THREADS, \ - THREAD_M_BLOCKS, \ - THREAD_N_BLOCKS, \ - THREAD_K_BLOCKS, \ - M_BLOCK_SIZE_8, \ - pipe_stages, \ - GROUP_BLOCKS, \ - IS_ZP_FLOAT>; \ - } - -// COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false) -// this is the most common cases -// BIGGROUP: cases for big group size (group_blocks in [-1, 8]) -// FZP: cases for float-zero-point (is_zp_float = true) -// ACT: cases for act order case (group_blocks == 0) -// FP4: cases for nvfp4(e2m1) (group_blocks == 1) -#define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define COMMON_GET_IF(W_TYPE) \ - COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \ - COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \ - COMMON_GET_IF_M1(W_TYPE, 4, 8, 128) \ - COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \ - COMMON_GET_IF_M234(W_TYPE, 8, 4, 128) \ - COMMON_GET_IF_M234(W_TYPE, 4, 8, 128) - -#define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define BIGGROUP_GET_IF(W_TYPE) \ - BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \ - BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \ - BIGGROUP_GET_IF_M1(W_TYPE, 4, 8, 128) \ - BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \ - BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) \ - BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128) - -#define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) - -#define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) - -#define FP4_GET_IF(W_TYPE) \ - FP4_GET_IF_M1(W_TYPE, 8, 8, 256) \ - FP4_GET_IF_M1(W_TYPE, 8, 4, 128) \ - FP4_GET_IF_M1(W_TYPE, 4, 8, 128) \ - FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \ - FP4_GET_IF_M234(W_TYPE, 8, 4, 128) \ - FP4_GET_IF_M234(W_TYPE, 4, 8, 128) - -// We currently have 4-bit models only with group_blocks == 4 -#define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) - -#define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) - -#define FZP_GET_IF(W_TYPE) \ - FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \ - FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \ - FZP_GET_IF_M1(W_TYPE, 4, 8, 128) \ - FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \ - FZP_GET_IF_M234(W_TYPE, 8, 4, 128) \ - FZP_GET_IF_M234(W_TYPE, 4, 8, 128) - -// We currently have 4-bit models only with group_blocks == 4 -#define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) - -#define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) - -#define ACT_GET_IF(W_TYPE) \ - ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \ - ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \ - ACT_GET_IF_M1(W_TYPE, 4, 8, 128) \ - ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \ - ACT_GET_IF_M234(W_TYPE, 8, 4, 128) \ - ACT_GET_IF_M234(W_TYPE, 4, 8, 128) - -template -MarlinFuncPtr get_marlin_kernel( - const sglang::ScalarType q_type, - int thread_m_blocks, - int thread_n_blocks, - int thread_k_blocks, - bool m_block_size_8, - bool has_act_order, - bool has_zp, - int group_blocks, - int num_threads, - bool is_zp_float) { - int num_bits = q_type.size_bits(); - auto kernel = MarlinDefault; - if (false) { - } - - COMMON_GET_IF(sglang::kU4) - COMMON_GET_IF(sglang::kU4B8) - COMMON_GET_IF(sglang::kU8B128) - - FP4_GET_IF(sglang::kFE2M1f) - - BIGGROUP_GET_IF(sglang::kFE4M3fn) - - ACT_GET_IF(sglang::kU4B8) - ACT_GET_IF(sglang::kU8B128) - - if (std::is_same::value) { - if (false) { - } - FZP_GET_IF(sglang::kU4) - } - - return kernel; -} - -template -exec_config_t determine_exec_config( - const sglang::ScalarType& q_type, - int prob_m, - int prob_n, - int prob_k, - int thread_m_blocks, - bool m_block_size_8, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - bool has_zp, - bool is_zp_float, - int max_shared_mem, - int sms) { - exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; - thread_config_t* thread_configs = thread_m_blocks > 1 ? large_batch_thread_configs : small_batch_thread_configs; - int thread_configs_size = thread_m_blocks > 1 ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) - : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); - - for (int i = 0; i < thread_configs_size; i++) { - thread_config_t th_config = thread_configs[i]; - - if (!is_valid_config( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem)) { - continue; - } - - int cache_size = get_kernel_cache_size( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float); - - int group_blocks = 0; - if (!has_act_order) { - group_blocks = group_size == -1 ? -1 : group_size / 16; - } - - auto kernel = get_marlin_kernel( - q_type, - thread_m_blocks, - th_config.thread_n / 16, - th_config.thread_k / 16, - m_block_size_8, - has_act_order, - has_zp, - group_blocks, - th_config.num_threads, - is_zp_float); - - if (kernel == MarlinDefault) continue; - - // int m_tiles = div_ceil(prob_m, thread_m_blocks * 16); - // int n_tiles = prob_n / th_config.thread_n; - // int k_tiles = prob_k / th_config.thread_k; - - return {1, th_config}; - } - - return exec_cfg; -} - -template -void marlin_mm( - const void* A, - const void* B, - void* C, - void* C_tmp, - void* s, - void* s2, - void* zp, - void* g_idx, - void* perm, - void* a_tmp, - int prob_m, - int prob_n, - int prob_k, - int lda, - void* workspace, - sglang::ScalarType const& q_type, - bool has_act_order, - bool is_k_full, - bool has_zp, - int num_groups, - int group_size, - int dev, - cudaStream_t stream, - int thread_k_init, - int thread_n_init, - int sms, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float) { - if (has_zp) { - TORCH_CHECK( - q_type == sglang::kU4 || q_type == sglang::kU8, - "q_type must be u4 or u8 when has_zp = True. Got = ", - q_type.str()); - } else { - TORCH_CHECK( - q_type == sglang::kU4B8 || q_type == sglang::kU8B128 || q_type == sglang::kFE4M3fn || q_type == sglang::kFE2M1f, - "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " - "has_zp = False. Got = ", - q_type.str()); - } - - TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]"); - - int group_blocks = 0; - if (has_act_order) { - if (is_k_full) { - TORCH_CHECK(group_size != -1); - group_blocks = group_size / 16; - TORCH_CHECK( - prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); - } else { - TORCH_CHECK(group_size == 0); - group_blocks = 0; - } - } else { - if (group_size == -1) { - group_blocks = -1; - } else { - group_blocks = group_size / 16; - TORCH_CHECK( - prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); - } - } - - int num_bits = q_type.size_bits(); - const int4* A_ptr = (const int4*)A; - const int4* B_ptr = (const int4*)B; - int4* C_ptr = (int4*)C; - int4* C_tmp_ptr = (int4*)C_tmp; - const int4* s_ptr = (const int4*)s; - const uint16_t* s2_ptr = (const uint16_t*)s2; - const int4* zp_ptr = (const int4*)zp; - const int* g_idx_ptr = (const int*)g_idx; - const int* perm_ptr = (const int*)perm; - int4* a_tmp_ptr = (int4*)a_tmp; - - int* locks = (int*)workspace; - - if (has_act_order) { - // Permute A columns - int block_rows = div_ceil(prob_m, sms); - // avoid ">>>" being formatted to "> > >" - // clang-format off - permute_cols_kernel<<>>( - A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows); - // clang-format on - A_ptr = a_tmp_ptr; - lda = prob_k; - - // If we have a full K, then we can run the non-act-order version of Marlin - // (since the weight rows are reordered by increasing group ids, and by - // having a full K, we have full original groups) - if (is_k_full) has_act_order = false; - } - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); - - int max_par = 16; - if (prob_n <= 4096) max_par = 16 * 8; - int max_shared_mem_new = max_shared_mem; - int rest_m = prob_m; - int max_thread_m_blocks = 4; - while (rest_m) { - int par_count = rest_m / (max_thread_m_blocks * 16); - if (par_count > max_par) par_count = max_par; - int prob_m_split = par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; - - int thread_k = thread_k_init; - int thread_n = thread_n_init; - - int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks); - int m_block_size_8 = prob_m_split <= 8; - - // Set thread config - exec_config_t exec_cfg; - thread_config_t thread_tfg; - if (thread_k != -1 && thread_n != -1) { - thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; - exec_cfg = exec_config_t{1, thread_tfg}; - TORCH_CHECK(prob_n % thread_n == 0, "prob_n = ", prob_n, " is not divisible by thread_n = ", thread_n); - TORCH_CHECK(prob_k % thread_k == 0, "prob_k = ", prob_k, " is not divisible by thread_k = ", thread_k); - } else { - // Auto config - exec_cfg = determine_exec_config( - q_type, - prob_m_split, - prob_n, - prob_k, - thread_m_blocks, - m_block_size_8, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem, - sms); - thread_tfg = exec_cfg.tb_cfg; - if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) { - max_thread_m_blocks--; - continue; - } - } - - int num_threads = thread_tfg.num_threads; - thread_k = thread_tfg.thread_k; - thread_n = thread_tfg.thread_n; - int blocks = sms * exec_cfg.blocks_per_sm; - if (exec_cfg.blocks_per_sm > 1) max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; - - int thread_k_blocks = thread_k / 16; - int thread_n_blocks = thread_n / 16; - - TORCH_CHECK( - is_valid_config( - thread_tfg, - thread_m_blocks, - prob_m_split, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem_new), - "Invalid thread config: thread_m_blocks = ", - thread_m_blocks, - ", thread_k = ", - thread_tfg.thread_k, - ", thread_n = ", - thread_tfg.thread_n, - ", num_threads = ", - thread_tfg.num_threads, - " for MKN = [", - prob_m, - ", ", - prob_k, - ", ", - prob_n, - "] and num_bits = ", - num_bits, - ", prob_m_split = ", - prob_m_split, - ", group_size = ", - group_size, - ", has_act_order = ", - has_act_order, - ", is_k_full = ", - is_k_full, - ", has_zp = ", - has_zp, - ", is_zp_float = ", - is_zp_float, - ", max_shared_mem_new = ", - max_shared_mem_new); - - auto kernel = get_marlin_kernel( - q_type, - thread_m_blocks, - thread_n_blocks, - thread_k_blocks, - m_block_size_8, - has_act_order, - has_zp, - group_blocks, - num_threads, - is_zp_float); - - if (kernel == MarlinDefault) { - TORCH_CHECK( - false, - "Unsupported shapes: MNK = [", - prob_m, - ", ", - prob_n, - ", ", - prob_k, - "]", - ", has_act_order = ", - has_act_order, - ", num_groups = ", - num_groups, - ", group_size = ", - group_size, - ", prob_m_split = ", - prob_m_split, - ", thread_m_blocks = ", - thread_m_blocks, - ", thread_n_blocks = ", - thread_n_blocks, - ", thread_k_blocks = ", - thread_k_blocks, - ", num_threads = ", - num_threads, - ", num_bits = ", - num_bits); - } - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem_new); - - bool part_use_atomic_add = use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; - - // avoid ">>>" being formatted to "> > >" - // clang-format off - kernel<<>>( - A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr, num_groups, - prob_m_split, prob_n, prob_k, lda, locks, part_use_atomic_add, - use_fp32_reduce, max_shared_mem_new); - // clang-format on - - A_ptr += prob_m_split * (lda / 8); - C_ptr += prob_m_split * (prob_n / 8); - rest_m -= prob_m_split; - } -} - -} // namespace marlin - -torch::Tensor gptq_marlin_gemm( - torch::Tensor& a, - std::optional c_or_none, - torch::Tensor& b_q_weight, - torch::Tensor& b_scales, - std::optional const& global_scale_or_none, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, - torch::Tensor& workspace, - sglang::ScalarTypeId const& b_q_type_id, - int64_t size_m, - int64_t size_n, - int64_t size_k, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float) { - sglang::ScalarType const b_q_type = sglang::ScalarType::from_id(b_q_type_id); - int pack_factor = 32 / b_q_type.size_bits(); - - // Verify A - TORCH_CHECK(a.size(0) == size_m, "Shape mismatch: a.size(0) = ", a.size(0), ", size_m = ", size_m); - TORCH_CHECK(a.size(1) == size_k, "Shape mismatch: a.size(1) = ", a.size(1), ", size_k = ", size_k); - - // Verify B - TORCH_CHECK( - size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, - "size_k = ", - size_k, - " is not divisible by tile_size = ", - MARLIN_NAMESPACE_NAME::tile_size); - TORCH_CHECK( - (size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_weight.size(0), - "Shape mismatch: b_q_weight.size(0) = ", - b_q_weight.size(0), - ", size_k = ", - size_k, - ", tile_size = ", - MARLIN_NAMESPACE_NAME::tile_size); - TORCH_CHECK( - b_q_weight.size(1) % MARLIN_NAMESPACE_NAME::tile_size == 0, - "b_q_weight.size(1) = ", - b_q_weight.size(1), - " is not divisible by tile_size = ", - MARLIN_NAMESPACE_NAME::tile_size); - int actual_size_n = (b_q_weight.size(1) / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor; - TORCH_CHECK(size_n == actual_size_n, "size_n = ", size_n, ", actual_size_n = ", actual_size_n); - - // Verify device and strides - TORCH_CHECK(a.device().is_cuda(), "A is not on GPU"); - TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1"); - // We use int4 (16 bytes) to load A, so A must aligned to 16 bytes - TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8"); - TORCH_CHECK(((uint64_t)a.data_ptr()) % 16 == 0, "A must aligned to 16 bytes"); - - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - - TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU"); - TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous"); - - // thread_k: `k` size of a thread_tile in `weights` (can usually be left as - // auto -1) - int thread_k = -1; - // thread_n: `n` size of a thread_tile in `weights` (can usually be left as - // auto -1) - int thread_n = -1; - // sms: number of SMs to use for the kernel - int sms = -1; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, a.get_device()); - - // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(a)); - auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device()); - torch::Tensor c; - if (c_or_none.has_value()) { - c = c_or_none.value(); - TORCH_CHECK(c.device().is_cuda(), "c is not on GPU"); - TORCH_CHECK(c.is_contiguous(), "c is not contiguous"); - TORCH_CHECK(c.size(0) == size_m, "Shape mismatch: c.size(0) = ", c.size(0), ", size_m = ", size_m); - TORCH_CHECK(c.size(1) == size_n, "Shape mismatch: c.size(1) = ", c.size(1), ", size_n = ", size_n); - } else { - c = torch::empty({size_m, size_n}, options); - } - if (size_m == 0) return c; - - // Alloc C tmp buffer that is going to be used for the global reduce - torch::Tensor c_tmp; - auto options_fp32 = torch::TensorOptions().dtype(at::kFloat).device(a.device()); - if (use_fp32_reduce) { - int max_m_block_size = (size_m + 16 - 1) / 16 * 16; - max_m_block_size = min(max_m_block_size, 64); - int max_c_tmp_size = sms * max_m_block_size * MARLIN_NAMESPACE_NAME::max_thread_n; - c_tmp = torch::empty({max_c_tmp_size}, options_fp32); - } else { - c_tmp = torch::empty({0}, options_fp32); - } - - // Detect groupsize and act_order - int num_groups = -1; - int group_size = -1; - - int rank = b_scales.sizes().size(); - TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2"); - TORCH_CHECK(b_scales.size(1) == size_n, "b_scales dim 1 = ", b_scales.size(1), " is not size_n = ", size_n); - num_groups = b_scales.size(0); - - torch::Tensor g_idx, perm, a_tmp; - if (g_idx_or_none.has_value() && perm_or_none.has_value()) { - g_idx = g_idx_or_none.value(); - perm = perm_or_none.value(); - - TORCH_CHECK(g_idx.device().is_cuda(), "g_idx is not on GPU"); - TORCH_CHECK(g_idx.is_contiguous(), "g_idx is not contiguous"); - TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); - - // Verify g_idx and perm - TORCH_CHECK( - (g_idx.size(-1) == 0 && perm.size(-1) == 0) || (g_idx.size(-1) == size_k && perm.size(-1) == size_k), - "Unexpected g_idx.size(-1) = ", - g_idx.size(-1), - " and perm.size(-1) = ", - perm.size(-1), - ", where size_k = ", - size_k); - } else { - g_idx = torch::empty({0}, options); - perm = torch::empty({0}, options); - a_tmp = torch::empty({0}, options); - } - bool has_act_order = g_idx.size(-1) > 0 && perm.size(-1) > 0; - - if (has_act_order) { - a_tmp = torch::empty({size_m, size_k}, options); - if (is_k_full) { - TORCH_CHECK(num_groups > 1, "For act_order, num_groups must be > 1"); - TORCH_CHECK(size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by num_groups = ", num_groups); - group_size = size_k / num_groups; - } else { - group_size = 0; - } - - } else { - a_tmp = torch::empty({0}, options); - if (num_groups > 1) { - TORCH_CHECK( - size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by b_scales.size(0) = ", b_scales.size(0)); - group_size = size_k / num_groups; - } else { - group_size = -1; - } - } - - torch::Tensor global_scale; - if (global_scale_or_none.has_value()) { - global_scale = global_scale_or_none.value(); - TORCH_CHECK(b_q_type == sglang::kFE2M1f, "global_scale can only be used for float4_e2m1f."); - } else { - global_scale = torch::empty({0}, options); - TORCH_CHECK(!(b_q_type == sglang::kFE2M1f), "the global_scale parameter must be passed for float4_e2m1f."); - } - - torch::Tensor b_zeros; - if (b_zeros_or_none.has_value()) { - b_zeros = b_zeros_or_none.value(); - TORCH_CHECK(b_zeros.device().is_cuda(), "b_zeros is not on GPU"); - TORCH_CHECK(b_zeros.is_contiguous(), "b_zeros is not contiguous"); - } else { - b_zeros = torch::empty({0}, options); - } - bool has_zp = b_zeros.size(-1) > 0; - if (has_zp) { - TORCH_CHECK( - b_q_type == sglang::kU4 || b_q_type == sglang::kU8, - "b_q_type must be u4 or u8 when has_zp = True. Got = ", - b_q_type.str()); - } else { - TORCH_CHECK( - b_q_type == sglang::kU4B8 || b_q_type == sglang::kU8B128 || b_q_type == sglang::kFE4M3fn || - b_q_type == sglang::kFE2M1f, - "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or " - "float4_e2m1f when " - "has_zp = False. Got = ", - b_q_type.str()); - } - - if (has_zp && is_zp_float) { - TORCH_CHECK( - a.scalar_type() == at::ScalarType::Half, - "Computation type must be float16 (half) when using float zero " - "points."); - } - - // Verify b_zeros - if (has_zp) { - int rank = b_zeros.sizes().size(); - TORCH_CHECK(rank == 2, "b_zeros rank = ", rank, " is not 2"); - if (is_zp_float) { - TORCH_CHECK(b_zeros.size(1) == size_n, "b_zeros dim 1 = ", b_zeros.size(1), " is not size_n = ", size_n); - TORCH_CHECK( - num_groups == b_zeros.size(0), "b_zeros dim 0 = ", b_zeros.size(0), " is not num_groups = ", num_groups); - TORCH_CHECK(num_groups != -1, "num_groups must be != -1"); - } else { - TORCH_CHECK( - b_zeros.size(0) == num_groups, "b_zeros dim 0 = ", b_zeros.size(0), " is not num_groups = ", num_groups); - TORCH_CHECK( - b_zeros.size(1) == size_n / pack_factor, - "b_zeros dim 1 = ", - b_zeros.size(1), - " is not size_n / pack_factor = ", - size_n / pack_factor); - } - } - - // Verify workspace size - TORCH_CHECK( - size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0, - "size_n = ", - size_n, - ", is not divisible by min_thread_n = ", - MARLIN_NAMESPACE_NAME::min_thread_n); - - int min_workspace_size = sms; - TORCH_CHECK( - workspace.numel() >= min_workspace_size, - "workspace.numel = ", - workspace.numel(), - " is below min_workspace_size = ", - min_workspace_size); - - int dev = a.get_device(); - if (a.scalar_type() == at::ScalarType::Half) { - void* scales_ptr; - if (b_q_type == sglang::kFE2M1f) { - scales_ptr = b_scales.data_ptr(); - } else { - scales_ptr = b_scales.data_ptr(); - } - - marlin::marlin_mm( - a.data_ptr(), - b_q_weight.data_ptr(), - c.data_ptr(), - c_tmp.data_ptr(), - scales_ptr, - global_scale.data_ptr(), - b_zeros.data_ptr(), - g_idx.data_ptr(), - perm.data_ptr(), - a_tmp.data_ptr(), - size_m, - size_n, - size_k, - a.stride(0), - workspace.data_ptr(), - b_q_type, - has_act_order, - is_k_full, - has_zp, - num_groups, - group_size, - dev, - at::cuda::getCurrentCUDAStream(dev), - thread_k, - thread_n, - sms, - use_atomic_add, - use_fp32_reduce, - is_zp_float); - } else if (a.scalar_type() == at::ScalarType::BFloat16) { - void* scales_ptr; - if (b_q_type == sglang::kFE2M1f) { - scales_ptr = b_scales.data_ptr(); - } else { - scales_ptr = b_scales.data_ptr(); - } - - marlin::marlin_mm( - a.data_ptr(), - b_q_weight.data_ptr(), - c.data_ptr(), - c_tmp.data_ptr(), - scales_ptr, - global_scale.data_ptr(), - b_zeros.data_ptr(), - g_idx.data_ptr(), - perm.data_ptr(), - a_tmp.data_ptr(), - size_m, - size_n, - size_k, - a.stride(0), - workspace.data_ptr(), - b_q_type, - has_act_order, - is_k_full, - has_zp, - num_groups, - group_size, - dev, - at::cuda::getCurrentCUDAStream(dev), - thread_k, - thread_n, - sms, - use_atomic_add, - use_fp32_reduce, - is_zp_float); - } else { - TORCH_CHECK(false, "gpt_marlin_gemm only supports bfloat16 and float16"); - } - - return c; -} - -#endif diff --git a/sgl-kernel/csrc/gemm/marlin/gptq_marlin_repack.cu b/sgl-kernel/csrc/gemm/marlin/gptq_marlin_repack.cu deleted file mode 100644 index 3bdf84161..000000000 --- a/sgl-kernel/csrc/gemm/marlin/gptq_marlin_repack.cu +++ /dev/null @@ -1,329 +0,0 @@ -#include "marlin.cuh" - -namespace marlin { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 -template -__global__ void gptq_marlin_repack_kernel( - uint32_t const* __restrict__ b_q_weight_ptr, - uint32_t const* __restrict__ perm_ptr, - uint32_t* __restrict__ out_ptr, - int size_k, - int size_n) { - return; -} -#else -template -__global__ void gptq_marlin_repack_kernel( - uint32_t const* __restrict__ b_q_weight_ptr, - uint32_t const* __restrict__ perm_ptr, - uint32_t* __restrict__ out_ptr, - int size_k, - int size_n) { - constexpr int pack_factor = 32 / num_bits; - - int k_tiles = size_k / tile_k_size; - int n_tiles = size_n / tile_n_size; - int block_k_tiles = div_ceil(k_tiles, gridDim.x); - - auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) { - return; - } - - int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - extern __shared__ int4 sh[]; - - constexpr int perm_size = tile_k_size / 4; - - int4* sh_perm_ptr = sh; - int4* sh_pipe_ptr = sh_perm_ptr; - if constexpr (has_perm) { - sh_pipe_ptr += perm_size; - } - - constexpr int tile_ints = tile_k_size / pack_factor; - - constexpr int stage_n_threads = tile_n_size / 4; - constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints; - constexpr int stage_size = stage_k_threads * stage_n_threads; - - auto load_perm_to_shared = [&](int k_tile_id) { - int first_k_int4 = (k_tile_id * tile_k_size) / 4; - - int4 const* perm_int4_ptr = reinterpret_cast(perm_ptr); - - if (threadIdx.x < perm_size) { - sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; - } - __syncthreads(); - }; - - auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - cp_async_fence(); - return; - } - - int first_n = n_tile_id * tile_n_size; - - int4* sh_ptr = sh_pipe_ptr + stage_size * pipe; - - if constexpr (has_perm) { - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - uint32_t const* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - - int src_k = sh_perm_int_ptr[k_id]; - int src_k_packed = src_k / pack_factor; - - cp_async4( - &sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast(&(b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); - } - - } else { - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - int first_k = k_tile_id * tile_k_size; - int first_k_packed = first_k / pack_factor; - - cp_async4( - &sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast(&(b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + (n_id * 4)]))); - } - } - - cp_async_fence(); - }; - - auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - return; - } - - auto warp_id = threadIdx.x / 32; - auto th_id = threadIdx.x % 32; - - if (warp_id >= 4) { - return; - } - - int tc_col = th_id / 4; - int tc_row = (th_id % 4) * 2; - - constexpr int tc_offsets[4] = {0, 1, 8, 9}; - - int cur_n = warp_id * 16 + tc_col; - - constexpr int sh_stride = 64; - constexpr uint32_t mask = (1 << num_bits) - 1; - - int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; - uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - - uint32_t* sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - - uint32_t vals[8]; - - if constexpr (has_perm) { - for (int i = 0; i < 4; i++) { - int k_idx = tc_row + tc_offsets[i]; - - uint32_t src_k = sh_perm_int_ptr[k_idx]; - uint32_t src_k_pos = src_k % pack_factor; - - uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; - uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; - - uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; - uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; - - vals[i] = b1_cur_val; - vals[4 + i] = b2_cur_val; - } - - } else { - uint32_t b1_vals[tile_ints]; - uint32_t b2_vals[tile_ints]; - -#pragma unroll - for (int i = 0; i < tile_ints; i++) { - b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; - b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; - } - -#pragma unroll - for (int i = 0; i < 4; i++) { - int cur_elem = tc_row + tc_offsets[i]; - int cur_int = cur_elem / pack_factor; - int cur_pos = cur_elem % pack_factor; - - vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; - vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; - } - } - - constexpr int tile_size = tile_k_size * tile_n_size / pack_factor; - int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; - - // Result of: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (num_bits == 4) { - constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - - } else { - constexpr int pack_idx[4] = {0, 2, 1, 3}; - - uint32_t res1 = 0; - uint32_t res2 = 0; -#pragma unroll - for (int i = 0; i < 4; i++) { - res1 |= vals[pack_idx[i]] << (i * 8); - res2 |= vals[4 + pack_idx[i]] << (i * 8); - } - - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; - } - }; - - auto start_pipes = [&](int k_tile_id, int n_tile_id) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages - 1; pipe++) { - fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); - } - - wait_for_stage(); - }; -#pragma unroll - for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { - int n_tile_id = 0; - - if constexpr (has_perm) { - load_perm_to_shared(k_tile_id); - } - - start_pipes(k_tile_id, n_tile_id); - - while (n_tile_id < n_tiles) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages; pipe++) { - fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += repack_stages; - } - } -} -#endif -} // namespace marlin - -#define CALL_IF(NUM_BITS, HAS_PERM) \ - else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \ - cudaFuncSetAttribute( \ - marlin::gptq_marlin_repack_kernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, \ - max_shared_mem); \ - marlin::gptq_marlin_repack_kernel \ - <<>>( \ - b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ - } - -torch::Tensor -gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, int64_t size_k, int64_t size_n, int64_t num_bits) { - // Verify compatibility with marlin tile of 16x64 - TORCH_CHECK( - size_k % marlin::tile_k_size == 0, - "size_k = ", - size_k, - " is not divisible by tile_k_size = ", - marlin::tile_k_size); - TORCH_CHECK( - size_n % marlin::tile_n_size == 0, - "size_n = ", - size_n, - " is not divisible by tile_n_size = ", - marlin::tile_n_size); - - TORCH_CHECK(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits); - int const pack_factor = 32 / num_bits; - - // Verify B - TORCH_CHECK( - (size_k / pack_factor) == b_q_weight.size(0), - "Shape mismatch: b_q_weight.size(0) = ", - b_q_weight.size(0), - ", size_k = ", - size_k, - ", pack_factor = ", - pack_factor); - TORCH_CHECK(b_q_weight.size(1) == size_n, "b_q_weight.size(1) = ", b_q_weight.size(1), " is not size_n = ", size_n); - - // Verify device and strides - TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU"); - TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous"); - TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt"); - - TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU"); - TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous"); - TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt"); - - // Alloc buffers - const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight)); - auto options = torch::TensorOptions().dtype(b_q_weight.dtype()).device(b_q_weight.device()); - torch::Tensor out = torch::empty({size_k / marlin::tile_size, size_n * marlin::tile_size / pack_factor}, options); - - // Detect if there is act_order - bool has_perm = perm.size(0) != 0; - - // Get ptrs - uint32_t const* b_q_weight_ptr = reinterpret_cast(b_q_weight.data_ptr()); - uint32_t const* perm_ptr = reinterpret_cast(perm.data_ptr()); - uint32_t* out_ptr = reinterpret_cast(out.data_ptr()); - - // Get dev info - int dev = b_q_weight.get_device(); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); - int blocks; - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev); - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - TORCH_CHECK(max_shared_mem > 0); - - if (false) { - } - CALL_IF(4, false) - CALL_IF(4, true) - CALL_IF(8, false) - CALL_IF(8, true) - else { - TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits, ", has_perm = ", has_perm); - } - - return out; -} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 3b10d90a9..9072072cf 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -261,25 +261,6 @@ void bmm_fp8( void dsv3_router_gemm(torch::Tensor& output, const torch::Tensor& mat_a, const torch::Tensor& mat_b); void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b); -torch::Tensor gptq_marlin_gemm( - torch::Tensor& a, - std::optional c_or_none, - torch::Tensor& b_q_weight, - torch::Tensor& b_scales, - std::optional const& global_scale_or_none, - std::optional const& b_zeros_or_none, - std::optional const& g_idx_or_none, - std::optional const& perm_or_none, - torch::Tensor& workspace, - sglang::ScalarTypeId const& b_q_type_id, - int64_t size_m, - int64_t size_n, - int64_t size_k, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float); - torch::Tensor gptq_gemm( torch::Tensor a, torch::Tensor b_q_weight, @@ -291,11 +272,6 @@ torch::Tensor gptq_gemm( void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit); -torch::Tensor -gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm, int64_t size_k, int64_t size_n, int64_t num_bits); - -torch::Tensor awq_marlin_repack(torch::Tensor& b_q_weight, int64_t size_k, int64_t size_n, int64_t num_bits); - /* * From csrc/moe */ diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index dc97be8c0..a24713a24 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -48,7 +48,6 @@ from sgl_kernel.gemm import ( fp8_blockwise_scaled_mm, fp8_scaled_mm, gptq_gemm, - gptq_marlin_gemm, gptq_shuffle, int8_scaled_mm, qserve_w4a8_per_chn_gemm, @@ -78,11 +77,6 @@ from sgl_kernel.mamba import ( causal_conv1d_update_cpu, chunk_gated_delta_rule_cpu, ) -from sgl_kernel.marlin import ( - awq_marlin_moe_repack, - awq_marlin_repack, - gptq_marlin_repack, -) from sgl_kernel.memory import set_kv_buffer_kernel, weak_ref_tensor from sgl_kernel.moe import ( apply_shuffle_mul_sum, diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py index d98a68323..e9e11c834 100644 --- a/sgl-kernel/python/sgl_kernel/gemm.py +++ b/sgl-kernel/python/sgl_kernel/gemm.py @@ -1,7 +1,6 @@ from typing import Optional, Tuple import torch -from sgl_kernel.scalar_type import ScalarType from sgl_kernel.utils import _get_cache_buf @@ -506,46 +505,6 @@ def scaled_fp4_experts_quant( # GPTQ kernels -def gptq_marlin_gemm( - a: torch.Tensor, - c: Optional[torch.Tensor], - b_q_weight: torch.Tensor, - b_scales: torch.Tensor, - global_scale: Optional[torch.Tensor], - b_zeros: Optional[torch.Tensor], - g_idx: Optional[torch.Tensor], - perm: Optional[torch.Tensor], - workspace: torch.Tensor, - b_q_type: ScalarType, - size_m: int, - size_n: int, - size_k: int, - is_k_full: bool = True, - use_atomic_add: bool = False, - use_fp32_reduce: bool = False, - is_zp_float: bool = False, -) -> torch.Tensor: - return torch.ops.sgl_kernel.gptq_marlin_gemm( - a, - c, - b_q_weight, - b_scales, - global_scale, - b_zeros, - g_idx, - perm, - workspace, - b_q_type.id, - size_m, - size_n, - size_k, - is_k_full, - use_atomic_add, - use_fp32_reduce, - is_zp_float, - ) - - def gptq_gemm( a: torch.Tensor, b_q_weight: torch.Tensor, diff --git a/sgl-kernel/python/sgl_kernel/marlin.py b/sgl-kernel/python/sgl_kernel/marlin.py deleted file mode 100644 index 823d3fd4b..000000000 --- a/sgl-kernel/python/sgl_kernel/marlin.py +++ /dev/null @@ -1,44 +0,0 @@ -import torch - - -def gptq_marlin_repack( - b_q_weight, - perm, - size_k, - size_n, - num_bits, -) -> torch.Tensor: - return torch.ops.sgl_kernel.gptq_marlin_repack( - b_q_weight, - perm, - size_k, - size_n, - num_bits, - ) - - -def awq_marlin_repack( - b_q_weight: torch.Tensor, size_k: int, size_n: int, num_bits: int -) -> torch.Tensor: - return torch.ops.sgl_kernel.awq_marlin_repack(b_q_weight, size_k, size_n, num_bits) - - -def awq_marlin_moe_repack( - b_q_weight: torch.Tensor, - perm: torch.Tensor, - size_k: int, - size_n: int, - num_bits: int, -) -> torch.Tensor: - num_experts = b_q_weight.shape[0] - assert size_k % 16 == 0 - output = torch.empty( - (num_experts, size_k // 16, size_n * (num_bits // 2)), - device=b_q_weight.device, - dtype=b_q_weight.dtype, - ) - for e in range(num_experts): - output[e] = torch.ops.sgl_kernel.awq_marlin_repack( - b_q_weight[e], size_k, size_n, num_bits - ) - return output diff --git a/sgl-kernel/tests/test_marlin_gemm.py b/sgl-kernel/tests/test_marlin_gemm.py deleted file mode 100644 index 9c4d48a39..000000000 --- a/sgl-kernel/tests/test_marlin_gemm.py +++ /dev/null @@ -1,121 +0,0 @@ -import pytest -import torch -from sgl_kernel import gptq_marlin_gemm -from sgl_kernel.scalar_type import scalar_types - -from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace -from sglang.test.test_marlin_utils import awq_marlin_quantize, marlin_quantize - -MNK_FACTORS = [ - (1, 1, 1), - (1, 4, 8), - (1, 7, 5), - (13, 17, 67), - (26, 37, 13), - (67, 13, 11), - (257, 13, 11), - (658, 13, 11), -] - - -# uint4 for awq -# uint4b8 for gptq -@pytest.mark.parametrize("k_chunk", [128]) -@pytest.mark.parametrize("n_chunk", [64, 256]) -@pytest.mark.parametrize("quant_type", [scalar_types.uint4, scalar_types.uint4b8]) -@pytest.mark.parametrize("group_size", [-1, 32, 64, 128]) -@pytest.mark.parametrize("mnk_factors", MNK_FACTORS) -@pytest.mark.parametrize("act_order", [False, True]) -@pytest.mark.parametrize("is_k_full", [False, True]) -@pytest.mark.parametrize("use_atomic_add", [False, True]) -@pytest.mark.parametrize("use_fp32_reduce", [False, True]) -def test_gptq_marlin_gemm( - k_chunk, - n_chunk, - quant_type, - group_size, - mnk_factors, - act_order, - is_k_full, - use_atomic_add, - use_fp32_reduce, -): - m_factor, n_factor, k_factor = mnk_factors - has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8] - - size_m = m_factor - size_k = k_chunk * k_factor - size_n = n_chunk * n_factor - - if act_order: - if group_size == -1: - return - if group_size == size_k: - return - if has_zp: - return - - if size_k % group_size != 0: - return - - a_input = torch.randn((size_m, size_k), dtype=torch.float16, device="cuda") - b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") - - if has_zp: - # AWQ style, unsigned + runtime zero-point - if group_size == 16: - return - w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize( - b_weight, quant_type, group_size - ) - g_idx = None - sort_indices = None - marlin_s2 = None - else: - # GPTQ style, unsigned + symmetric bias - if group_size == 16: - return - w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize( - b_weight, quant_type, group_size, act_order - ) - marlin_zp = None - marlin_s2 = None - - workspace = marlin_make_workspace(w_ref.device) - - # marlin gemm - output = gptq_marlin_gemm( - a_input, - None, - marlin_q_w, - marlin_s, - marlin_s2, - marlin_zp, - g_idx, - sort_indices, - workspace, - quant_type, - a_input.shape[0], - b_weight.shape[1], - a_input.shape[1], - is_k_full=is_k_full, - use_atomic_add=use_atomic_add, - use_fp32_reduce=use_fp32_reduce, - is_zp_float=False, - ) - # ref gemm - output_ref = torch.matmul(a_input, w_ref) - - torch.cuda.synchronize() - - max_diff = torch.mean(torch.abs(output - output_ref)) / torch.mean( - torch.abs(output_ref) - ) - - assert max_diff < 0.04 - - -if __name__ == "__main__": - import subprocess - - subprocess.call(["pytest", "--tb=short", str(__file__)]) diff --git a/sgl-kernel/tests/test_marlin_repack.py b/sgl-kernel/tests/test_marlin_repack.py deleted file mode 100644 index 508099106..000000000 --- a/sgl-kernel/tests/test_marlin_repack.py +++ /dev/null @@ -1,148 +0,0 @@ -import numpy as np -import pytest -import torch -from sgl_kernel import awq_marlin_repack, gptq_marlin_repack -from sgl_kernel.scalar_type import scalar_types - -from sglang.srt.layers.quantization.utils import ( - gptq_quantize_weights, - pack_cols, - pack_rows, - quantize_weights, - sort_weights, -) -from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights - -GPTQ_MARLIN_TILE = 16 -MARLIN_K_CHUNKS = [128] -MARLIN_N_CHUNKS = [64, 256] - -MNK_FACTORS = [ - (1, 1, 1), - (1, 4, 8), - (1, 7, 5), - (13, 17, 67), - (26, 37, 13), - (67, 13, 11), - (257, 13, 11), - (658, 13, 11), -] - - -def awq_pack( - q_w: torch.Tensor, - num_bits: int, - size_k: int, - size_n: int, -): - assert q_w.shape == (size_k, size_n) - - # Interleave column dim (for the dequantize code) and pack it to int32 - if num_bits == 4: - interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) - elif num_bits == 8: - interleave = np.array([0, 2, 1, 3]) - else: - raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) - - q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() - q_w = q_w.reshape((-1, size_n)).contiguous() - - return pack_cols(q_w, num_bits, size_k, size_n) - - -@pytest.mark.parametrize("num_bits", [4, 8]) -@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)]) -@pytest.mark.parametrize("group_size", [16, 32]) -def test_awq_marlin_repack_correct(num_bits, k_tiles, n_tiles, group_size): - tile_k, tile_n = 16, 64 - size_k = k_tiles * tile_k - size_n = n_tiles * tile_n - pack_factor = 32 // num_bits - - b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") - - w_ref, q_w, s, zp = quantize_weights( - b_weight, scalar_types.uint4, group_size, zero_points=True - ) - - q_w_awq = awq_pack(q_w, num_bits, size_k, size_n) - - weight_perm = get_weight_perm(num_bits) - q_w_marlin = marlin_weights(q_w, size_k, size_n, num_bits, weight_perm) - - out_gpu = awq_marlin_repack(q_w_awq, size_k, size_n, num_bits) - assert out_gpu.is_cuda and out_gpu.dtype == torch.int32 - - expected_cols = size_n * tile_k // pack_factor - assert list(out_gpu.shape) == [size_k // tile_k, expected_cols] - - torch.cuda.synchronize() - - torch.testing.assert_close(out_gpu, q_w_marlin) - - -@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS) -@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS) -@pytest.mark.parametrize("quant_type", [scalar_types.uint4b8]) -@pytest.mark.parametrize("group_size", [-1, 32, 64, 128]) -@pytest.mark.parametrize("act_order", [False, True]) -@pytest.mark.parametrize("mnk_factors", MNK_FACTORS) -def test_gptq_marlin_repack( - k_chunk, n_chunk, quant_type, group_size, act_order, mnk_factors -): - m_factor, n_factor, k_factor = mnk_factors - - size_k = k_chunk * k_factor - size_n = n_chunk * n_factor - - # Filter act_order - if act_order: - if group_size == -1: - return - if group_size == size_k: - return - - # Normalize group_size - if group_size == -1: - group_size = size_k - assert group_size <= size_k - - if size_k % group_size != 0: - pytest.skip("size_k must be divisible by group_size") - - # Create input - b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") - - # Quantize (and apply act_order if provided) - w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights( - b_weight, quant_type, group_size, act_order - ) - - q_w_gptq = pack_rows(q_w, quant_type.size_bits, size_k, size_n) - - # For act_order, sort the "weights" and "g_idx" so that group ids are - # increasing - sort_indices = torch.empty(0, dtype=torch.int, device=b_weight.device) - if act_order: - q_w, g_idx, sort_indices = sort_weights(q_w, g_idx) - - marlin_layout_perm = get_weight_perm(quant_type.size_bits) - q_w_marlin_ref = marlin_weights( - q_w, size_k, size_n, quant_type.size_bits, marlin_layout_perm - ) - - # Run Marlin repack GPU kernel - q_w_marlin = gptq_marlin_repack( - q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits - ) - - torch.cuda.synchronize() - - torch.testing.assert_close(q_w_marlin, q_w_marlin_ref) - - -if __name__ == "__main__": - import subprocess - - subprocess.call(["pytest", "--tb=short", str(__file__)])