@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 25965ee583

[Kernel] Remove AllSpark INT8 W8A16 GEMM backend (#58001)

Signed-off-by: mgoin <mgoin64@gmail.com>

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