@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6a92ff93e1
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 21:00 UTC

[Misc][Kernel]: Add GPTQAllSpark Quantization (#12931)

YajieWang committed Mar 1, 2025 at 14:30 UTC 6a92ff93e1221ef0a5ff1528c2fdf05cecfb94ec
12 files changed +2005 -4
CMakeLists.txt
+16
@@ -317,6 +317,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
317 " in CUDA target architectures")
318 endif()
319
320 + # Only build AllSpark kernels if we are building for at least some compatible archs.
321 + cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
322 + if (ALLSPARK_ARCHS)
323 + set(ALLSPARK_SRCS
324 + "csrc/quantization/gptq_allspark/allspark_repack.cu"
325 + "csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
326 + set_gencode_flags_for_srcs(
327 + SRCS "${ALLSPARK_SRCS}"
328 + CUDA_ARCHS "${ALLSPARK_ARCHS}")
329 + list(APPEND VLLM_EXT_SRC "${ALLSPARK_SRCS}")
330 + message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
331 + else()
332 + message(STATUS "Not building AllSpark kernels as no compatible archs found"
333 + " in CUDA target architectures")
334 + endif()
335 +
336 # The cutlass_scaled_mm kernels for Hopper (c3x, i.e. CUTLASS 3.x) require
337 # CUDA 12.0 or later (and only work on Hopper, 9.0a for now).
338 cuda_archs_loose_intersection(SCALED_MM_3X_ARCHS "9.0a" "${CUDA_ARCHS}")
benchmarks/kernels/benchmark_marlin.py
+45 -2
@@ -10,6 +10,8 @@ from vllm import _custom_ops as ops
10 from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
11 GPTQ_MARLIN_24_MAX_PARALLEL, GPTQ_MARLIN_24_MIN_THREAD_N,
12 GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES, GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES)
13 +from vllm.model_executor.layers.quantization.utils.allspark_utils import (
14 + ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD, ALLSPARK_SUPPORTED_QUANT_TYPES)
15 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
16 GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N,
17 MARLIN_SUPPORTED_GROUP_SIZES, query_marlin_supported_quant_types)
@@ -18,12 +20,12 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
20 from vllm.model_executor.layers.quantization.utils.marlin_utils_test_24 import (
21 marlin_24_quantize)
22 from vllm.model_executor.layers.quantization.utils.quant_utils import (
21 - gptq_pack, gptq_quantize_weights, sort_weights)
23 + gptq_pack, gptq_quantize_weights, quantize_weights, sort_weights)
24 from vllm.scalar_type import ScalarType
25 from vllm.utils import FlexibleArgumentParser
26
27 DEFAULT_MODELS = ["meta-llama/Llama-2-7b-hf/TP1"]
26 -DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
28 +DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
29
30 ACT_ORDER_OPTS = [False, True]
31 K_FULL_OPTS = [False, True]
@@ -81,6 +83,27 @@ def bench_run(results: List[benchmark.Measurement], model: str,
83 GPTQ_MARLIN_24_MAX_PARALLEL)
84 marlin_zp = torch.zeros_like(marlin_s, dtype=torch.int)
85
86 + # AllSpark W8A16 quant
87 + as_supported_case = (quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES
88 + and group_size == -1 and not act_order and is_k_full)
89 + if as_supported_case:
90 + properties = torch.cuda.get_device_properties(b.device.index)
91 + sm_count = properties.multi_processor_count
92 + sm_version = properties.major * 10 + properties.minor
93 +
94 + supported_arch = (sm_version >= 80 and sm_version < 90)
95 + as_supported_case = as_supported_case and supported_arch
96 + if supported_arch:
97 + has_zp = False
98 + w_ref, qw, s, zp = quantize_weights(b, quant_type, group_size,
99 + has_zp)
100 + qw = qw.to(torch.uint8)
101 +
102 + qw_reorder, s_reorder, zp_reorder = \
103 + ops.allspark_repack_weight(
104 + qw, s, zp, has_zp)
105 + CUBLAS_M_THRESHOLD = ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD
106 +
107 globals = {
108 # Gen params
109 "quant_type": quant_type,
@@ -109,10 +132,19 @@ def bench_run(results: List[benchmark.Measurement], model: str,
132 # GPTQ params
133 "q_w_gptq": q_w_gptq,
134 "repack_sort_indices": repack_sort_indices,
135 + # AllSpark W8A16 params
136 + "qw_reorder": qw_reorder if as_supported_case else None,
137 + "s_reorder": s_reorder if as_supported_case else None,
138 + "zp_reorder": zp_reorder if as_supported_case else None,
139 + "sm_count": sm_count if as_supported_case else None,
140 + "sm_version": sm_version if as_supported_case else None,
141 + "CUBLAS_M_THRESHOLD":
142 + CUBLAS_M_THRESHOLD if as_supported_case else None,
143 # Kernels
144 "gptq_marlin_gemm": ops.gptq_marlin_gemm,
145 "gptq_marlin_24_gemm": ops.gptq_marlin_24_gemm,
146 "gptq_marlin_repack": ops.gptq_marlin_repack,
147 + "allspark_w8a16_gemm": ops.allspark_w8a16_gemm,
148 }
149
150 min_run_time = 1
@@ -172,6 +204,17 @@ def bench_run(results: List[benchmark.Measurement], model: str,
204 description="gptq_marlin_repack",
205 ).blocked_autorange(min_run_time=min_run_time))
206
207 + if as_supported_case:
208 + results.append(
209 + benchmark.Timer(
210 + stmt=
211 + "output = allspark_w8a16_gemm(a, qw_reorder, s_reorder, zp_reorder, size_n, group_size, sm_count, sm_version, CUBLAS_M_THRESHOLD, False, True)", # noqa: E501
212 + globals=globals,
213 + label=label,
214 + sub_label=sub_label,
215 + description="allspark_w8a16_gemm_fp32",
216 + ).blocked_autorange(min_run_time=min_run_time))
217 +
218
219 def main(args):
220 print("Benchmarking models:")
csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu new
+1008
@@ -0,0 +1,1008 @@
1 +#include "allspark_utils.cuh"
2 +#include <torch/all.h>
3 +#include "core/registration.h"
4 +#include <cublas_v2.h>
5 +
6 +at::Tensor as_g_workspace;
7 +
8 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
9 +
10 +torch::Tensor allspark_w8a16_gemm(
11 + torch::Tensor const& a, torch::Tensor const& b_qweight,
12 + torch::Tensor const& b_scales, c10::optional<torch::Tensor> const& b_qzeros,
13 + int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
14 + int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
15 + TORCH_CHECK_NOT_IMPLEMENTED(
16 + false, "allspark_w8a16_gemm(..) requires CUDA_ARCH >= 8.0");
17 + return torch::empty({1, 1});
18 +}
19 +
20 +#else
21 +namespace allspark {
22 +/*
23 + * GemmTile manage data movement from Global Memory to Shared Memory
24 + * requiring N % 8 == 0， K % 16 == 0 by loading uint
25 + * BN is obtained by padding the original N to a multiple of 32
26 + * weight B is rearranged as N32K16 order,
27 + * i.e. a initial data block of size 32(n)x16(k) is reordered as n8k4n4k4，
28 + * in order to put data loaded by the same thread of 32x16 data block together
29 + * continuously (see
30 + * https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type)
31 + */
32 +template <typename FType, typename QType, int Mtile, int Ntile, int NStage,
33 + int BLOCK>
34 +struct GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
35 + // element num loaded by a LDG inst.
36 + static constexpr int LDG_ELEMENT_CNT_A = 8;
37 + static constexpr int LDG_ELEMENT_CNT_B = 16;
38 + static constexpr int WARP_SIZE = 32;
39 + static constexpr int M_SIZE_ONE_LOAD = (BLOCK * LDG_ELEMENT_CNT_A) / 32;
40 + static constexpr int N_SIZE_ONE_LOAD = (BLOCK * LDG_ELEMENT_CNT_B) / 32;
41 +
42 + __device__ GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK(
43 + const SM8x_GEMM_W8A16_Splitk_Params<FType, QType>& k_params,
44 + const uint32_t& A_smem_addr, const uint32_t& BQ_smem_addr,
45 + const uint32_t& A_stage_stride, const uint32_t& BQ_stage_stride)
46 + : params(k_params),
47 + A_smem_base_addr(A_smem_addr),
48 + BQ_smem_base_addr(BQ_smem_addr),
49 + A_smem_stage_stride(A_stage_stride),
50 + BQ_smem_stage_stride(BQ_stage_stride) {
51 + this_block_A_base_ptr = params.A_ptr + blockIdx.x * Mtile * params.K +
52 + blockIdx.z * params.SplitK;
53 + // here B is rearranged as N32K16 order, i.e. 4 continuous N-direction
54 + // 8(N)x16(K) size data blocks are packed together
55 + this_block_B_base_ptr = params.B_ptr + blockIdx.y * Ntile * params.K +
56 + blockIdx.z * params.SplitK * 4;
57 +
58 + const int lane_id = threadIdx.x % WARP_SIZE;
59 +
60 + // For matrix A, a block load/store Mtile(row) x 32(col) elements in
61 + // multiple iters, 8x4 warp load/store 8(row) x 32(col) elements per iter
62 + const int Aldg_row_base_idx = threadIdx.x / 4;
63 + Aldg_col_idx = (threadIdx.x % 4) * LDG_ELEMENT_CNT_A;
64 + const int Aldg_base_offset = Aldg_row_base_idx * params.K + Aldg_col_idx;
65 +
66 + // For matrix B, a block load/store elements of (Ntile / 4) row x 128 col
67 + // elements of N32K16 packing in multiple iters, 4x8 warp load/store 4(row)
68 + // * 128(col) per iter
69 + Bldg_col_idx = (threadIdx.x % 8) * LDG_ELEMENT_CNT_B;
70 + const int Bldg_row_base_idx = threadIdx.x / 8;
71 + const int Bldg_base_offset =
72 + Bldg_row_base_idx * params.K * 4 + Bldg_col_idx;
73 +
74 + this_block_A_base_ptr += Aldg_base_offset;
75 + this_block_B_base_ptr += Bldg_base_offset;
76 +
77 + const int sts_a_base_offset =
78 + (threadIdx.x / 4) * 32 +
79 + ((lane_id % 4) ^ ((lane_id / 4) % 4) ^ ((lane_id / 4) / 4)) *
80 + LDG_ELEMENT_CNT_A;
81 + const int sts_bq_base_offset =
82 + Bldg_row_base_idx * 32 * 4 +
83 + ((threadIdx.x % 8) ^ (((threadIdx.x / 8) % 2) * 4)) * LDG_ELEMENT_CNT_B;
84 +
85 + A_smem_base_addr += sts_a_base_offset * sizeof(FType);
86 + BQ_smem_base_addr += sts_bq_base_offset * sizeof(uint8_t);
87 +
88 + A_ldg_guard = 0;
89 + B_ldg_guard = 0;
90 + #pragma unroll
91 + for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD; ++i) {
92 + int m_idx = blockIdx.x * Mtile + Aldg_row_base_idx + i * M_SIZE_ONE_LOAD;
93 + if (m_idx < params.M) {
94 + A_ldg_guard |= (1u << i);
95 + }
96 + }
97 +
98 + const int N_padded = (params.N + 31) / 32 * 32;
99 + #pragma unroll
100 + for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD; ++i) {
101 + int n_idx = blockIdx.y * Ntile + (Bldg_row_base_idx / 8) * 32 +
102 + i * N_SIZE_ONE_LOAD;
103 + if (n_idx < N_padded) {
104 + B_ldg_guard |= (1u << i);
105 + }
106 + }
107 + }
108 +
109 + __device__ void ldgsts_first_ktiles(const int& first_k_tile,
110 + const int& k_tiles) {
111 + // load first k_tile
112 + // load A
113 + const int A_src_size = Aldg_col_idx < first_k_tile ? 16 : 0;
114 + #pragma unroll
115 + for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD; ++i) {
116 + cp_async<16>(
117 + A_smem_base_addr + (i * M_SIZE_ONE_LOAD * 32) * sizeof(FType),
118 + this_block_A_base_ptr + i * M_SIZE_ONE_LOAD * params.K, A_src_size,
119 + (A_ldg_guard & (1u << i)) != 0);
120 + }
121 +
122 + // load B
123 + const int B_src_size = (Bldg_col_idx / 4) < first_k_tile ? 16 : 0;
124 + #pragma unroll
125 + for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD; ++i) {
126 + cp_async<16>(
127 + BQ_smem_base_addr + (i * N_SIZE_ONE_LOAD * 32) * sizeof(uint8_t),
128 + this_block_B_base_ptr + i * N_SIZE_ONE_LOAD * params.K, B_src_size,
129 + (B_ldg_guard & (1u << i)) != 0);
130 + }
131 +
132 + cp_async_commit_group();
133 + this_block_A_base_ptr += first_k_tile;
134 + this_block_B_base_ptr += (first_k_tile * 4);
135 +
136 + // load second to (N-stage - 1) k_tiles
137 + for (int stage_idx = 1; stage_idx < NStage - 1; ++stage_idx) {
138 + if (stage_idx < k_tiles) {
139 + #pragma unroll
140 + for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD;
141 + ++i) {
142 + cp_async<16>(A_smem_base_addr + stage_idx * A_smem_stage_stride +
143 + (i * M_SIZE_ONE_LOAD * 32) * sizeof(FType),
144 + this_block_A_base_ptr + i * M_SIZE_ONE_LOAD * params.K,
145 + 16, (A_ldg_guard & (1u << i)) != 0);
146 + }
147 +
148 + #pragma unroll
149 + for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD;
150 + ++i) {
151 + cp_async<16>(BQ_smem_base_addr + stage_idx * BQ_smem_stage_stride +
152 + (i * N_SIZE_ONE_LOAD * 32) * sizeof(uint8_t),
153 + this_block_B_base_ptr + i * N_SIZE_ONE_LOAD * params.K,
154 + 16, (B_ldg_guard & (1u << i)) != 0);
155 + }
156 +
157 + this_block_A_base_ptr += 32;
158 + this_block_B_base_ptr += (32 * 4);
159 + }
160 + cp_async_commit_group();
161 + }
162 + }
163 +
164 + __device__ void ldgsts(const int& sts_stage_idx) {
165 + const int a_stage_offset = sts_stage_idx * A_smem_stage_stride;
166 + const int bq_stage_offset = sts_stage_idx * BQ_smem_stage_stride;
167 + #pragma unroll
168 + for (int i = 0; i < (Mtile + M_SIZE_ONE_LOAD - 1) / M_SIZE_ONE_LOAD; ++i) {
169 + cp_async<16>(A_smem_base_addr + a_stage_offset +
170 + (i * M_SIZE_ONE_LOAD * 32) * sizeof(FType),
171 + this_block_A_base_ptr + i * M_SIZE_ONE_LOAD * params.K, 16,
172 + (A_ldg_guard & (1u << i)) != 0);
173 + }
174 +
175 + #pragma unroll
176 + for (int i = 0; i < (Ntile + N_SIZE_ONE_LOAD - 1) / N_SIZE_ONE_LOAD; ++i) {
177 + cp_async<16>(BQ_smem_base_addr + bq_stage_offset +
178 + (i * N_SIZE_ONE_LOAD * 32) * sizeof(uint8_t),
179 + this_block_B_base_ptr + i * N_SIZE_ONE_LOAD * params.K, 16,
180 + (B_ldg_guard & (1u << i)) != 0);
181 + }
182 +
183 + cp_async_commit_group();
184 + this_block_A_base_ptr += 32;
185 + this_block_B_base_ptr += (32 * 4);
186 + }
187 +
188 + const FType* this_block_A_base_ptr = nullptr;
189 + const QType* this_block_B_base_ptr = nullptr;
190 +
191 + int Aldg_col_idx;
192 + int Bldg_col_idx;
193 +
194 + uint32_t A_ldg_guard;
195 + uint32_t B_ldg_guard;
196 +
197 + uint32_t A_smem_base_addr, BQ_smem_base_addr;
198 + const uint32_t A_smem_stage_stride, BQ_smem_stage_stride;
199 +
200 + const SM8x_GEMM_W8A16_Splitk_Params<FType, QType>& params;
201 +};
202 +
203 +/*
204 + * requiring N % 8 == 0
205 + */
206 +template <typename FType, typename QType, int Mtile, int Ntile, int BLOCK,
207 + bool EnableFuse, bool has_zp>
208 +struct ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK {
209 + static constexpr int WARP_SIZE = 32;
210 + static constexpr int WARP_CNT = BLOCK / WARP_SIZE;
211 + static constexpr int WARP_NTILE = Ntile / WARP_CNT;
212 + static constexpr int WARP_NITER = WARP_NTILE / 8; // hmma16816
213 + static_assert(WARP_NTILE == 32 or WARP_NTILE == 64,
214 + "now only support WARP_NTILE = 32 or 64!");
215 +
216 + __device__ ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK(
217 + const SM8x_GEMM_W8A16_Splitk_Params<FType, QType>& k_params,
218 + const uint32_t& A_smem_addr, const uint32_t& BQ_smem_addr,
219 + const uint32_t& A_stage_stride, const uint32_t& BQ_stage_stride)
220 + : params(k_params),
221 + A_smem_base_addr(A_smem_addr),
222 + BQ_smem_base_addr(BQ_smem_addr),
223 + A_smem_stage_stride(A_stage_stride),
224 + BQ_smem_stage_stride(BQ_stage_stride) {
225 + warp_id = threadIdx.x / WARP_SIZE;
226 + lane_id = threadIdx.x % WARP_SIZE;
227 +
228 + load_a_base_offset[0] =
229 + (lane_id % 16) * 32 +
230 + ((lane_id / 16) ^ (lane_id % 4) ^ ((lane_id / 4) % 2)) * 8;
231 + load_a_base_offset[1] =
232 + (lane_id % 16) * 32 +
233 + ((lane_id / 16 + 2) ^ (lane_id % 4) ^ ((lane_id / 4) % 2)) * 8;
234 +
235 + load_b_base_offset[0] =
236 + (lane_id / 4 + warp_id * (WARP_NTILE / 4)) * 32 * 4 +
237 + (lane_id % 4) * 16 + ((lane_id / 4) % 2) * 16 * 4;
238 + load_b_base_offset[1] =
239 + (lane_id / 4 + warp_id * (WARP_NTILE / 4)) * 32 * 4 +
240 + (lane_id % 4) * 16 + (((lane_id / 4) % 2) ^ 1) * 16 * 4;
241 +
242 + sts_c_base_offset = warp_id * Mtile * WARP_NTILE +
243 + (lane_id / 4) * WARP_NTILE + (lane_id % 4) * 2;
244 +
245 + if (EnableFuse) {
246 + this_block_C_base_ptr =
247 + params.C_ptr + blockIdx.x * Mtile * params.N + blockIdx.y * Ntile;
248 + } else {
249 + this_block_C_base_ptr =
250 + params.C_split_ptr + blockIdx.z * params.M * params.N +
251 + blockIdx.x * Mtile * params.N + blockIdx.y * Ntile;
252 + }
253 + int store_thds_in_row = WARP_NTILE / 8;
254 + store_c_row_base_idx = lane_id / store_thds_in_row;
255 + store_c_col_idx = warp_id * WARP_NTILE + (lane_id % store_thds_in_row) * 8;
256 + store_c_base_offset = store_c_row_base_idx * params.N + store_c_col_idx;
257 +
258 + #pragma unroll
259 + for (int i = 0; i < Mtile / 16; ++i) {
260 + #pragma unroll
261 + for (int j = 0; j < WARP_NITER; ++j) {
262 + #pragma unroll
263 + for (int k = 0; k < 4; ++k) {
264 + C_frag[i][j][k] = 0.f;
265 + }
266 + }
267 + }
268 + params_n_idx =
269 + blockIdx.y * Ntile + warp_id * WARP_NTILE + (lane_id / 4) * 4;
270 + }
271 +
272 + __device__ void lds(const int& smem_stage_idx, const int& reg_buf_idx,
273 + const int& k_phase_idx) {
274 + uint32_t A_smem_addr =
275 + A_smem_base_addr + A_smem_stage_stride * smem_stage_idx;
276 + uint32_t B_smem_addr =
277 + BQ_smem_base_addr + BQ_smem_stage_stride * smem_stage_idx;
278 +
279 + #pragma unroll
280 + for (int i = 0; i < Mtile / 16; ++i) {
281 + ldsm_4(A_frag[reg_buf_idx][i][0], A_frag[reg_buf_idx][i][1],
282 + A_frag[reg_buf_idx][i][2], A_frag[reg_buf_idx][i][3],
283 + A_smem_addr + (load_a_base_offset[k_phase_idx] + i * 16 * 32) *
284 + sizeof(FType));
285 + }
286 + #pragma unroll
287 + for (int i = 0; i < WARP_NTILE / 32; ++i) {
288 + lds128(BQ_frag[reg_buf_idx][4 * i + 0], BQ_frag[reg_buf_idx][4 * i + 1],
289 + BQ_frag[reg_buf_idx][4 * i + 2], BQ_frag[reg_buf_idx][4 * i + 3],
290 + B_smem_addr + (load_b_base_offset[k_phase_idx] + i * 32 * 32) *
291 + sizeof(uint8_t));
292 + }
293 +
294 + // dequant B
295 + #pragma unroll
296 + for (int i = 0; i < WARP_NITER / 2; ++i) {
297 + cvt_8bx4_to_16bx4_bias128(BQ_frag[reg_buf_idx][2 * i],
298 + BF_frag[reg_buf_idx][2 * i]);
299 + if (has_zp) {
300 + BF_frag[reg_buf_idx][2 * i][0] =
301 + __hsub2(BF_frag[reg_buf_idx][2 * i][0], num2num2(B_zero[i].x));
302 + BF_frag[reg_buf_idx][2 * i][1] =
303 + __hsub2(BF_frag[reg_buf_idx][2 * i][1], num2num2(B_zero[i].x));
304 + }
305 +
306 + BF_frag[reg_buf_idx][2 * i][0] =
307 + __hmul2(BF_frag[reg_buf_idx][2 * i][0], num2num2(B_scale[i].x));
308 + BF_frag[reg_buf_idx][2 * i][1] =
309 + __hmul2(BF_frag[reg_buf_idx][2 * i][1], num2num2(B_scale[i].x));
310 +
311 + cvt_8bx4_to_16bx4_bias128(BQ_frag[reg_buf_idx][2 * i + 1],
312 + BF_frag[reg_buf_idx][2 * i + 1]);
313 + if (has_zp) {
314 + BF_frag[reg_buf_idx][2 * i + 1][0] =
315 + __hsub2(BF_frag[reg_buf_idx][2 * i + 1][0], num2num2(B_zero[i].y));
316 + BF_frag[reg_buf_idx][2 * i + 1][1] =
317 + __hsub2(BF_frag[reg_buf_idx][2 * i + 1][1], num2num2(B_zero[i].y));
318 + }
319 +
320 + BF_frag[reg_buf_idx][2 * i + 1][0] =
321 + __hmul2(BF_frag[reg_buf_idx][2 * i + 1][0], num2num2(B_scale[i].y));
322 + BF_frag[reg_buf_idx][2 * i + 1][1] =
323 + __hmul2(BF_frag[reg_buf_idx][2 * i + 1][1], num2num2(B_scale[i].y));
324 + }
325 + }
326 +
327 + __device__ void ldg_params() {
328 + const int N_padded = (params.N + 31) / 32 * 32;
329 + // load B scale and zero_point
330 + #pragma unroll
331 + for (int i = 0; i < WARP_NTILE / 32; ++i) {
332 + ldg64_ca(B_scale[2 * i + 0], B_scale[2 * i + 1],
333 + params.B_scale_ptr + params_n_idx + i * 32,
334 + (params_n_idx + i * 32) < N_padded);
335 + if (has_zp) {
336 + ldg64_ca(B_zero[2 * i + 0], B_zero[2 * i + 1],
337 + params.B_zero_ptr + params_n_idx + i * 32,
338 + (params_n_idx + i * 32) < N_padded);
339 + }
340 + }
341 + }
342 +
343 + __device__ void mma(const int& reg_buf_idx) {
344 + #pragma unroll
345 + for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
346 + #pragma unroll
347 + for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
348 + hmma16816_f32<FType>(
349 + C_frag[m_idx][n_idx], A_frag[reg_buf_idx][m_idx],
350 + reinterpret_cast<uint32_t(&)[2]>(BF_frag[reg_buf_idx][n_idx]));
351 + }
352 + }
353 + }
354 +
355 + __device__ void fused_splitk_reduce() {
356 + // need splitk-reduce if enable splitk
357 + if (gridDim.z > 1) {
358 + int blk_red_idx = blockIdx.x * gridDim.y + blockIdx.y;
359 + // Wait for all previous blocks in the splitk direction to accumulate the
360 + // results into C_tmp
361 + if (threadIdx.x == 0) {
362 + uint32_t* red_count_ptr = params.red_count_ptr + blk_red_idx;
363 + uint32_t count;
364 + do {
365 + // make sure the ld.cg inside the do-wile loop
366 + __threadfence_block();
367 + asm volatile("ld.global.cg.b32 %0, [%1];"
368 + : "=r"(count)
369 + : "l"(red_count_ptr));
370 + } while (count != blockIdx.z);
371 + }
372 + __syncthreads();
373 +
374 + int C_tmp_base_offset = blk_red_idx * Mtile * Ntile + threadIdx.x * 4;
375 + if (blockIdx.z != 0) {
376 + // expecting that temporary register here reuses the previous A&B frag
377 + // register
378 + float temp_frag[Mtile / 16][WARP_NITER][4];
379 + #pragma unroll
380 + for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
381 + #pragma unroll
382 + for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
383 + int offset =
384 + C_tmp_base_offset + (m_idx * WARP_NITER + n_idx) * BLOCK * 4;
385 + *reinterpret_cast<int4*>(temp_frag[m_idx][n_idx]) =
386 + *reinterpret_cast<int4*>(params.C_tmp_ptr + offset);
387 + }
388 + }
389 + #pragma unroll
390 + for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
391 + #pragma unroll
392 + for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
393 + #pragma unroll
394 + for (int idx = 0; idx < 4; ++idx) {
395 + C_frag[m_idx][n_idx][idx] += temp_frag[m_idx][n_idx][idx];
396 + }
397 + }
398 + }
399 + }
400 +
401 + // first splitk - 1 blocks need to write partial results into C_tmp
402 + if (blockIdx.z != gridDim.z - 1) {
403 + #pragma unroll
404 + for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
405 + #pragma unroll
406 + for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
407 + int offset =
408 + C_tmp_base_offset + (m_idx * WARP_NITER + n_idx) * BLOCK * 4;
409 + asm volatile(
410 + "{st.global.cg.v4.b32 [%0], {%1, %2, %3, %4};}\n"
411 + :
412 + : "l"(params.C_tmp_ptr + offset), "f"(C_frag[m_idx][n_idx][0]),
413 + "f"(C_frag[m_idx][n_idx][1]), "f"(C_frag[m_idx][n_idx][2]),
414 + "f"(C_frag[m_idx][n_idx][3]));
415 + }
416 + }
417 + __threadfence();
418 + __syncthreads();
419 + if (threadIdx.x == 0) {
420 + uint32_t* red_count_ptr = params.red_count_ptr + blk_red_idx;
421 + atomicInc(red_count_ptr, gridDim.z);
422 + }
423 + }
424 + }
425 + }
426 +
427 + __device__ void stg(char* smem) {
428 + if (EnableFuse) {
429 + if (blockIdx.z != gridDim.z - 1) return;
430 + }
431 + uint32_t* C_sts_ptr =
432 + reinterpret_cast<uint32_t*>(smem + sts_c_base_offset * sizeof(FType));
433 + // C_tile sts
434 + #pragma unroll
435 + for (int m_idx = 0; m_idx < Mtile / 16; ++m_idx) {
436 + #pragma unroll
437 + for (int n_idx = 0; n_idx < WARP_NITER; ++n_idx) {
438 + #pragma unroll
439 + for (int k_idx = 0; k_idx < 2; ++k_idx) {
440 + FType low16 = static_cast<FType>(C_frag[m_idx][n_idx][k_idx * 2]);
441 + FType high16 =
442 + static_cast<FType>(C_frag[m_idx][n_idx][k_idx * 2 + 1]);
443 + uint32_t tmp = (reinterpret_cast<uint32_t&>(low16) & 0xffff) |
444 + (reinterpret_cast<uint32_t&>(high16) << 16);
445 + int sts_offset =
446 + m_idx * 16 * (WARP_NTILE / 2) +
447 + (((lane_id / (32 / WARP_NITER)) + n_idx) % WARP_NITER) * (8 / 2) +
448 + k_idx * 8 * (WARP_NTILE / 2);
449 + C_sts_ptr[sts_offset] = tmp;
450 + }
451 + }
452 + }
453 +
454 + __syncthreads();
455 +
456 + FType* C_base_ptr = this_block_C_base_ptr + store_c_base_offset;
457 + // C_tile lds and stg
458 + int m_base_idx = store_c_row_base_idx + blockIdx.x * Mtile;
459 + bool n_guard = (store_c_col_idx + blockIdx.y * Ntile) < params.N;
460 + if (WARP_NTILE == 32) {
461 + int lds_c_base_offset = warp_id * Mtile * WARP_NTILE +
462 + (lane_id / 4) * WARP_NTILE +
463 + ((lane_id % 4 + lane_id / 8) % 4) * 8;
464 + uint4* C_lds_ptr =
465 + reinterpret_cast<uint4*>(smem + lds_c_base_offset * sizeof(FType));
466 + #pragma unroll
467 + for (int i = 0; i < (Mtile / 16) * (WARP_NITER / 2); ++i) {
468 + uint4 stg_reg = C_lds_ptr[i * 8 * 4];
469 + stg128(stg_reg.x, stg_reg.y, stg_reg.z, stg_reg.w,
470 + C_base_ptr + i * 8 * params.N,
471 + (m_base_idx + i * 8) < params.M && n_guard);
472 + }
473 + } else if (WARP_NTILE == 64) {
474 + int lds_c_base_offset =
475 + warp_id * Mtile * WARP_NTILE + (lane_id / 8) * WARP_NTILE;
476 + #pragma unroll
477 + for (int i = 0; i < (Mtile / 16) * (WARP_NITER / 2); ++i) {
478 + int lds_c_offset = lds_c_base_offset + i * 4 * WARP_NTILE +
479 + ((lane_id % 8 + lane_id / 8 + (i % 2) * 4) % 8) * 8;
480 + uint4 stg_reg =
481 + *reinterpret_cast<uint4*>(smem + lds_c_offset * sizeof(FType));
482 + stg128(stg_reg.x, stg_reg.y, stg_reg.z, stg_reg.w,
483 + C_base_ptr + i * 4 * params.N,
484 + (m_base_idx + i * 4) < params.M && n_guard);
485 + }
486 + }
487 + }
488 +
489 + const SM8x_GEMM_W8A16_Splitk_Params<FType, QType>& params;
490 +
491 + int load_a_base_offset[2];
492 + int load_b_base_offset[2];
493 + int sts_c_base_offset;
494 +
495 + int store_c_base_offset;
496 +
497 + int store_c_row_base_idx, store_c_col_idx;
498 + FType* this_block_C_base_ptr = nullptr;
499 +
500 + int params_n_idx;
501 + const uint32_t A_smem_base_addr, BQ_smem_base_addr;
502 + const uint32_t A_smem_stage_stride, BQ_smem_stage_stride;
503 +
504 + int lane_id;
505 + int warp_id;
506 + // first 2 denotes double buffer, second dim denotes M direction
507 + uint32_t A_frag[2][Mtile / 16][4];
508 +
509 + typename HalfType<FType>::T2 B_scale[WARP_NITER / 2];
510 + typename HalfType<FType>::T2 B_zero[WARP_NITER / 2];
511 + uint32_t BQ_frag[2][WARP_NITER];
512 + // first 2 denotes double buffer, second dim denotes N direction, last 2
513 + // denotes K direction
514 + typename HalfType<FType>::T2 BF_frag[2][WARP_NITER][2];
515 + // first dim denotes M direction, second dim denotes N direction
516 + float C_frag[Mtile / 16][WARP_NITER][4];
517 +};
518 +
519 +/*
520 + * @brief W8A16 Perchannel Quantization GEMM,
521 + * requires N % 8 == 0, K % 16 == 0
522 + * accumulator precision: FP32
523 + * @tparam FType: DataType for A, B_scale, B_zero, and C, supports half or
524 + * nv_bfloat16
525 + * @tparam QType: DataType for B, support uint8(bias128)
526 + * @tparam Mtile: M-dimensional size of the gemm block tile, supports 16, 32,
527 + * 48 or 64
528 + * @tparam Ntile: N-dimensional size of the gemm block tile, supports 128 or
529 + * 256
530 + * @tparam NStage: Num of stages for async copy
531 + * @tparam BLOCK: BLOCK size
532 + * @tparam EnableFuse: If true, use fused splitk-reduce, otherwise use
533 + * non-fused splitk-reduce
534 + * @tparam has_zp: whether to use zero_point
535 + *
536 + * @fparam params struct consists of following parameters:
537 + * @param A_ptr: Matrix A value ptr, A = (M, K)
538 + * @param B_ptr: Matrix B value ptr, B = (N32_align, K) (N32K16 special
539 + * format), N32_align = (N + 32 - 1) / 32 * 32
540 + * @param B_scale_ptr: B_scale value ptr, B_scale = (N32_align,) (N32K16
541 + * special format)
542 + * @param B_zero_ptr: B_zero value ptr, B_zero = (N32_align,) (N32K16
543 + * special format)
544 + * @param C_ptr: Matrix C value ptr, C = (M, N)
545 + * @param M: dimnesion m
546 + * @param N: dimnesion n
547 + * @param K: dimnesion k
548 + * @param SplitK: split size along K-dimension
549 + * @param C_split_ptr: Matrix C_split value ptr, used only in non-fused
550 + * splitk-reduce
551 + * @param C_tmp_ptr: Matrix C_tmp value ptr, used only in fused
552 + * splitk-reduce
553 + * @param red_count_ptr: 1-D red_count value ptr, used only in fused
554 + * splitk-reduce
555 + */
556 +template <typename FType, typename QType, int Mtile, int Ntile, int NStage,
557 + int BLOCK, bool EnableFuse, bool has_zp>
558 +__global__ void __launch_bounds__(BLOCK)
559 + ampere_hgemm_W8A16_perc_f16_f16_MtilexNtilex32_hmma16816_multistage_AN_BTN32K16_CN_splitk_kernel(
560 + const SM8x_GEMM_W8A16_Splitk_Params<FType, QType> params) {
561 + // A smem size = 64 * 32 * 2B/elem * 4(stage) = 16KB
562 + // B smem size = 128 * 32 * 1B/elem * 4(stage) = 16KB
563 + constexpr int smem_size_one_stage = Mtile * 32 * 2 + Ntile * 32;
564 + __shared__ char smem[NStage * smem_size_one_stage];
565 + char* A_smem = smem;
566 + char* BQ_smem = smem + Mtile * 32 * 2 * NStage;
567 +
568 + uint32_t A_smem_addr = smem_u32addr(A_smem);
569 + uint32_t BQ_smem_addr = smem_u32addr(BQ_smem);
570 + uint32_t A_smem_stage_stride = Mtile * 32 * 2;
571 + uint32_t BQ_smem_stage_stride = Ntile * 32;
572 +
573 + // initialize the data move process from GM to SMEM for this block
574 + GmemTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK<
575 + FType, QType, Mtile, Ntile, NStage, BLOCK>
576 + gmem_tile(params, A_smem_addr, BQ_smem_addr, A_smem_stage_stride,
577 + BQ_smem_stage_stride);
578 +
579 + int sts_stage_idx = 0;
580 + int lds_stage_idx = 0;
581 +
582 + int tb_k_slice = blockIdx.z * params.SplitK + params.SplitK <= params.K
583 + ? params.SplitK
584 + : params.K - blockIdx.z * params.SplitK;
585 + int k_tiles = (tb_k_slice + 31) / 32;
586 + int first_k_tile = tb_k_slice - (k_tiles - 1) * 32;
587 +
588 + // load first three tiles to shared memory
589 + gmem_tile.ldgsts_first_ktiles(first_k_tile, k_tiles);
590 + sts_stage_idx += (NStage - 2);
591 + ComputeTile_W8A16_PerC_MtilexNtilex32_multistage_SM8x_SplitK<
592 + FType, QType, Mtile, Ntile, BLOCK, EnableFuse, has_zp>
593 + compute_tile(params, A_smem_addr, BQ_smem_addr, A_smem_stage_stride,
594 + BQ_smem_stage_stride);
595 + compute_tile.ldg_params();
596 + cp_asyc_wait_group<NStage - 2>();
597 + __syncthreads();
598 +
599 + compute_tile.lds(lds_stage_idx, 0, 0);
600 + int reg_buf_idx = 1;
601 +
602 + // main loop
603 + for (; k_tiles > NStage - 1; --k_tiles) {
604 + // load next A&B tile
605 + sts_stage_idx = sts_stage_idx < NStage - 1 ? sts_stage_idx + 1 : 0;
606 + gmem_tile.ldgsts(sts_stage_idx);
607 +
608 + #pragma unroll
609 + for (int k_phase_idx = 0; k_phase_idx < 2; k_phase_idx++) {
610 + // dequantize next B tile
611 + if (k_phase_idx == 1) {
612 + cp_asyc_wait_group<NStage - 2>();
613 + __syncthreads();
614 + lds_stage_idx = lds_stage_idx < NStage - 1 ? lds_stage_idx + 1 : 0;
615 + }
616 +
617 + compute_tile.lds(lds_stage_idx, reg_buf_idx, (k_phase_idx + 1) % 2);
618 +
619 + compute_tile.mma(reg_buf_idx ^ 1);
620 + reg_buf_idx ^= 1;
621 + }
622 + }
623 +
624 + // last NStage-1 tiles
625 + for (; k_tiles > 0; --k_tiles) {
626 + cp_async_commit_group();
627 + #pragma unroll
628 + for (int k_phase_idx = 0; k_phase_idx < 2; k_phase_idx++) {
629 + // dequantize next B tile
630 + if (k_phase_idx == 1) {
631 + cp_asyc_wait_group<NStage - 2>();
632 + __syncthreads();
633 + lds_stage_idx = lds_stage_idx < NStage - 1 ? lds_stage_idx + 1 : 0;
634 + }
635 +
636 + compute_tile.lds(lds_stage_idx, reg_buf_idx, (k_phase_idx + 1) % 2);
637 +
638 + compute_tile.mma(reg_buf_idx ^ 1);
639 + reg_buf_idx ^= 1;
640 + }
641 + }
642 +
643 + if (EnableFuse) {
644 + compute_tile.fused_splitk_reduce();
645 + }
646 + compute_tile.stg(smem);
647 +}
648 +
649 + #define __CALL_IF(MTILE, NTILE, NUM_THREADS, ENABLE_FUSE, HAS_ZP) \
650 + else if (Mtile == MTILE && Ntile == NTILE && BLOCK == NUM_THREADS && \
651 + enable_fuse == ENABLE_FUSE && has_zp == HAS_ZP) { \
652 + ampere_hgemm_W8A16_perc_f16_f16_MtilexNtilex32_hmma16816_multistage_AN_BTN32K16_CN_splitk_kernel< \
653 + FType, QType, MTILE, NTILE, 4, NUM_THREADS, ENABLE_FUSE, HAS_ZP> \
654 + <<<grid, block, 0, stream>>>(params); \
655 + }
656 +
657 +template <typename FType, typename QType>
658 +void ampere_hgemm_W8A16_perc_f16_f16_MtilexNtilex32_mma16816_multistage_AN_BTN32K16_CN_splitk(
659 + const FType* A, const QType* B, const FType* B_scale, const FType* B_zero,
660 + FType* C, const int M, const int N, const int K, void* workspace,
661 + const int sm_version, const BlockTileSplitkParams& fused_gemm_params,
662 + cudaStream_t stream) {
663 + int Mtile = fused_gemm_params.Mtile;
664 + int grid_x = (M + Mtile - 1) / Mtile;
665 + int Ntile = fused_gemm_params.Ntile;
666 + int grid_y = (N + Ntile - 1) / Ntile;
667 + int SplitK = fused_gemm_params.SplitK;
668 + int grid_z = (K + SplitK - 1) / SplitK;
669 +
670 + int BLOCK = (Ntile == 256) ? 256 : 128;
671 +
672 + dim3 grid(grid_x, grid_y, grid_z);
673 + dim3 block(BLOCK);
674 +
675 + bool enable_fuse = fused_gemm_params.EnableFuse;
676 + bool has_zp = B_zero != nullptr;
677 + if (enable_fuse) {
678 + float* C_tmp = reinterpret_cast<float*>(workspace);
679 + uint32_t* red_count = reinterpret_cast<uint32_t*>(
680 + (char*)workspace + grid_x * Mtile * grid_y * Ntile * sizeof(float));
681 + CHECK_CUDA(cudaMemsetAsync(red_count, 0, grid_x * grid_y * sizeof(uint32_t),
682 + stream));
683 + SM8x_GEMM_W8A16_Splitk_Params<FType, QType> params{
684 + A, B, B_scale, B_zero, C, M, N,
685 + K, SplitK, 0, -1, nullptr, C_tmp, red_count};
686 +
687 + if (false) {
688 + }
689 + // Select the template parameters for kernel launch
690 + // according to the above settings. Tuning is not supported.
691 + __CALL_IF(16, 256, 256, true, false)
692 + __CALL_IF(32, 256, 256, true, false)
693 + __CALL_IF(48, 256, 256, true, false)
694 + __CALL_IF(64, 128, 128, true, false)
695 + __CALL_IF(64, 256, 256, true, false)
696 + __CALL_IF(16, 256, 256, true, true)
697 + __CALL_IF(32, 256, 256, true, true)
698 + __CALL_IF(48, 256, 256, true, true)
699 + __CALL_IF(64, 128, 128, true, true)
700 + __CALL_IF(64, 256, 256, true, true)
701 + } else {
702 + FType* C_split = reinterpret_cast<FType*>(workspace);
703 + SM8x_GEMM_W8A16_Splitk_Params<FType, QType> params{
704 + A, B, B_scale, B_zero, C, M, N,
705 + K, SplitK, 0, -1, C_split, nullptr, nullptr};
706 +
707 + if (false) {
708 + }
709 + // Select the template parameters for kernel launch
710 + // according to the above settings. Tuning is not supported.
711 + __CALL_IF(16, 256, 256, false, false)
712 + __CALL_IF(32, 256, 256, false, false)
713 + __CALL_IF(48, 256, 256, false, false)
714 + __CALL_IF(64, 128, 128, false, false)
715 + __CALL_IF(64, 256, 256, false, false)
716 + __CALL_IF(16, 256, 256, false, true)
717 + __CALL_IF(32, 256, 256, false, true)
718 + __CALL_IF(48, 256, 256, false, true)
719 + __CALL_IF(64, 128, 128, false, true)
720 + __CALL_IF(64, 256, 256, false, true)
721 +
722 + // SplitK reduce
723 + f16_gemm_splitk_reduce(C_split, C, M, N, grid_z, stream);
724 + }
725 +}
726 +
727 +size_t allspark_qgemm_w8a16_perc_n32k16_ampere_workspace_size(
728 + int m, int n, int k, int sm_count,
729 + BlockTileSplitkParams& fused_gemm_params) {
730 + // Determine the block tile and splitk strategy
731 + int m16_times = (m + 16 - 1) / 16;
732 + int Mtile = m16_times <= 4 ? m16_times * 16 : 64;
733 + int grid_x = (m + Mtile - 1) / Mtile;
734 + int Ntile =
735 + (float(grid_x * ((n + 127) / 128)) / sm_count > 10) || (Mtile < 64) ? 256
736 + : 128;
737 + int grid_y = (n + Ntile - 1) / Ntile;
738 + int grid_z;
739 +
740 + // split-k
741 + const float SPLIT_THRESHOLD = 0.8;
742 + int n_slice;
743 + for (n_slice = 1; n_slice < k / 256; ++n_slice) {
744 + int n_block = grid_x * grid_y * n_slice;
745 + if (n_block >= sm_count * SPLIT_THRESHOLD &&
746 + (n_block % sm_count == 0 || n_block % sm_count >= sm_count * 0.5)) {
747 + break;
748 + }
749 + }
750 +
751 + int k_slice =
752 + (k / n_slice) % 32 == 0 ? k / n_slice : k / n_slice / 32 * 32 + 32;
753 + grid_z = (k + k_slice - 1) / k_slice;
754 + bool enable_fuse = float(grid_x * grid_y) / sm_count >= 0.5 ? 1 : 0;
755 +
756 + size_t ws_size;
757 + if (enable_fuse) {
758 + ws_size = grid_x * Mtile * grid_y * Ntile * sizeof(float) // For C_tmp
759 + + grid_x * grid_y * sizeof(uint32_t); // For red_count
760 + } else {
761 + ws_size = grid_z * m * n * sizeof(__half);
762 + }
763 +
764 + fused_gemm_params.Mtile = Mtile;
765 + fused_gemm_params.Ntile = Ntile;
766 + fused_gemm_params.SplitK = k_slice;
767 + fused_gemm_params.EnableFuse = enable_fuse;
768 + return ws_size;
769 +}
770 +
771 +// restore from N32K16 order to original N-major order
772 +// K % 16 == 0, N % 8 == 0
773 +// each block process 64(k) * 32(n) result elements
774 +template <typename FT, typename QT>
775 +__global__ void restore_N32_K16_dequantize_rhs_w8a16_perc_kernel(
776 + const QT* qdata, const FT* scales, const FT* zeros, FT* fdata,
777 + const int N_32align, const int N, const int K) {
778 + __shared__ FT smem[64 * 32];
779 + int warp_id = threadIdx.x / 32;
780 + int lane_id = threadIdx.x % 32;
781 + const int src_row_idx = blockIdx.x * 8 + lane_id / 4;
782 + const int src_col_idx =
783 + blockIdx.y * 64 * 4 + warp_id * 16 * 4 + (lane_id % 4) * 16;
784 + const int src_offset = src_row_idx * K * 4 + src_col_idx;
785 + int params_nidx = blockIdx.x * 32 + (lane_id / 4) * 4;
786 +
787 + QT qval_reg[16];
788 + const QT* pdata = qdata + src_offset;
789 + if (src_col_idx < (K * 4)) {
790 + *(reinterpret_cast<uint4*>(qval_reg)) =
791 + *(reinterpret_cast<const uint4*>(qdata + src_offset));
792 + }
793 + FT scale_reg[4];
794 + *(reinterpret_cast<uint2*>(scale_reg)) =
795 + *(reinterpret_cast<const uint2*>(scales + params_nidx));
796 + FT zero_reg[4] = {0};
797 + if (zeros != nullptr) {
798 + *(reinterpret_cast<uint2*>(zero_reg)) =
799 + *(reinterpret_cast<const uint2*>(zeros + params_nidx));
800 + }
801 + FT fval_reg[16];
802 +
803 + const int sts_base_offset =
804 + (warp_id * 16 + (lane_id % 4) * 2) * 32 + lane_id / 4;
805 + #pragma unroll
806 + for (int ni = 0; ni < 4; ++ni) {
807 + cvt_8bx4_to_16bx4_bias128(
808 + *reinterpret_cast<uint32_t*>(&qval_reg[ni * 4]),
809 + reinterpret_cast<typename HalfType<FT>::T2*>(&(fval_reg[ni * 4])));
810 + #pragma unroll
811 + for (int ki = 0; ki < 4; ++ki) {
812 + fval_reg[ni * 4 + ki] =
813 + (fval_reg[ni * 4 + ki] - zero_reg[ni]) * scale_reg[ni];
814 + int sts_offset = sts_base_offset + ((ki / 2) * 8 + (ki % 2)) * 32 +
815 + ((ni + lane_id % 4) % 4) * 8;
816 + smem[sts_offset] = fval_reg[ni * 4 + ki];
817 + }
818 + }
819 + __syncthreads();
820 +
821 + const int lds_base_offset =
822 + (threadIdx.x / 4) * 32 + ((threadIdx.x % 4 + threadIdx.x / 8) % 4) * 8;
823 + #pragma unroll
824 + for (int i = 0; i < 2; ++i) {
825 + *reinterpret_cast<uint4*>(fval_reg + i * 8) =
826 + *reinterpret_cast<uint4*>(smem + lds_base_offset + i * 32 * 32);
827 + }
828 +
829 + const int dst_row_base_kidx = blockIdx.y * 64 + threadIdx.x / 4;
830 + const int dst_col_nidx = blockIdx.x * 32 + (threadIdx.x % 4) * 8;
831 + #pragma unroll
832 + for (int i = 0; i < 2; ++i) {
833 + int dst_row_kidx = dst_row_base_kidx + i * 32;
834 + int dst_offset = dst_row_kidx * N + dst_col_nidx;
835 + if (dst_row_kidx < K && dst_col_nidx < N) {
836 + *reinterpret_cast<uint4*>(fdata + dst_offset) =
837 + *reinterpret_cast<uint4*>(fval_reg + i * 8);
838 + }
839 + }
840 +}
841 +
842 +template <typename FT, typename QT>
843 +void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
844 + const FT* zeros, FT* fdata,
845 + const int N_32align, const int N,
846 + const int K, const int GroupSize,
847 + cudaStream_t stream) {
848 + TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
849 + "Unsupported shape");
850 + if (GroupSize == -1) {
851 + const int BLOCK = 128;
852 + dim3 grid(N_32align / 32, ((K / 16) + 3) / 4);
853 + restore_N32_K16_dequantize_rhs_w8a16_perc_kernel<FT, QT>
854 + <<<grid, BLOCK, 0, stream>>>(qdata, scales, zeros, fdata, N_32align, N,
855 + K);
856 + }
857 + // TODO: Support SubChannel
858 + else {
859 + TORCH_CHECK(false, "Now only support PerChannel");
860 + }
861 +}
862 +
863 +template <typename FT, typename QT>
864 +void w8a16_gemm_dq_cublas(const FT* in, const QT* rhs_qdata_ptr,
865 + const FT* rhs_scales_ptr, const FT* rhs_zeros_ptr,
866 + FT* out, void* workspace, const int M,
867 + const int N_32align, const int N, const int K,
868 + const int group_size, cudaStream_t stream,
869 + cublasHandle_t handle) {
870 + static_assert(
871 + std::is_same<FT, half>::value || std::is_same<FT, nv_bfloat16>::value,
872 + "only float16 and bfloat16 is supported");
873 + // Dequant
874 + FT* rhs_fdata_ptr = static_cast<FT*>(workspace);
875 + restore_N32_K16_dequantize_rhs_w8a16(rhs_qdata_ptr, rhs_scales_ptr,
876 + rhs_zeros_ptr, rhs_fdata_ptr, N_32align,
877 + N, K, group_size, stream);
878 + // cuBLAS GEMM
879 + int lda = K;
880 + int ldb = N;
881 + int ldc = N;
882 + const float alpha = 1.0f;
883 + const float beta = 0.0f;
884 + cudaDataType_t cuda_type;
885 + if (std::is_same<FT, __half>::value) {
886 + cuda_type = CUDA_R_16F;
887 + } else {
888 + cuda_type = CUDA_R_16BF;
889 + }
890 + CHECK_CUBLAS(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha,
891 + rhs_fdata_ptr, cuda_type, ldb, in, cuda_type, lda,
892 + &beta, out, cuda_type, ldc, CUDA_R_32F,
893 + CUBLAS_GEMM_DEFAULT_TENSOR_OP));
894 +}
895 +
896 +template <typename FType, typename QType>
897 +void allspark_qgemm_w8a16_perc_ampere(
898 + const FType* A, const QType* B, const FType* B_scale, const FType* B_zero,
899 + FType* C, const int M, const int N_32align, const int N, const int K,
900 + void* workspace, const BlockTileSplitkParams& fused_gemm_params,
901 + const int group_size, int CUBLAS_M_THRESHOLD, const int sm_version,
902 + cudaStream_t stream, cublasHandle_t handle) {
903 + if (M > CUBLAS_M_THRESHOLD) {
904 + w8a16_gemm_dq_cublas<FType, QType>(A, B, B_scale, B_zero, C, workspace, M,
905 + N_32align, N, K, group_size, stream,
906 + handle);
907 + } else {
908 + ampere_hgemm_W8A16_perc_f16_f16_MtilexNtilex32_mma16816_multistage_AN_BTN32K16_CN_splitk<
909 + FType, QType>(A, B, B_scale, B_zero, C, M, N, K, workspace, sm_version,
910 + fused_gemm_params, stream);
911 + }
912 +}
913 +
914 +} // namespace allspark
915 +
916 +torch::Tensor allspark_w8a16_gemm(
917 + torch::Tensor const& a, torch::Tensor const& b_qweight,
918 + torch::Tensor const& b_scales, c10::optional<torch::Tensor> const& b_qzeros,
919 + int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
920 + int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
921 + // Verify device and strides
922 + TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
923 + TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
924 +
925 + TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
926 + TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
927 +
928 + TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
929 + TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
930 +
931 + if (has_zp) {
932 + TORCH_CHECK(b_qzeros.value().device().is_cuda(), "b_qzeros is not on GPU");
933 + TORCH_CHECK(b_qzeros.value().is_contiguous(), "b_qzeros is not contiguous");
934 + }
935 +
936 + int m = a.size(0);
937 + int n_32align = (n + 32 - 1) / 32 * 32;
938 + int k = a.size(1);
939 +
940 + // Verify shape
941 + TORCH_CHECK(b_qweight.size(0) == n_32align,
942 + "Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
943 + ", n_32align = ", n_32align);
944 + TORCH_CHECK(b_qweight.size(1) == k,
945 + "Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
946 + ", k = ", k);
947 +
948 + TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
949 +
950 + const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
951 + const void* a_ptr = reinterpret_cast<const void*>(a.data_ptr());
952 + const uint8_t* b_ptr = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
953 + const void* b_scale_ptr = reinterpret_cast<const void*>(b_scales.data_ptr());
954 + const void* b_zero_ptr = nullptr;
955 + if (b_qzeros.has_value()) {
956 + b_zero_ptr = reinterpret_cast<const void*>(b_qzeros.value().data_ptr());
957 + }
958 +
959 + auto c_options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
960 + torch::Tensor c = torch::empty({m, n}, c_options);
961 + void* c_ptr = reinterpret_cast<void*>(c.data_ptr());
962 +
963 + cudaStream_t stream = at::cuda::getCurrentCUDAStream();
964 + cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
965 +
966 + allspark::BlockTileSplitkParams fused_gemm_params;
967 +
968 + size_t ws_size = 0;
969 + if (m > CUBLAS_M_THRESHOLD) {
970 + ws_size = k * n * 2; // sizeof(f16)==2
971 + } else {
972 + ws_size = allspark::allspark_qgemm_w8a16_perc_n32k16_ampere_workspace_size(
973 + m, n, k, sm_count, fused_gemm_params);
974 + }
975 +
976 + auto ws_options = torch::TensorOptions().dtype(at::kChar).device(a.device());
977 + if (as_g_workspace.numel() <
978 + ws_size) { // ws_options: kChar, so numel() is bytes
979 + as_g_workspace = torch::empty({long(ws_size)}, ws_options);
980 + }
981 + void* ws = reinterpret_cast<void*>(as_g_workspace.data_ptr());
982 +
983 + if (a.dtype() == at::ScalarType::Half) {
984 + allspark::allspark_qgemm_w8a16_perc_ampere<__half, uint8_t>(
985 + reinterpret_cast<const __half*>(a_ptr), b_ptr,
986 + reinterpret_cast<const __half*>(b_scale_ptr),
987 + reinterpret_cast<const __half*>(b_zero_ptr),
988 + reinterpret_cast<__half*>(c_ptr), m, n_32align, n, k, ws,
989 + fused_gemm_params, group_size, CUBLAS_M_THRESHOLD, sm_version, stream,
990 + handle);
991 + } else if (a.dtype() == at::ScalarType::BFloat16) {
992 + allspark::allspark_qgemm_w8a16_perc_ampere<__nv_bfloat16, uint8_t>(
993 + reinterpret_cast<const __nv_bfloat16*>(a_ptr), b_ptr,
994 + reinterpret_cast<const __nv_bfloat16*>(b_scale_ptr),
995 + reinterpret_cast<const __nv_bfloat16*>(b_zero_ptr),
996 + reinterpret_cast<__nv_bfloat16*>(c_ptr), m, n_32align, n, k, ws,
997 + fused_gemm_params, group_size, CUBLAS_M_THRESHOLD, sm_version, stream,
998 + handle);
999 + }
1000 +
1001 + return c;
1002 +}
1003 +
1004 +#endif
1005 +
1006 +TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
1007 + m.impl("allspark_w8a16_gemm", &allspark_w8a16_gemm);
1008 +}
\ No newline at end of file
csrc/quantization/gptq_allspark/allspark_repack.cu new
+163
@@ -0,0 +1,163 @@
1 +#include "allspark_utils.cuh"
2 +#include <torch/all.h>
3 +#include "core/registration.h"
4 +
5 +namespace allspark {
6 +
7 +// Rearrange B to facilitate Ampere Tensor Core load data
8 +// reorder B from (K, N) to (N_32align / 4, K * 4)
9 +// K % 16 == 0, N % 16 == 0, N_32align % 32 == 0
10 +template <typename FType>
11 +__global__ void __launch_bounds__(128)
12 + rearrange_kn_weight_as_n32k16_order_ldg16_kernel(
13 + const uint8_t* B, const FType* B_scale, const FType* B_zero,
14 + uint8_t* B_result, FType* B_scale_result, FType* B_zero_result,
15 + const int K, const int N, const int N_32align) {
16 + const int lane_id = threadIdx.x % 32;
17 + const int warp_id = threadIdx.x / 32;
18 +
19 + if (blockIdx.x != gridDim.x - 1) {
20 + // Load B
21 + // per block process 64(k) * 128(n) B elements
22 + // per warp process 16(k) * 128 B elements
23 + const int src_row_base_idx =
24 + blockIdx.x * 64 + warp_id * 16 + ((lane_id % 8) / 2) * 2;
25 + const int src_col_idx =
26 + blockIdx.y * 128 + (lane_id / 8) * 32 + (lane_id % 2) * 16;
27 + uint8_t B_frag[4][16];
28 +#pragma unroll
29 + for (int i = 0; i < 4; ++i) {
30 + int src_row_idx = src_row_base_idx + (i / 2) * 8 + (i % 2);
31 + int src_offset = src_row_idx * N + src_col_idx;
32 + bool guard = src_row_idx < K && src_col_idx < N;
33 + ldg128_cg_0(*reinterpret_cast<uint32_t*>(B_frag[i]),
34 + *(reinterpret_cast<uint32_t*>(B_frag[i]) + 1),
35 + *(reinterpret_cast<uint32_t*>(B_frag[i]) + 2),
36 + *(reinterpret_cast<uint32_t*>(B_frag[i]) + 3), B + src_offset,
37 + guard);
38 + }
39 +
40 + // reorder B
41 + uint8_t B_reorder_frag[8][8];
42 +#pragma unroll
43 + for (int i = 0; i < 4; ++i) {
44 +#pragma unroll
45 + for (int j = 0; j < 16; ++j) {
46 + int dst_i = j % 8;
47 + int dst_j = i + (j / 8) * 4;
48 + B_reorder_frag[dst_i][dst_j] = B_frag[i][j];
49 + }
50 + }
51 +
52 + // Store B
53 + const int dst_row_base_idx = blockIdx.y * (128 / 4) + (lane_id / 8) * 8;
54 + const int dst_col_idx =
55 + blockIdx.x * (64 * 4) + warp_id * 64 + (lane_id % 8) * 8;
56 + for (int i = 0; i < 8; ++i) {
57 + int dst_row_idx = dst_row_base_idx + i;
58 + int dst_offset = dst_row_idx * K * 4 + dst_col_idx;
59 + bool guard = (dst_row_base_idx < N_32align / 4) && (dst_col_idx < K * 4);
60 + if (guard) {
61 + *reinterpret_cast<int2*>(B_result + dst_offset) =
62 + *reinterpret_cast<int2*>(B_reorder_frag[i]);
63 + }
64 + }
65 + } else {
66 + // Load B_scale and B_zero
67 + FType b_scale_reg, b_zero_reg;
68 + int src_offset = blockIdx.y * 128 + threadIdx.x;
69 + ldg16_cg_0(b_scale_reg, B_scale + src_offset, src_offset < N);
70 + if (B_zero != nullptr)
71 + ldg16_cg_0(b_zero_reg, B_zero + src_offset, src_offset < N);
72 + int dst_offset =
73 + blockIdx.y * 128 + warp_id * 32 + (lane_id % 8) * 4 + lane_id / 8;
74 + if (dst_offset < N_32align) {
75 + B_scale_result[dst_offset] = b_scale_reg;
76 + if (B_zero != nullptr) B_zero_result[dst_offset] = b_zero_reg;
77 + }
78 + }
79 +}
80 +
81 +template <typename FType>
82 +void rearrange_kn_weight_as_n32k16_order_ldg16(
83 + const uint8_t* B, const FType* B_scale, const FType* B_zero,
84 + uint8_t* B_result, FType* B_scale_result, FType* B_zero_result,
85 + const int64_t K, const int64_t N, const int64_t N_32align,
86 + cudaStream_t stream) {
87 + if (N % 16 != 0 || K % 16 != 0) {
88 + std::cerr << "Now only support N and K is multiples of 16" << std::endl;
89 + }
90 + const int BLOCK = 128;
91 + int grid_x = (K + 64 - 1) / 64 + 1;
92 + int grid_y = (N + 128 - 1) / 128;
93 + dim3 grid(grid_x, grid_y);
94 +
95 + rearrange_kn_weight_as_n32k16_order_ldg16_kernel<FType>
96 + <<<grid, BLOCK, 0, stream>>>(B, B_scale, B_zero, B_result, B_scale_result,
97 + B_zero_result, K, N, N_32align);
98 +}
99 +} // namespace allspark
100 +
101 +void rearrange_kn_weight_as_n32k16_order(
102 + torch::Tensor const& b_qweight, torch::Tensor const& b_scales,
103 + c10::optional<torch::Tensor> const& b_zeros, bool has_zp,
104 + torch::Tensor& b_qweight_reorder, torch::Tensor& b_scales_reorder,
105 + c10::optional<torch::Tensor> const& b_zeros_reorder, const int64_t K,
106 + const int64_t N, const int64_t N_32align) {
107 + // Verify device and strides
108 + TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
109 + TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
110 +
111 + TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
112 + TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
113 +
114 + TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
115 + "b_qweight_reorder is not on GPU");
116 + TORCH_CHECK(b_qweight_reorder.is_contiguous(),
117 + "b_qweight_reorder is not contiguous");
118 +
119 + TORCH_CHECK(b_scales_reorder.device().is_cuda(),
120 + "b_scales_reorder is not on GPU");
121 + TORCH_CHECK(b_scales_reorder.is_contiguous(),
122 + "b_scales_reorder is not contiguous");
123 +
124 + if (has_zp) {
125 + TORCH_CHECK(b_zeros.value().device().is_cuda(), "b_zeros is not on GPU");
126 + TORCH_CHECK(b_zeros.value().is_contiguous(), "b_zeros is not contiguous");
127 +
128 + TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
129 + "b_zeros_reorder is not on GPU");
130 + TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
131 + "b_zeros_reorder is not contiguous");
132 + }
133 +
134 + const uint8_t* matB = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
135 + const void* b_scale = b_scales.data_ptr();
136 + const void* b_zero = has_zp ? b_zeros.value().data_ptr() : nullptr;
137 +
138 + uint8_t* matB_reorder =
139 + reinterpret_cast<uint8_t*>(b_qweight_reorder.data_ptr());
140 + void* b_scale_reorder = b_scales_reorder.data_ptr();
141 + void* b_zero_reorder = has_zp ? b_zeros_reorder.value().data_ptr() : nullptr;
142 +
143 + cudaStream_t stream = at::cuda::getCurrentCUDAStream();
144 + if (b_scales.dtype() == at::ScalarType::Half) {
145 + allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(
146 + matB, reinterpret_cast<const __half*>(b_scale),
147 + reinterpret_cast<const __half*>(b_zero), matB_reorder,
148 + reinterpret_cast<__half*>(b_scale_reorder),
149 + reinterpret_cast<__half*>(b_zero_reorder), K, N, N_32align, stream);
150 + } else if (b_scales.dtype() == at::ScalarType::BFloat16) {
151 + allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__nv_bfloat16>(
152 + matB, reinterpret_cast<const __nv_bfloat16*>(b_scale),
153 + reinterpret_cast<const __nv_bfloat16*>(b_zero), matB_reorder,
154 + reinterpret_cast<__nv_bfloat16*>(b_scale_reorder),
155 + reinterpret_cast<__nv_bfloat16*>(b_zero_reorder), K, N, N_32align,
156 + stream);
157 + }
158 +}
159 +
160 +TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
161 + m.impl("rearrange_kn_weight_as_n32k16_order",
162 + &rearrange_kn_weight_as_n32k16_order);
163 +}
csrc/quantization/gptq_allspark/allspark_utils.cuh new
+408
@@ -0,0 +1,408 @@
1 +#pragma once
2 +
3 +#include <torch/all.h>
4 +#include <c10/cuda/CUDAGuard.h>
5 +#include <ATen/cuda/CUDAContext.h>
6 +#include <cuda_runtime.h>
7 +#include <cuda_fp16.h>
8 +#include <cuda_bf16.h>
9 +#include <iostream>
10 +
11 +namespace allspark {
12 +
13 +#define CHECK_CUDA(cmd) \
14 + do { \
15 + cudaError_t cuda_status = cmd; \
16 + if (cuda_status != cudaSuccess) { \
17 + std::string err_str = cudaGetErrorString(cuda_status); \
18 + std::cerr << "Failed: " << __FILE__ << ":" << __LINE__ << " " \
19 + << err_str; \
20 + exit(-1); \
21 + } \
22 + } while (0)
23 +
24 +#define CHECK_CUBLAS(cmd) \
25 + do { \
26 + cublasStatus_t cublas_status = cmd; \
27 + if (cublas_status != CUBLAS_STATUS_SUCCESS) { \
28 + std::cerr << "Failed: " << __FILE__ << ":" << __LINE__ << " " \
29 + << cublas_status << std::endl; \
30 + exit(-1); \
31 + } \
32 + } while (0)
33 +
34 +template <typename FType, typename QType>
35 +struct SM8x_GEMM_W8A16_Splitk_Params {
36 + const FType* A_ptr;
37 + const QType* B_ptr;
38 + const FType* B_scale_ptr;
39 + const FType* B_zero_ptr;
40 + FType* C_ptr;
41 + int M;
42 + int N;
43 + int K;
44 + int SplitK;
45 + int GroupCnt;
46 + int GroupSize;
47 + FType* C_split_ptr; // for non-fused splitk reduce
48 + float* C_tmp_ptr; // for fused splitk reduce
49 + uint32_t* red_count_ptr; // for fused splitk reduce
50 +};
51 +
52 +struct alignas(16) BlockTileSplitkParams {
53 + int Mtile;
54 + int Ntile;
55 + int SplitK;
56 + bool EnableFuse;
57 +};
58 +
59 +template <typename FType, int BLOCK, int N_MATRIX>
60 +__global__ void f16_gemm_splitk_reduce_kernel(const FType* C_split, FType* C,
61 + uint32_t n, uint32_t n_matrix,
62 + uint32_t matrix_size) {
63 + int idx = blockIdx.x * BLOCK + threadIdx.x;
64 +
65 + if (idx >= matrix_size) {
66 + return;
67 + }
68 +
69 + FType sum(0);
70 +
71 + int n_mat = N_MATRIX > 0 ? N_MATRIX : (int)n_matrix;
72 + for (int i = 0; i < n_mat; ++i) {
73 + sum += C_split[idx + i * matrix_size];
74 + }
75 +
76 + C[idx] = sum;
77 +}
78 +
79 +template <typename FType>
80 +void f16_gemm_splitk_reduce(const FType* C_split, FType* C, const uint32_t m,
81 + const uint32_t n, const uint32_t n_matrix,
82 + cudaStream_t stream) {
83 + const int BLOCK = 128;
84 + uint32_t matrix_size = m * n;
85 + int grid = (matrix_size + BLOCK - 1) / BLOCK;
86 +
87 + void (*kernel)(const FType*, FType*, uint32_t, uint32_t, uint32_t) = nullptr;
88 +
89 + switch (n_matrix) {
90 + case 4:
91 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 4>;
92 + break;
93 + case 5:
94 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 5>;
95 + break;
96 + case 6:
97 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 6>;
98 + break;
99 + case 7:
100 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 7>;
101 + break;
102 + case 8:
103 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 8>;
104 + break;
105 + case 9:
106 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 9>;
107 + break;
108 + case 10:
109 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 10>;
110 + break;
111 + case 11:
112 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 11>;
113 + break;
114 + case 12:
115 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, 12>;
116 + break;
117 + default:
118 + kernel = f16_gemm_splitk_reduce_kernel<FType, BLOCK, -1>;
119 + break;
120 + }
121 +
122 + kernel<<<grid, BLOCK, 0, stream>>>(C_split, C, n, n_matrix, matrix_size);
123 +}
124 +
125 +template <typename T>
126 +struct HalfType;
127 +template <>
128 +struct HalfType<half> {
129 + using T1 = __half;
130 + using T2 = __half2;
131 +};
132 +template <>
133 +struct HalfType<__nv_bfloat16> {
134 + using T1 = __nv_bfloat16;
135 + using T2 = __nv_bfloat162;
136 +};
137 +
138 +// convert 64-bit pointer to 32-bit smem addr
139 +__device__ __forceinline__ uint32_t smem_u32addr(const void* smem_ptr) {
140 + uint32_t addr;
141 + asm("{.reg .u64 u64addr;\n"
142 + " cvta.to.shared.u64 u64addr, %1;\n"
143 + " cvt.u32.u64 %0, u64addr;}\n"
144 + : "=r"(addr)
145 + : "l"(smem_ptr));
146 +
147 + return addr;
148 +}
149 +
150 +template <typename T>
151 +__device__ __forceinline__ void ldg16_cg_0(T& r0, const void* ptr, bool guard) {
152 + static_assert(sizeof(T) == 2, "ldg16_cg_0: invalid T");
153 +
154 + asm volatile(
155 + "{.reg .pred p;\n"
156 + " setp.ne.b32 p, %2, 0;\n"
157 + " @!p mov.b16 %0, 0;\n"
158 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDACC_VER_MINOR__ >= 4 && \
159 + __CUDA_ARCH__ >= 750
160 + " @p ld.global.cg.L2::128B.b16 {%0}, [%1];}\n"
161 +#else
162 + " @p ld.global.ca.b16 {%0}, [%1];}\n"
163 +#endif
164 + : "=h"(reinterpret_cast<uint16_t&>(r0))
165 + : "l"(ptr), "r"((int)guard));
166 +}
167 +
168 +template <typename T>
169 +__device__ __forceinline__ void ldg64_ca(T& r0, T& r1, const void* ptr,
170 + bool guard) {
171 + static_assert(sizeof(T) == 4, "ldg64_ca: invalid T");
172 +
173 + asm volatile(
174 + "{.reg .pred p;\n"
175 + " setp.ne.b32 p, %3, 0;\n"
176 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDACC_VER_MINOR__ >= 4 && \
177 + __CUDA_ARCH__ >= 750
178 + " @p ld.global.ca.L2::128B.v2.b32 {%0, %1}, [%2];}\n"
179 +#else
180 + " @p ld.global.ca.v2.b32 {%0, %1}, [%2];}\n"
181 +#endif
182 + : "=r"(reinterpret_cast<uint32_t&>(r0)),
183 + "=r"(reinterpret_cast<uint32_t&>(r1))
184 + : "l"(ptr), "r"((int)guard));
185 +}
186 +
187 +template <typename T>
188 +__device__ __forceinline__ void ldg128_cg_0(T& r0, T& r1, T& r2, T& r3,
189 + const void* ptr, bool guard) {
190 + static_assert(sizeof(T) == 4, "ldg128_cg_0: invalid T");
191 +
192 + asm volatile(
193 + "{.reg .pred p;\n"
194 + " setp.ne.b32 p, %5, 0;\n"
195 + " @!p mov.b32 %0, 0;\n"
196 + " @!p mov.b32 %1, 0;\n"
197 + " @!p mov.b32 %2, 0;\n"
198 + " @!p mov.b32 %3, 0;\n"
199 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDACC_VER_MINOR__ >= 4 && \
200 + __CUDA_ARCH__ >= 750
201 + " @p ld.global.cg.L2::128B.v4.b32 {%0, %1, %2, %3}, [%4];}\n"
202 +#else
203 + " @p ld.global.cg.v4.b32 {%0, %1, %2, %3}, [%4];}\n"
204 +#endif
205 + : "=r"(reinterpret_cast<uint32_t&>(r0)),
206 + "=r"(reinterpret_cast<uint32_t&>(r1)),
207 + "=r"(reinterpret_cast<uint32_t&>(r2)),
208 + "=r"(reinterpret_cast<uint32_t&>(r3))
209 + : "l"(ptr), "r"((int)guard));
210 +}
211 +
212 +template <typename T>
213 +__device__ __forceinline__ void lds128(T& reg0, T& reg1, T& reg2, T& reg3,
214 + const uint32_t addr) {
215 + static_assert(sizeof(T) == 4, "lds128: invalid T");
216 +
217 + asm volatile("ld.shared.v4.b32 {%0, %1, %2, %3}, [%4];\n"
218 + : "=r"(reinterpret_cast<uint32_t&>(reg0)),
219 + "=r"(reinterpret_cast<uint32_t&>(reg1)),
220 + "=r"(reinterpret_cast<uint32_t&>(reg2)),
221 + "=r"(reinterpret_cast<uint32_t&>(reg3))
222 + : "r"(addr));
223 +}
224 +
225 +template <typename T>
226 +__device__ __forceinline__ void stg128(const T& r0, const T& r1, const T& r2,
227 + const T& r3, const void* ptr,
228 + bool guard) {
229 + static_assert(sizeof(T) == 4, "stg128: invalid T");
230 +
231 + asm volatile(
232 + "{.reg .pred p;\n"
233 + " setp.ne.b32 p, %1, 0;\n"
234 + " @p st.global.v4.b32 [%0], {%2, %3, %4, %5};}\n"
235 + :
236 + : "l"(ptr), "r"((int)guard), "r"(reinterpret_cast<const uint32_t&>(r0)),
237 + "r"(reinterpret_cast<const uint32_t&>(r1)),
238 + "r"(reinterpret_cast<const uint32_t&>(r2)),
239 + "r"(reinterpret_cast<const uint32_t&>(r3)));
240 +}
241 +
242 +template <typename T>
243 +__device__ __forceinline__ void ldsm_4(T& r0, T& r1, T& r2, T& r3,
244 + const uint32_t& addr) {
245 + static_assert(sizeof(T) == 4, "ldsm_4: invalid T");
246 +#if (__CUDA_ARCH__ >= 750) && (__CUDACC_VER_MAJOR__ >= 11)
247 + asm volatile(
248 + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
249 + : "=r"(reinterpret_cast<uint32_t&>(r0)),
250 + "=r"(reinterpret_cast<uint32_t&>(r1)),
251 + "=r"(reinterpret_cast<uint32_t&>(r2)),
252 + "=r"(reinterpret_cast<uint32_t&>(r3))
253 + : "r"(addr));
254 +#endif
255 +}
256 +
257 +template <typename FType>
258 +__device__ __forceinline__ void hmma16816_f32(float (&d)[4],
259 + const uint32_t (&a)[4],
260 + const uint32_t (&b)[2]);
261 +
262 +template <>
263 +__device__ __forceinline__ void hmma16816_f32<__half>(float (&d)[4],
264 + const uint32_t (&a)[4],
265 + const uint32_t (&b)[2]) {
266 +#if (__CUDA_ARCH__ >= 800) && (__CUDACC_VER_MAJOR__ >= 11)
267 + asm volatile(
268 + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, "
269 + "{%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
270 + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
271 + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
272 +#endif
273 +}
274 +
275 +template <>
276 +__device__ __forceinline__ void hmma16816_f32<__nv_bfloat16>(
277 + float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) {
278 +#if (__CUDA_ARCH__ >= 800) && (__CUDACC_VER_MAJOR__ >= 11)
279 + asm volatile(
280 + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, "
281 + "{%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n"
282 + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
283 + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]));
284 +#endif
285 +}
286 +
287 +template <int SIZE_IN_BYTES>
288 +__device__ __forceinline__ void cp_async(const uint32_t smem_addr,
289 + const void* gmem_ptr,
290 + const int src_in_bytes, bool guard) {
291 + static_assert(
292 + (SIZE_IN_BYTES == 4 || SIZE_IN_BYTES == 8 || SIZE_IN_BYTES == 16),
293 + "Size is not supported");
294 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
295 + asm volatile(
296 + "{.reg.pred p;\n"
297 + " setp.ne.b32 p, %4, 0;\n"
298 + #if __CUDACC_VER_MINOR__ >= 4
299 + " @p cp.async.cg.shared.global.L2::256B [%0], [%1], %2, %3;}\n"
300 + #else
301 + " @p cp.async.cg.shared.global [%0], [%1], %2, %3;}\n"
302 + #endif
303 + ::"r"(smem_addr),
304 + "l"(gmem_ptr), "n"(SIZE_IN_BYTES), "r"(src_in_bytes), "r"((int)guard));
305 +#endif
306 +}
307 +
308 +template <int SIZE_IN_BYTES>
309 +__device__ __forceinline__ void cp_async_ca(const uint32_t smem_addr,
310 + const void* gmem_ptr,
311 + const int src_in_bytes,
312 + bool guard) {
313 + static_assert(
314 + (SIZE_IN_BYTES == 4 || SIZE_IN_BYTES == 8 || SIZE_IN_BYTES == 16),
315 + "Size is not supported");
316 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
317 + asm volatile(
318 + "{.reg.pred p;\n"
319 + " setp.ne.b32 p, %4, 0;\n"
320 + #if __CUDACC_VER_MINOR__ >= 4
321 + " @p cp.async.ca.shared.global.L2::256B [%0], [%1], %2, %3;}\n"
322 + #else
323 + " @p cp.async.ca.shared.global [%0], [%1], %2, %3;}\n"
324 + #endif
325 + ::"r"(smem_addr),
326 + "l"(gmem_ptr), "n"(SIZE_IN_BYTES), "r"(src_in_bytes), "r"((int)guard));
327 +#endif
328 +}
329 +
330 +__device__ __forceinline__ void cp_async_commit_group() {
331 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
332 + asm volatile("cp.async.commit_group;\n");
333 +#endif
334 +}
335 +
336 +template <int N>
337 +__device__ __forceinline__ void cp_asyc_wait_group() {
338 +#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
339 + asm volatile("cp.async.wait_group %0;\n" : : "n"(N));
340 +#endif
341 +}
342 +
343 +template <typename T>
344 +__device__ __forceinline__ void cvt_8bx4_to_16bx4_bias128(const uint32_t& idata,
345 + T* fdata);
346 +
347 +template <>
348 +// fast conversion: 4xuint8 to 4xhalf, subtracting bias = 128
349 +__device__ __forceinline__ void cvt_8bx4_to_16bx4_bias128<__half2>(
350 + const uint32_t& idata, __half2* fdata) {
351 + uint32_t i10, i32;
352 + asm volatile(
353 + "prmt.b32 %0, %2, 0x64, 0x4140;"
354 + "prmt.b32 %1, %2, 0x64, 0x4342;"
355 + : "=r"(i10), "=r"(i32)
356 + : "r"(idata));
357 +
358 + static constexpr uint32_t MAGIC_NUM = 0x64806480;
359 + fdata[0] = __hsub2(reinterpret_cast<const __half2&>(i10),
360 + reinterpret_cast<const __half2&>(MAGIC_NUM));
361 + fdata[1] = __hsub2(reinterpret_cast<const __half2&>(i32),
362 + reinterpret_cast<const __half2&>(MAGIC_NUM));
363 +}
364 +
365 +template <>
366 +// fast conversion: 4xuint8 to 4xbfloat16, subtracting bias = 128
367 +// reference from marlin fast implementation
368 +__device__ __forceinline__ void cvt_8bx4_to_16bx4_bias128<__nv_bfloat162>(
369 + const uint32_t& idata, __nv_bfloat162* fdata) {
370 + float fp32_imd[4];
371 + uint32_t* fp32_imd_casted = reinterpret_cast<uint32_t*>(fp32_imd);
372 + asm volatile(
373 + "prmt.b32 %0, %4, 0x4B000000, 0x7650;"
374 + "prmt.b32 %1, %4, 0x4B000000, 0x7651;"
375 + "prmt.b32 %2, %4, 0x4B000000, 0x7652;"
376 + "prmt.b32 %3, %4, 0x4B000000, 0x7653;"
377 + : "=r"(fp32_imd_casted[0]), "=r"(fp32_imd_casted[1]),
378 + "=r"(fp32_imd_casted[2]), "=r"(fp32_imd_casted[3])
379 + : "r"(idata));
380 +
381 + fp32_imd[0] -= 8388736.f;
382 + fp32_imd[1] -= 8388736.f;
383 + fp32_imd[2] -= 8388736.f;
384 + fp32_imd[3] -= 8388736.f;
385 +
386 + uint32_t* bf16_res = reinterpret_cast<uint32_t*>(fdata);
387 + asm volatile(
388 + "prmt.b32 %0, %2, %3, 0x7632;"
389 + "prmt.b32 %1, %4, %5, 0x7632;"
390 + : "=r"(bf16_res[0]), "=r"(bf16_res[1])
391 + : "r"(fp32_imd_casted[0]), "r"(fp32_imd_casted[1]),
392 + "r"(fp32_imd_casted[2]), "r"(fp32_imd_casted[3]));
393 +}
394 +
395 +static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) {
396 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
397 + assert(false);
398 +#else
399 + return __bfloat162bfloat162(x);
400 +#endif
401 + __builtin_unreachable(); // Suppress missing return statement warning
402 +}
403 +
404 +static __device__ half2 inline num2num2(const half x) {
405 + return __half2half2(x);
406 +}
407 +
408 +} // namespace allspark
\ No newline at end of file
csrc/torch_bindings.cpp
+19
@@ -447,6 +447,25 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
447 "Tensor!? azp) -> ()");
448 ops.impl("dynamic_scaled_int8_quant", torch::kCUDA,
449 &dynamic_scaled_int8_quant);
450 +
451 +#ifndef USE_ROCM
452 + // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
453 + ops.def(
454 + "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
455 + "Tensor? b_zeros, "
456 + "bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
457 + "Tensor!? b_zeros_reorder, "
458 + "int K, int N, int N_32align) -> ()");
459 + // conditionally compiled so impl in source file
460 +
461 + // AllSpark quantization ops
462 + ops.def(
463 + "allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
464 + "Tensor? b_qzeros, "
465 + "SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
466 + "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
467 + // conditionally compiled so impl in source file
468 +#endif
469 }
470
471 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
tests/kernels/test_allspark_gemm.py new
+100
@@ -0,0 +1,100 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +import torch
4 +
5 +from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
6 +from vllm import _custom_ops as ops
7 +from vllm.model_executor.layers.quantization.utils.allspark_utils import (
8 + ALLSPARK_AMPERE_K_ALIGN, ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
9 + ALLSPARK_AMPERE_N_ALIGN)
10 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 + quantize_weights)
12 +from vllm.platforms import current_platform
13 +from vllm.scalar_type import scalar_types
14 +
15 +
16 +def is_gptq_allspark_supported(min_capability: int,
17 + max_capability: int) -> bool:
18 + if not current_platform.is_cuda():
19 + return False
20 +
21 + capability = current_platform.get_device_capability()
22 + assert capability is not None
23 +
24 + return capability.to_int() >= min_capability \
25 + and capability.to_int() <= max_capability
26 +
27 +
28 +MNK_FACTORS = [
29 + (1, 4, 8),
30 + (13, 17, 67),
31 + (26, 37, 13),
32 + (48, 16, 24),
33 + (67, 13, 88),
34 + (257, 13, 11),
35 + (658, 13, 11),
36 + (1033, 9, 17),
37 +]
38 +
39 +DTYPES = [torch.float16, torch.bfloat16]
40 +HAS_ZP_OPTS = [False, True]
41 +
42 +
43 +def compute_max_diff(output, output_ref):
44 + return torch.mean(torch.abs(output - output_ref)) / torch.mean(
45 + torch.abs(output_ref))
46 +
47 +
48 +def rand_data(shape, dtype=torch.float16):
49 + return torch.randn(shape, dtype=dtype, device="cuda")
50 +
51 +
52 +@pytest.mark.skipif(
53 + not is_gptq_allspark_supported(80, 89),
54 + reason="AllSpark Ampere kernel is not supported on this GPU type.")
55 +@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
56 +@pytest.mark.parametrize("group_size", [-1])
57 +@pytest.mark.parametrize("has_zp", HAS_ZP_OPTS)
58 +@pytest.mark.parametrize("dtype", DTYPES)
59 +def test_gptq_allspark_gemm_ampere(mnk_factors, group_size, has_zp, dtype):
60 + m_factor, n_factor, k_factor = mnk_factors
61 + m = m_factor
62 + n = n_factor * ALLSPARK_AMPERE_N_ALIGN
63 + k = k_factor * ALLSPARK_AMPERE_K_ALIGN
64 +
65 + input = rand_data((m, k), dtype=dtype)
66 + weight = rand_data((k, n), dtype=dtype)
67 +
68 + # Quantize (and apply act_order if provided)
69 + w_ref, qw, s, zp = quantize_weights(weight, scalar_types.uint8b128,
70 + group_size, has_zp)
71 +
72 + qw = qw.to(torch.uint8)
73 + if has_zp:
74 + zp = zp.to(dtype)
75 + properties = torch.cuda.get_device_properties(qw.device.index)
76 + sm_count = properties.multi_processor_count
77 + sm_version = properties.major * 10 + properties.minor
78 +
79 + n_32align = (n + 32 - 1) // 32 * 32
80 +
81 + qw_reorder, s_reorder, zp_reorder = ops.allspark_repack_weight(
82 + qw, s, zp, has_zp)
83 + opcheck(torch.ops._C.rearrange_kn_weight_as_n32k16_order,
84 + (qw, s, zp, has_zp, qw_reorder, s_reorder, zp_reorder, k, n,
85 + n_32align))
86 +
87 + opcheck(torch.ops._C.allspark_w8a16_gemm,
88 + (input, qw_reorder, s_reorder, zp_reorder, n, group_size, sm_count,
89 + sm_version, ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD, has_zp, True),
90 + test_utils=DEFAULT_OPCHECK_TEST_UTILS)
91 + output = ops.allspark_w8a16_gemm(input, qw_reorder, s_reorder, zp_reorder,
92 + n, group_size, sm_count, sm_version,
93 + ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
94 + has_zp, True)
95 +
96 + output_ref = torch.matmul(input, w_ref)
97 + torch.cuda.synchronize()
98 + max_diff = compute_max_diff(output, output_ref)
99 +
100 + assert max_diff < 0.04
tests/quantization/test_compressed_tensors.py
-2
@@ -215,8 +215,6 @@ def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
215 assert qkv_proj.scheme.group_size == (-1
216 if group is None else group)
217
218 - assert qkv_proj.weight_packed.dtype is torch.int32
219 - assert qkv_proj.weight_scale.dtype is torch.float16
218 assert qkv_proj.scheme.pack_factor == pack_factor
219
220 llm.apply_model(check_model)
vllm/_custom_ops.py
+77
@@ -404,6 +404,22 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
404 memory_format=torch.contiguous_format)
405
406
407 +if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
408 +
409 + @register_fake("_C::allspark_w8a16_gemm")
410 + def _allspark_w8a16_gemm_fake(a: torch.Tensor, b_qweight: torch.Tensor,
411 + b_scales: torch.Tensor,
412 + b_qzeros: Optional[torch.Tensor],
413 + n: torch.SymInt, group_size: torch.SymInt,
414 + sm_count: torch.SymInt,
415 + sm_version: torch.SymInt,
416 + CUBLAS_M_THRESHOLD: torch.SymInt,
417 + has_zp: bool,
418 + n32k16_reorder: bool) -> torch.Tensor:
419 + m = a.size(0)
420 + return torch.empty((m, n), device=a.device, dtype=a.dtype)
421 +
422 +
423 if hasattr(torch.ops._C, "ggml_dequantize"):
424
425 @register_fake("_C::ggml_dequantize")
@@ -881,6 +897,67 @@ def scaled_fp8_quant(
897 return output, scale
898
899
900 +# gptq allspark
901 +def allspark_repack_weight(
902 + qweight: torch.Tensor,
903 + scale: torch.Tensor,
904 + zero_point: Optional[torch.Tensor] = None,
905 + has_zp: bool = False
906 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
907 + """
908 + Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format
909 + for Ampere W8A16 Fused Gemm kernel
910 +
911 + Args:
912 + qweight: uint8 weight tensor, original k x n format.
913 + scale: fp16/bf16 weight scale tensor, 1 x n format.
914 + zero_point: fp16/bf16 weight zero_point tensor, 1 x n format.
915 + Must be provided for asymmetric quantization.
916 + has_zp: if use symmetric quantization, has_zp = False.
917 + if use asymmetric quantization, has_zp = True.
918 +
919 + Returns:
920 + Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] :
921 + rearranged weight, scale, and optionally zero_point.
922 + """
923 + K = qweight.shape[0]
924 + N = qweight.shape[1]
925 + N_32align = (N + 32 - 1) // 32 * 32
926 +
927 + qweight_reorder = torch.empty((N_32align, K),
928 + device=qweight.device,
929 + dtype=qweight.dtype)
930 + scale_reorder = torch.empty((1, N_32align),
931 + device=scale.device,
932 + dtype=scale.dtype)
933 + zero_point_reorder = None
934 + if has_zp:
935 + assert zero_point is not None, (
936 + "zero_point must be provided for asymmetric quantization.")
937 + zero_point_reorder = torch.empty((1, N_32align),
938 + device=zero_point.device,
939 + dtype=zero_point.dtype)
940 +
941 + torch.ops._C.rearrange_kn_weight_as_n32k16_order(
942 + qweight, scale, zero_point, has_zp, qweight_reorder, scale_reorder,
943 + zero_point_reorder, K, N, N_32align)
944 +
945 + return qweight_reorder, scale_reorder, zero_point_reorder
946 +
947 +
948 +def allspark_w8a16_gemm(a: torch.Tensor, b_qweight: torch.Tensor,
949 + b_scales: torch.Tensor,
950 + b_qzeros: Optional[torch.Tensor], n: int,
951 + group_size: int, sm_count: int, sm_version: int,
952 + CUBLAS_M_THRESHOLD: int, has_zp: bool,
953 + n32k16_reorder: bool) -> torch.Tensor:
954 +
955 + return torch.ops._C.allspark_w8a16_gemm(a, b_qweight, b_scales, b_qzeros,
956 + n, group_size, sm_count,
957 + sm_version, CUBLAS_M_THRESHOLD,
958 + has_zp, n32k16_reorder)
959 +
960 +
961 # int8
962 def scaled_int8_quant(
963 input: torch.Tensor,
vllm/model_executor/layers/quantization/kernels/mixed_precision/__init__.py
+3
@@ -3,6 +3,8 @@
3 from typing import List, Optional, Type
4
5 import vllm.envs as envs
6 +from vllm.model_executor.layers.quantization.kernels.mixed_precision.allspark import ( # noqa: E501
7 + AllSparkLinearKernel)
8 from vllm.model_executor.layers.quantization.kernels.mixed_precision.exllama import ( # noqa: E501
9 ExllamaLinearKernel)
10 from vllm.model_executor.layers.quantization.kernels.mixed_precision.machete import ( # noqa: E501
@@ -16,6 +18,7 @@ from vllm.platforms import current_platform
18 # in priority/performance order (when available)
19 _POSSIBLE_KERNELS: List[Type[MPLinearKernel]] = [
20 MacheteLinearKernel,
21 + AllSparkLinearKernel,
22 MarlinLinearKernel,
23 ExllamaLinearKernel,
24 ]
vllm/model_executor/layers/quantization/kernels/mixed_precision/allspark.py new
+115
@@ -0,0 +1,115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from typing import Optional, Tuple
4 +
5 +import torch
6 +
7 +from vllm import _custom_ops as ops
8 +from vllm.model_executor.layers.quantization.utils import replace_parameter
9 +from vllm.model_executor.layers.quantization.utils.allspark_utils import (
10 + ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD, check_allspark_supported_dtype_shape)
11 +from vllm.model_executor.parameter import (BasevLLMParameter,
12 + permute_param_layout_)
13 +
14 +from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
15 +
16 +
17 +class AllSparkLinearKernel(MPLinearKernel):
18 +
19 + @classmethod
20 + def get_min_capability(cls) -> int:
21 + return 80
22 +
23 + @classmethod
24 + def can_implement(cls,
25 + c: MPLinearLayerConfig) -> Tuple[bool, Optional[str]]:
26 + if c.has_g_idx:
27 + return False, "Act reordering currently not supported by AllSpark"
28 +
29 + if c.zero_points:
30 + return False, "Zero points currently not supported by AllSpark"
31 +
32 + return check_allspark_supported_dtype_shape(
33 + c.partition_weight_shape[0], # in_features
34 + c.partition_weight_shape[1], # out_features
35 + c.group_size,
36 + c.weight_type,
37 + c.act_type)
38 +
39 + # note assumes that
40 + # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
41 + # `weight_scale` is: {input_dim = 0, output_dim = 1}
42 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
43 + device = getattr(layer, self.w_q_name).device
44 + c = self.config
45 +
46 + # prepare the parameters required for the kernel
47 + properties = torch.cuda.get_device_properties(device.index)
48 + sm_count = properties.multi_processor_count
49 + sm_version = properties.major * 10 + properties.minor
50 + gemm_args = {}
51 + gemm_args['sm_count'] = sm_count
52 + gemm_args['sm_version'] = sm_version
53 +
54 + self.gemm_args = gemm_args
55 +
56 + # transform param weight, scale
57 + old_weight_param = getattr(layer, self.w_q_name)
58 + old_scale_param = getattr(layer, self.w_s_name)
59 +
60 + assert isinstance(old_weight_param, BasevLLMParameter)
61 + permute_param_layout_(old_weight_param,
62 + input_dim=0,
63 + output_dim=1,
64 + packed_dim=0)
65 +
66 + assert isinstance(old_scale_param, BasevLLMParameter)
67 + permute_param_layout_(old_scale_param, input_dim=0, output_dim=1)
68 +
69 + # unpack weight from K / 4 x N int32 to K x N uint8
70 + new_weight_param = torch.nn.Parameter(old_weight_param.data,
71 + requires_grad=False)
72 + new_weight_param.data = new_weight_param.data.t().contiguous().view(
73 + dtype=torch.uint8)
74 + new_weight_param.data = new_weight_param.data.t().contiguous()
75 +
76 + new_scale_param = torch.nn.Parameter(old_scale_param.data,
77 + requires_grad=False)
78 +
79 + # reorder K x N weight as N32K16 format for Ampere W8A16
80 + new_weight_param.data, new_scale_param.data, _ = \
81 + ops.allspark_repack_weight(
82 + new_weight_param.data, new_scale_param.data, None,
83 + c.zero_points)
84 +
85 + replace_parameter(layer, self.w_q_name, new_weight_param.data)
86 + replace_parameter(layer, self.w_s_name, new_scale_param.data)
87 +
88 + def apply_weights(self,
89 + layer: torch.nn.Module,
90 + x: torch.Tensor,
91 + bias: Optional[torch.Tensor] = None) -> torch.Tensor:
92 + c = self.config
93 + gemm_args = self.gemm_args
94 + w_q, w_s, _, _ = self._get_weight_params(layer)
95 +
96 + reshaped_x = x.reshape(-1, x.shape[-1])
97 + out_shape = x.shape[:-1] + (c.partition_weight_shape[1], )
98 +
99 + output = ops.allspark_w8a16_gemm(
100 + a=reshaped_x,
101 + b_qweight=w_q,
102 + b_scales=w_s,
103 + b_qzeros=None,
104 + n=c.partition_weight_shape[1],
105 + group_size=c.group_size,
106 + sm_count=gemm_args['sm_count'],
107 + sm_version=gemm_args['sm_version'],
108 + CUBLAS_M_THRESHOLD=ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
109 + has_zp=c.zero_points,
110 + n32k16_reorder=True)
111 +
112 + if bias is not None:
113 + output.add_(bias) # In-place add
114 +
115 + return output.reshape(out_shape)
vllm/model_executor/layers/quantization/utils/allspark_utils.py new
+51
@@ -0,0 +1,51 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import torch
4 +
5 +from vllm.platforms import current_platform
6 +from vllm.scalar_type import ScalarType, scalar_types
7 +
8 +ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD = 1024
9 +ALLSPARK_SUPPORTED_QUANT_TYPES = [scalar_types.uint8b128]
10 +ALLSPARK_AMPERE_N_ALIGN = 16
11 +ALLSPARK_AMPERE_K_ALIGN = 16
12 +
13 +
14 +def check_allspark_supported_dtype_shape(input_size_per_partition: int,
15 + output_size_per_partition: int,
16 + group_size: int,
17 + weight_dtype: ScalarType,
18 + act_dtype: torch.dtype):
19 + capability_tuple = current_platform.get_device_capability()
20 + device_capability = (-1 if capability_tuple is None else
21 + capability_tuple.to_int())
22 +
23 + # For Ampere GPU
24 + if device_capability >= 80 and device_capability < 90:
25 + if group_size != -1:
26 + return False, \
27 + "For Ampere GPU, AllSpark does not support group_size "\
28 + f"= {group_size}. Only group_size = -1 are supported."
29 +
30 + if weight_dtype not in ALLSPARK_SUPPORTED_QUANT_TYPES:
31 + return False, "For Ampere GPU, AllSpark does not support "\
32 + f"quant type ({weight_dtype}). Only quant type "\
33 + f"({ALLSPARK_SUPPORTED_QUANT_TYPES}) are supported."
34 +
35 + if input_size_per_partition % ALLSPARK_AMPERE_K_ALIGN != 0 \
36 + or output_size_per_partition % ALLSPARK_AMPERE_N_ALIGN != 0:
37 + return False, \
38 + "AllSpark needs input_size_per_partition % "\
39 + f"{ALLSPARK_AMPERE_K_ALIGN} = 0 and "\
40 + f"output_size_per_partition % {ALLSPARK_AMPERE_N_ALIGN} = 0 "\
41 + "for Ampere GPU optimized kernels."
42 +
43 + if act_dtype != torch.float16 and act_dtype != torch.bfloat16:
44 + return False, \
45 + "AllSpark only supports act_dtype = float16 or bfloat16,"\
46 + f"for Ampere GPU, but got act_dtype = {act_dtype}."
47 + else:
48 + return False, "AllSpark currently does not support "\
49 + f"device_capability = {device_capability}."
50 +
51 + return True, None