diff --git a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py new file mode 100644 index 000000000..5ec65efec --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py @@ -0,0 +1,104 @@ +import os + +import torch +import triton +import triton.testing +from sgl_kernel.scalar_type import scalar_types + +from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack as jit_fn +from sglang.srt.layers.quantization.utils import gptq_quantize_weights, pack_rows + +try: + from sgl_kernel import gptq_marlin_repack as aot_fn + + AOT_AVAILABLE = True +except ImportError: + AOT_AVAILABLE = False + +IS_CI = ( + os.getenv("CI", "false").lower() == "true" + or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" +) + +# Fixed problem dimensions +SIZE_N = 4096 +NUM_BITS = 4 +QUANT_TYPE = scalar_types.uint4b8 +GROUP_SIZE = 128 + +# Pre-compute quantized weight for each size_k in the sweep +_cache = {} + + +def _get_inputs(size_k): + if size_k not in _cache: + size_n = SIZE_N + b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda") + _, q_w, _, _, _ = gptq_quantize_weights( + b_weight, QUANT_TYPE, GROUP_SIZE, act_order=False + ) + q_w_gptq = pack_rows(q_w, NUM_BITS, size_k, size_n) + sort_indices = torch.empty(0, dtype=torch.int, device="cuda") + _cache[size_k] = (q_w_gptq, sort_indices) + return _cache[size_k] + + +def check_correctness(): + if not AOT_AVAILABLE: + print("sgl_kernel AOT not available, skipping correctness check") + return + size_k = 4096 + q_w_gptq, sort_indices = _get_inputs(size_k) + out_jit = jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS) + out_aot = aot_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS) + torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0) + print("Correctness check passed (JIT vs AOT)") + + +if IS_CI: + k_range = [128, 1024, 4096] +else: + k_range = [128, 256, 512, 1024, 2048, 4096, 8192] + +if AOT_AVAILABLE: + line_vals = ["jit", "aot"] + line_names = ["JIT Kernel", "AOT Kernel"] + styles = [("blue", "-"), ("green", "-")] +else: + line_vals = ["jit"] + line_names = ["JIT Kernel"] + styles = [("blue", "-")] + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["size_k"], + x_vals=k_range, + line_arg="provider", + line_vals=line_vals, + line_names=line_names, + styles=styles, + ylabel="us", + plot_name="gptq-marlin-repack-performance", + args={}, + ) +) +def benchmark(size_k, provider): + q_w_gptq, sort_indices = _get_inputs(size_k) + + quantiles = [0.5, 0.2, 0.8] + + if provider == "jit": + fn = lambda: jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS) + elif provider == "aot": + fn = lambda: aot_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS) + else: + raise ValueError(f"Unknown provider: {provider}") + + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles) + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +if __name__ == "__main__": + check_correctness() + benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh b/python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh new file mode 100644 index 000000000..73bce7903 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh @@ -0,0 +1,362 @@ +/* + * 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 + */ + +#pragma once + +#include + +#include + +#include "marlin.cuh" + +namespace device::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 device::marlin + +#define CALL_IF_REPACK(NUM_BITS, HAS_PERM) \ + else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \ + host::RuntimeDeviceCheck(cudaFuncSetAttribute( \ + device::marlin::gptq_marlin_repack_kernel, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + max_shared_mem)); \ + host::LaunchKernel(blocks, device::marlin::repack_threads, stream, static_cast(max_shared_mem))( \ + device::marlin::gptq_marlin_repack_kernel, \ + b_q_weight_ptr, \ + perm_ptr, \ + out_ptr, \ + size_k, \ + size_n); \ + } + +void gptq_marlin_repack( + tvm::ffi::TensorView b_q_weight, + tvm::ffi::TensorView perm, + tvm::ffi::TensorView out, + int64_t size_k, + int64_t size_n, + int64_t num_bits) { + using namespace host; + + // Validate num_bits + RuntimeCheck(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits); + int const pack_factor = 32 / static_cast(num_bits); + + // Validate size alignment + RuntimeCheck( + size_k % device::marlin::tile_k_size == 0, + "size_k = ", + size_k, + " is not divisible by tile_k_size = ", + device::marlin::tile_k_size); + RuntimeCheck( + size_n % device::marlin::tile_n_size == 0, + "size_n = ", + size_n, + " is not divisible by tile_n_size = ", + device::marlin::tile_n_size); + + // Validate b_q_weight + auto bqw_dim0 = SymbolicSize{"bqw_dim0"}; + auto bqw_dim1 = SymbolicSize{"bqw_dim1"}; + bqw_dim0.set_value(size_k / pack_factor); + bqw_dim1.set_value(size_n); + auto device_ = SymbolicDevice{}; + device_.set_options(); + TensorMatcher({bqw_dim0, bqw_dim1}).with_dtype().with_device(device_).verify(b_q_weight); + + // Validate out + auto out_dim0 = SymbolicSize{"out_dim0"}; + auto out_dim1 = SymbolicSize{"out_dim1"}; + out_dim0.set_value(size_k / device::marlin::tile_size); + out_dim1.set_value(size_n * device::marlin::tile_size / pack_factor); + TensorMatcher({out_dim0, out_dim1}).with_dtype().with_device(device_).verify(out); + + // 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 + DLDevice dl_device = device_.unwrap(); + int dev = dl_device.device_id; + cudaStream_t stream = LaunchKernel::resolve_device(dl_device); + int blocks; + RuntimeDeviceCheck(cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev)); + + int max_shared_mem = 0; + RuntimeDeviceCheck(cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)); + RuntimeCheck(max_shared_mem > 0, "max_shared_mem must be > 0"); + + if (false) { + } + CALL_IF_REPACK(4, false) + CALL_IF_REPACK(4, true) + CALL_IF_REPACK(8, false) + CALL_IF_REPACK(8, true) + else { + Panic("Unsupported repack config: num_bits = ", num_bits, ", has_perm = ", has_perm); + } +} + +#undef CALL_IF_REPACK diff --git a/python/sglang/jit_kernel/gptq_marlin_repack.py b/python/sglang/jit_kernel/gptq_marlin_repack.py new file mode 100644 index 000000000..f04a2ce81 --- /dev/null +++ b/python/sglang/jit_kernel/gptq_marlin_repack.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import cache_once, load_jit + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +# Constants matching device::marlin:: in marlin.cuh +_TILE_SIZE = 16 + + +@cache_once +def _jit_gptq_marlin_repack_module() -> Module: + return load_jit( + "gptq_marlin_repack", + cuda_files=["gemm/marlin/gptq_marlin_repack.cuh"], + cuda_wrappers=[("gptq_marlin_repack", "gptq_marlin_repack")], + ) + + +def gptq_marlin_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, +) -> torch.Tensor: + pack_factor = 32 // num_bits + + # Allocate output tensor + out = torch.empty( + (size_k // _TILE_SIZE, size_n * _TILE_SIZE // pack_factor), + dtype=b_q_weight.dtype, + device=b_q_weight.device, + ) + + module = _jit_gptq_marlin_repack_module() + module.gptq_marlin_repack(b_q_weight, perm, out, size_k, size_n, num_bits) + return out diff --git a/python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py b/python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py new file mode 100644 index 000000000..dadfd0ecf --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py @@ -0,0 +1,101 @@ +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 +from sglang.srt.layers.quantization.utils import ( + gptq_quantize_weights, + pack_rows, + sort_weights, +) +from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights + +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), +] + + +@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 JIT repack kernel + jit_output = 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 + + subprocess.call(["pytest", "--tb=short", str(__file__)]) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py index 1d28412e8..415bb7960 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py @@ -43,7 +43,7 @@ from sglang.srt.utils import is_cuda _is_cuda = is_cuda() if _is_cuda: - from sgl_kernel import gptq_marlin_repack + from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack ScalarType, scalar_types = get_scalar_types() diff --git a/python/sglang/srt/layers/quantization/gptq.py b/python/sglang/srt/layers/quantization/gptq.py index a93988e7b..baf4c3d94 100644 --- a/python/sglang/srt/layers/quantization/gptq.py +++ b/python/sglang/srt/layers/quantization/gptq.py @@ -61,7 +61,9 @@ if TYPE_CHECKING: _is_cuda = is_cuda() if _is_cuda: - from sgl_kernel import gptq_gemm, gptq_marlin_repack, gptq_shuffle + from sgl_kernel import gptq_gemm, gptq_shuffle + + from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack _is_npu = is_npu() diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp8.py b/python/sglang/srt/layers/quantization/marlin_utils_fp8.py index ee5af6858..e699b6798 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils_fp8.py +++ b/python/sglang/srt/layers/quantization/marlin_utils_fp8.py @@ -17,9 +17,8 @@ from sglang.srt.utils import is_cuda _is_cuda = is_cuda() if _is_cuda: - from sgl_kernel import gptq_marlin_repack - from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm + from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack ScalarType, scalar_types = get_scalar_types()