@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b6efafd9e4
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, 06:00 UTC

[Perf] Vectorize static / dynamic INT8 quant kernels (#19233)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Jun 12, 2025 at 09:51 UTC b6efafd9e43fb2ecc18e7789206f11754a50c7d1
4 files changed +415 -101
benchmarks/kernels/bench_int8_gemm.py new
+200
@@ -0,0 +1,200 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import argparse
4 +import copy
5 +import itertools
6 +
7 +import torch
8 +from weight_shapes import WEIGHT_SHAPES
9 +
10 +from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
11 +from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant
12 +from vllm.triton_utils import triton
13 +
14 +
15 +@triton.testing.perf_report(
16 + triton.testing.Benchmark(
17 + x_names=["batch_size"],
18 + x_vals=[1, 16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384],
19 + x_log=False,
20 + line_arg="provider",
21 + line_vals=[
22 + "torch-bf16",
23 + # "int8-tensor-w-token-a",
24 + "int8-tensor-w-tensor-a",
25 + "int8-channel-w-token-a",
26 + # "int8-channel-w-tensor-a",
27 + # "int8-tensor-w-token-a-noquant",
28 + "int8-tensor-w-tensor-a-noquant",
29 + "int8-channel-w-token-a-noquant",
30 + # "int8-channel-w-tensor-a-noquant",
31 + ],
32 + line_names=[
33 + "torch-bf16",
34 + # "int8-tensor-w-token-a",
35 + "int8-tensor-w-tensor-a",
36 + "int8-channel-w-token-a",
37 + # "int8-channel-w-tensor-a",
38 + # "int8-tensor-w-token-a-noquant",
39 + "int8-tensor-w-tensor-a-noquant",
40 + "int8-channel-w-token-a-noquant",
41 + # "int8-channel-w-tensor-a-noquant",
42 + ],
43 + ylabel="TFLOP/s (larger is better)",
44 + plot_name="BF16 vs INT8 GEMMs",
45 + args={},
46 + )
47 +)
48 +def benchmark(batch_size, provider, N, K):
49 + M = batch_size
50 + device = "cuda"
51 + dtype = torch.bfloat16
52 + a = torch.randn((M, K), device=device, dtype=dtype)
53 + b = torch.randn((N, K), device=device, dtype=dtype)
54 +
55 + quantiles = [0.5, 0.2, 0.8]
56 +
57 + if "torch-bf16" in provider:
58 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
59 + lambda: torch.nn.functional.linear(a, b), quantiles=quantiles
60 + )
61 +
62 + elif "int8" in provider:
63 + # Weights are always quantized ahead of time
64 + if "noquant" in provider:
65 + # For "no quant", we don't measure the time for activations
66 + if "tensor-w-token-a" in provider:
67 + # Dynamic per-token quant for A, static per-tensor quant for B
68 + scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
69 + b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b)
70 + assert scale_b_int8.numel() == 1
71 + a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a)
72 +
73 + elif "tensor-w-tensor-a" in provider:
74 + # Static per-tensor quantization with fixed scales for both A and B
75 + scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
76 + scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
77 + b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b)
78 + assert scale_b_int8.numel() == 1
79 + a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a)
80 +
81 + elif "channel-w-token-a" in provider:
82 + # Dynamic per-channel quantization for weights, per-token quant for A
83 + b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b)
84 + assert scale_b_int8.numel() == N
85 + a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a)
86 +
87 + elif "channel-w-tensor-a" in provider:
88 + # Dynamic per-channel quantization for weights, per-tensor quant for A
89 + scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
90 + b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b)
91 + assert scale_b_int8.numel() == N
92 + a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a)
93 +
94 + def run_quant():
95 + return vllm_scaled_mm(a_int8, b_int8, scale_a_int8, scale_b_int8, dtype)
96 +
97 + else:
98 + # Quantize the activations during the GEMM call
99 + if "tensor-w-token-a" in provider:
100 + # Dynamic per-token quant for A, static per-tensor quant for B
101 + scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
102 + b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b)
103 + assert scale_b_int8.numel() == 1
104 +
105 + def run_quant():
106 + a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a)
107 + return vllm_scaled_mm(
108 + a_int8, b_int8, scale_a_int8, scale_b_int8, dtype
109 + )
110 +
111 + elif "tensor-w-tensor-a" in provider:
112 + # Static per-tensor quantization with fixed scales for both A and B
113 + scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
114 + scale_b = torch.tensor([1.0], device=device, dtype=torch.float32)
115 + b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b, scale_b)
116 + assert scale_b_int8.numel() == 1
117 +
118 + def run_quant():
119 + a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a)
120 + return vllm_scaled_mm(
121 + a_int8, b_int8, scale_a_int8, scale_b_int8, dtype
122 + )
123 +
124 + elif "channel-w-token-a" in provider:
125 + # Dynamic per-channel quant for weights, per-token quant for A
126 + b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b)
127 + assert scale_b_int8.numel() == N
128 +
129 + def run_quant():
130 + a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a)
131 + return vllm_scaled_mm(
132 + a_int8, b_int8, scale_a_int8, scale_b_int8, dtype
133 + )
134 +
135 + elif "channel-w-tensor-a" in provider:
136 + # Dynamic per-channel quant for weights, static per-tensor quant for A
137 + scale_a = torch.tensor([1.0], device=device, dtype=torch.float32)
138 + b_int8, scale_b_int8, _ = vllm_scaled_int8_quant(b)
139 + assert scale_b_int8.numel() == N
140 +
141 + def run_quant():
142 + a_int8, scale_a_int8, _ = vllm_scaled_int8_quant(a, scale_a)
143 + return vllm_scaled_mm(
144 + a_int8, b_int8, scale_a_int8, scale_b_int8, dtype
145 + )
146 +
147 + b_int8 = b_int8.t()
148 +
149 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
150 + lambda: run_quant(), quantiles=quantiles
151 + )
152 +
153 + # Calculate TFLOP/s, two flops per multiply-add
154 + tflops = lambda ms: (2 * M * N * K) * 1e-12 / (ms * 1e-3)
155 + return tflops(ms), tflops(max_ms), tflops(min_ms)
156 +
157 +
158 +def prepare_shapes(args):
159 + KN_model_names = []
160 + models_tps = list(itertools.product(args.models, args.tp_sizes))
161 + for model, tp_size in models_tps:
162 + assert model in WEIGHT_SHAPES
163 + for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
164 + KN[tp_split_dim] = KN[tp_split_dim] // tp_size
165 + KN.append(model)
166 + KN_model_names.append(KN)
167 + return KN_model_names
168 +
169 +
170 +if __name__ == "__main__":
171 + parser = argparse.ArgumentParser()
172 + parser.add_argument(
173 + "--models",
174 + nargs="+",
175 + type=str,
176 + default=["meta-llama/Llama-3.1-8B-Instruct"],
177 + choices=[*WEIGHT_SHAPES.keys()],
178 + help="List of models to benchmark",
179 + )
180 + parser.add_argument(
181 + "--tp-sizes",
182 + nargs="+",
183 + type=int,
184 + default=[1],
185 + help="List of tensor parallel sizes",
186 + )
187 + args = parser.parse_args()
188 +
189 + KN_model_names = prepare_shapes(args)
190 + for K, N, model_name in KN_model_names:
191 + print(f"{model_name}, N={N} K={K}, BF16 vs INT8 GEMMs TFLOP/s:")
192 + benchmark.run(
193 + print_data=True,
194 + show_plots=True,
195 + save_path=f"bench_int8_res_n{N}_k{K}",
196 + N=N,
197 + K=K,
198 + )
199 +
200 + print("Benchmark finished!")
csrc/quantization/compressed_tensors/int8_quant_kernels.cu
+139 -101
@@ -1,15 +1,17 @@
1 #include <ATen/cuda/CUDAContext.h>
2 #include <torch/all.h>
3 +
4 #include <cmath>
5
6 #include "../../dispatch_utils.h"
7 +#include "../vectorization_utils.cuh"
8
9 #ifndef USE_ROCM
8 - #include <cub/util_type.cuh>
10 #include <cub/cub.cuh>
11 + #include <cub/util_type.cuh>
12 #else
11 - #include <hipcub/util_type.hpp>
13 #include <hipcub/hipcub.hpp>
14 + #include <hipcub/util_type.hpp>
15 #endif
16
17 static inline __device__ int8_t float_to_int8_rn(float x) {
@@ -103,134 +105,170 @@ static inline __device__ int8_t int32_to_int8(int32_t x) {
105
106 namespace vllm {
107
106 -template <typename scalar_t, typename scale_type>
108 +template <typename scalar_t, typename scale_t>
109 __global__ void static_scaled_int8_quant_kernel(
108 - scalar_t const* __restrict__ input, int8_t* __restrict__ out,
109 - scale_type const* scale_ptr, const int hidden_size) {
110 - int const tid = threadIdx.x;
111 - int64_t const token_idx = blockIdx.x;
112 - scale_type const scale = *scale_ptr;
110 + const scalar_t* __restrict__ input, int8_t* __restrict__ output,
111 + const scale_t* scale_ptr, const int hidden_size) {
112 + const int tid = threadIdx.x;
113 + const int stride = blockDim.x;
114 + const int64_t token_idx = blockIdx.x;
115 + const float scale = *scale_ptr;
116
117 // Must be performed using 64-bit math to avoid integer overflow.
115 - out += token_idx * hidden_size;
116 - input += token_idx * hidden_size;
118 + const scalar_t* row_in = input + token_idx * hidden_size;
119 + int8_t* row_out = output + token_idx * hidden_size;
120
118 - for (int i = tid; i < hidden_size; i += blockDim.x) {
119 - out[i] = float_to_int8_rn(static_cast<float>(input[i]) / scale);
120 - }
121 + vectorize_with_alignment<16>(
122 + row_in, row_out, hidden_size, tid, stride,
123 + [=] __device__(int8_t& dst, const scalar_t& src) {
124 + dst = float_to_int8_rn(static_cast<float>(src) / scale);
125 + });
126 }
127
123 -template <typename scalar_t, typename scale_type, typename azp_type>
128 +template <typename scalar_t, typename scale_t, typename azp_t>
129 __global__ void static_scaled_int8_azp_quant_kernel(
125 - scalar_t const* __restrict__ input, int8_t* __restrict__ out,
126 - scale_type const* scale_ptr, azp_type const* azp_ptr,
127 - const int hidden_size) {
128 - int const tid = threadIdx.x;
129 - int64_t const token_idx = blockIdx.x;
130 - scale_type const scale = *scale_ptr;
131 - azp_type const azp = *azp_ptr;
130 + const scalar_t* __restrict__ input, int8_t* __restrict__ output,
131 + const scale_t* scale_ptr, const azp_t* azp_ptr, const int hidden_size) {
132 + const int tid = threadIdx.x;
133 + const int stride = blockDim.x;
134 + const int64_t token_idx = blockIdx.x;
135 + const float scale = *scale_ptr;
136 + const azp_t azp = *azp_ptr;
137 + const float inv_s = 1.0f / scale;
138
139 // Must be performed using 64-bit math to avoid integer overflow.
134 - out += token_idx * hidden_size;
135 - input += token_idx * hidden_size;
136 -
137 - for (int i = tid; i < hidden_size; i += blockDim.x) {
138 - auto const val = static_cast<float>(input[i]);
139 - auto const quant_val = int32_to_int8(float_to_int32_rn(val / scale) + azp);
140 - out[i] = quant_val;
141 - }
140 + const scalar_t* row_in = input + token_idx * hidden_size;
141 + int8_t* row_out = output + token_idx * hidden_size;
142 +
143 + vectorize_with_alignment<16>(
144 + row_in, row_out, hidden_size, tid, stride,
145 + [=] __device__(int8_t& dst, const scalar_t& src) {
146 + const auto v = static_cast<float>(src) * inv_s;
147 + dst = int32_to_int8(float_to_int32_rn(v) + azp);
148 + });
149 }
150
144 -template <typename scalar_t, typename scale_type>
151 +template <typename scalar_t, typename scale_t>
152 __global__ void dynamic_scaled_int8_quant_kernel(
146 - scalar_t const* __restrict__ input, int8_t* __restrict__ out,
147 - scale_type* scale, const int hidden_size) {
148 - int const tid = threadIdx.x;
149 - int64_t const token_idx = blockIdx.x;
150 - float absmax_val = 0.0f;
151 - float const zero = 0.0f;
153 + const scalar_t* __restrict__ input, int8_t* __restrict__ output,
154 + scale_t* scale_out, const int hidden_size) {
155 + const int tid = threadIdx.x;
156 + const int stride = blockDim.x;
157 + const int64_t token_idx = blockIdx.x;
158
159 // Must be performed using 64-bit math to avoid integer overflow.
154 - out += token_idx * hidden_size;
155 - input += token_idx * hidden_size;
156 -
157 - for (int i = tid; i < hidden_size; i += blockDim.x) {
158 - float val = static_cast<float>(input[i]);
159 - val = val > zero ? val : -val;
160 - absmax_val = val > absmax_val ? val : absmax_val;
160 + const scalar_t* row_in = input + token_idx * hidden_size;
161 + int8_t* row_out = output + token_idx * hidden_size;
162 +
163 + // calculate for absmax
164 + float thread_max = 0.f;
165 + for (int i = tid; i < hidden_size; i += stride) {
166 + const auto v = fabsf(static_cast<float>(row_in[i]));
167 + thread_max = fmaxf(thread_max, v);
168 }
162 -
163 - using BlockReduce = cub::BlockReduce<float, 1024>;
164 - __shared__ typename BlockReduce::TempStorage reduceStorage;
165 - float const block_absmax_val_maybe =
166 - BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
167 - __shared__ float block_absmax_val;
169 + using BlockReduce = cub::BlockReduce<float, 256>;
170 + __shared__ typename BlockReduce::TempStorage tmp;
171 + float block_max = BlockReduce(tmp).Reduce(thread_max, cub::Max{}, blockDim.x);
172 + __shared__ float absmax;
173 if (tid == 0) {
169 - block_absmax_val = block_absmax_val_maybe;
170 - scale[token_idx] = block_absmax_val / 127.0f;
174 + absmax = block_max;
175 + scale_out[blockIdx.x] = absmax / 127.f;
176 }
177 __syncthreads();
178
174 - float const tmp_scale = 127.0f / block_absmax_val;
175 - for (int i = tid; i < hidden_size; i += blockDim.x) {
176 - out[i] = float_to_int8_rn(static_cast<float>(input[i]) * tmp_scale);
179 + float inv_s = (absmax == 0.f) ? 0.f : 127.f / absmax;
180 +
181 + // 2. quantize
182 + vectorize_with_alignment<16>(
183 + row_in, row_out, hidden_size, tid, stride,
184 + [=] __device__(int8_t& dst, const scalar_t& src) {
185 + dst = float_to_int8_rn(static_cast<float>(src) * inv_s);
186 + });
187 +}
188 +
189 +// MinMax structure to hold min and max values in one go
190 +struct MinMax {
191 + float min, max;
192 +
193 + __host__ __device__ MinMax()
194 + : min(std::numeric_limits<float>::max()),
195 + max(std::numeric_limits<float>::lowest()) {}
196 +
197 + __host__ __device__ explicit MinMax(float v) : min(v), max(v) {}
198 +
199 + // add a value to the MinMax
200 + __host__ __device__ MinMax& operator+=(float v) {
201 + min = fminf(min, v);
202 + max = fmaxf(max, v);
203 + return *this;
204 }
205 +
206 + // merge two MinMax objects
207 + __host__ __device__ MinMax& operator&=(const MinMax& other) {
208 + min = fminf(min, other.min);
209 + max = fmaxf(max, other.max);
210 + return *this;
211 + }
212 +};
213 +
214 +__host__ __device__ inline MinMax operator+(MinMax a, float v) {
215 + return a += v;
216 +}
217 +__host__ __device__ inline MinMax operator&(MinMax a, const MinMax& b) {
218 + return a &= b;
219 }
220
180 -template <typename scalar_t, typename scale_type, typename azp_type>
221 +template <typename scalar_t, typename scale_t, typename azp_t>
222 __global__ void dynamic_scaled_int8_azp_quant_kernel(
182 - scalar_t const* __restrict__ input, int8_t* __restrict__ out,
183 - scale_type* scale, azp_type* azp, const int hidden_size) {
184 - int64_t const token_idx = blockIdx.x;
223 + const scalar_t* __restrict__ input, int8_t* __restrict__ output,
224 + scale_t* scale_out, azp_t* azp_out, const int hidden_size) {
225 + const int tid = threadIdx.x;
226 + const int stride = blockDim.x;
227 + const int64_t token_idx = blockIdx.x;
228
229 // Must be performed using 64-bit math to avoid integer overflow.
187 - out += token_idx * hidden_size;
188 - input += token_idx * hidden_size;
189 -
190 - // Scan for the min and max value for this token
191 - float max_val = std::numeric_limits<float>::min();
192 - float min_val = std::numeric_limits<float>::max();
193 - for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
194 - auto val = static_cast<float>(input[i]);
195 - max_val = std::max(max_val, val);
196 - min_val = std::min(min_val, val);
197 - }
230 + const scalar_t* row_in = input + token_idx * hidden_size;
231 + int8_t* row_out = output + token_idx * hidden_size;
232
199 - // Reduce the max and min values across the block
200 - using BlockReduce = cub::BlockReduce<float, 1024>;
201 - __shared__ typename BlockReduce::TempStorage reduceStorage;
202 - max_val = BlockReduce(reduceStorage).Reduce(max_val, cub::Max{}, blockDim.x);
203 - __syncthreads(); // Make sure min doesn't mess with max shared memory
204 - min_val = BlockReduce(reduceStorage).Reduce(min_val, cub::Min{}, blockDim.x);
205 -
206 - __shared__ scale_type scale_sh;
207 - __shared__ azp_type azp_sh;
208 -
209 - // Compute the scale and zero point and store them, only on the first thread
210 - if (threadIdx.x == 0) {
211 - float const scale_val = (max_val - min_val) / 255.0f;
212 - // Use rounding to even (same as torch.round)
213 - auto const azp_float = std::nearbyint(-128.0f - min_val / scale_val);
214 - auto const azp_val = static_cast<azp_type>(azp_float);
215 -
216 - // Store the scale and azp into shared and global
217 - scale[token_idx] = scale_sh = scale_val;
218 - azp[token_idx] = azp_sh = azp_val;
233 + // 1. calculate min & max
234 + MinMax thread_mm;
235 + for (int i = tid; i < hidden_size; i += stride) {
236 + thread_mm += static_cast<float>(row_in[i]);
237 }
238
221 - // Wait for the scale and azp to be computed
222 - __syncthreads();
239 + using BlockReduce = cub::BlockReduce<MinMax, 256>;
240 + __shared__ typename BlockReduce::TempStorage tmp;
241
224 - float const scale_val = scale_sh;
225 - azp_type const azp_val = azp_sh;
242 + MinMax mm = BlockReduce(tmp).Reduce(
243 + thread_mm,
244 + [] __device__(MinMax a, const MinMax& b) {
245 + a &= b;
246 + return a;
247 + },
248 + blockDim.x);
249
227 - // Quantize the values
228 - for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
229 - auto const val = static_cast<float>(input[i]);
230 - auto const quant_val =
231 - int32_to_int8(float_to_int32_rn(val / scale_val) + azp_val);
232 - out[i] = quant_val;
250 + __shared__ float scale_sh;
251 + __shared__ azp_t azp_sh;
252 + if (tid == 0) {
253 + float s = (mm.max - mm.min) / 255.f;
254 + float zp = nearbyintf(-128.f - mm.min / s); // round-to-even
255 + scale_sh = s;
256 + azp_sh = azp_t(zp);
257 + scale_out[blockIdx.x] = s;
258 + azp_out[blockIdx.x] = azp_sh;
259 }
260 + __syncthreads();
261 +
262 + const float inv_s = 1.f / scale_sh;
263 + const azp_t azp = azp_sh;
264 +
265 + // 2. quantize
266 + vectorize_with_alignment<16>(
267 + row_in, row_out, hidden_size, tid, stride,
268 + [=] __device__(int8_t& dst, const scalar_t& src) {
269 + const auto v = static_cast<float>(src) * inv_s;
270 + dst = int32_to_int8(float_to_int32_rn(v) + azp);
271 + });
272 }
273
274 } // namespace vllm
@@ -247,7 +285,7 @@ void static_scaled_int8_quant(torch::Tensor& out, // [..., hidden_size]
285 int const hidden_size = input.size(-1);
286 int const num_tokens = input.numel() / hidden_size;
287 dim3 const grid(num_tokens);
250 - dim3 const block(std::min(hidden_size, 1024));
288 + dim3 const block(std::min(hidden_size, 256));
289 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
290 VLLM_DISPATCH_FLOATING_TYPES(
291 input.scalar_type(), "static_scaled_int8_quant_kernel", [&] {
@@ -278,7 +316,7 @@ void dynamic_scaled_int8_quant(
316 int const hidden_size = input.size(-1);
317 int const num_tokens = input.numel() / hidden_size;
318 dim3 const grid(num_tokens);
281 - dim3 const block(std::min(hidden_size, 1024));
319 + dim3 const block(std::min(hidden_size, 256));
320 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
321 VLLM_DISPATCH_FLOATING_TYPES(
322 input.scalar_type(), "dynamic_scaled_int8_quant_kernel", [&] {
csrc/quantization/vectorization_utils.cuh new
+75
@@ -0,0 +1,75 @@
1 +#pragma once
2 +#include "vectorization.cuh"
3 +
4 +namespace vllm {
5 +
6 +template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
7 +struct DefaultVecOp {
8 + ScaOp scalar_op;
9 +
10 + __device__ __forceinline__ void operator()(
11 + vec_n_t<OutT, VEC_SIZE>& dst, const vec_n_t<InT, VEC_SIZE>& src) const {
12 +#pragma unroll
13 + for (int i = 0; i < VEC_SIZE; ++i) {
14 + scalar_op(dst.val[i], src.val[i]);
15 + }
16 + }
17 +};
18 +
19 +template <int VEC_SIZE, typename InT, typename OutT, typename VecOp,
20 + typename ScaOp>
21 +__device__ inline void vectorize_with_alignment(
22 + const InT* in, OutT* out, int len, int tid, int stride,
23 + VecOp&& vec_op, // vec_n_t<InT,16> -> vec_n_t<OutT,16>
24 + ScaOp&& scalar_op) { // InT -> OutT
25 + static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
26 + "VEC_SIZE must be a positive power-of-two");
27 + constexpr int WIDTH = VEC_SIZE * sizeof(InT); // eg: 64 B
28 + uintptr_t addr = reinterpret_cast<uintptr_t>(in);
29 +
30 + int misalignment_offset = addr & (WIDTH - 1); // addr % 64
31 + int alignment_bytes = WIDTH - misalignment_offset; // 64 - (addr % 64)
32 + int prefix_elems = alignment_bytes & (WIDTH - 1); // handle 64
33 + prefix_elems /= sizeof(InT);
34 + prefix_elems = min(prefix_elems, len); // 0 ≤ prefix < 16
35 +
36 + // 1. prefill the when it is unsafe to vectorize
37 + for (int i = tid; i < prefix_elems; i += stride) {
38 + scalar_op(out[i], in[i]);
39 + }
40 +
41 + in += prefix_elems;
42 + out += prefix_elems;
43 + len -= prefix_elems;
44 +
45 + int num_vec = len / VEC_SIZE;
46 + using vin_t = vec_n_t<InT, VEC_SIZE>;
47 + using vout_t = vec_n_t<OutT, VEC_SIZE>;
48 + auto* v_in = reinterpret_cast<const vin_t*>(in);
49 + auto* v_out = reinterpret_cast<vout_t*>(out);
50 +
51 + // 2. vectorize the main part
52 + for (int i = tid; i < num_vec; i += stride) {
53 + vout_t tmp;
54 + vec_op(tmp, v_in[i]);
55 + v_out[i] = tmp;
56 + }
57 +
58 + // 3. handle the tail
59 + int tail_start = num_vec * VEC_SIZE;
60 + for (int i = tid + tail_start; i < len; i += stride) {
61 + scalar_op(out[i], in[i]);
62 + }
63 +}
64 +
65 +template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
66 +__device__ __forceinline__ void vectorize_with_alignment(const InT* in,
67 + OutT* out, int len,
68 + int tid, int stride,
69 + ScaOp&& scalar_op) {
70 + using Vec = DefaultVecOp<VEC_SIZE, InT, OutT, std::decay_t<ScaOp>>;
71 + vectorize_with_alignment<VEC_SIZE>(in, out, len, tid, stride, Vec{scalar_op},
72 + std::forward<ScaOp>(scalar_op));
73 +}
74 +
75 +} // namespace vllm
tests/kernels/quantization/test_int8_quant.py
+1
@@ -11,6 +11,7 @@ from vllm.platforms import current_platform
11
12 DTYPES = [torch.half, torch.bfloat16, torch.float]
13 HIDDEN_SIZES = [16, 67, 768, 5137, 8193] # Arbitrary values for testing
14 +HIDDEN_SIZES += list(range(1024, 1033)) # vectorized conversion edge cases
15 NUM_TOKENS = [1, 7, 83, 4096] # Arbitrary values for testing
16 SEEDS = [0]
17 SCALE = [0.1, 2.1]