Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BugFix] [Kernel] Fix GPU SEGV occurring in int8 kernels #9391

Merged
merged 14 commits into from
Oct 17, 2024
Merged

[BugFix] [Kernel] Fix GPU SEGV occurring in int8 kernels #9391

Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
cc0abbc
fix int8 gpu segv
rasmith Oct 15, 2024
6cd8e38
use uint64_t and ensure iterator can fit in 32-bit register
rasmith Oct 15, 2024
ed62d62
revert custom_ops
rasmith Oct 15, 2024
9026085
main
rasmith Oct 15, 2024
e207343
add <algorithm> for upstream pr check
rasmith Oct 15, 2024
fcddad8
Make literal into long in std::min
rasmith Oct 15, 2024
d02a21b
change hidden_size back to int
rasmith Oct 16, 2024
b9df1da
change hidden_size back to int
rasmith Oct 16, 2024
c50ce42
change hidden_size back to int
rasmith Oct 16, 2024
f547c43
change hidden_size back to int
rasmith Oct 16, 2024
1e91784
add comment to clarify use of 64-bit math
rasmith Oct 16, 2024
e63904b
use current custom ops
rasmith Oct 16, 2024
6253fe9
clang format
rasmith Oct 16, 2024
fa5fd71
Merge branch 'vllm-project:main' into ransmith_int8_segv_fix
rasmith Oct 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 28 additions & 14 deletions csrc/quantization/compressed_tensors/int8_quant_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,12 +96,15 @@ __global__ void static_scaled_int8_quant_kernel(
scalar_t const* __restrict__ input, int8_t* __restrict__ out,
scale_type const* scale_ptr, const int hidden_size) {
int const tid = threadIdx.x;
int const token_idx = blockIdx.x;
int64_t const token_idx = blockIdx.x;
scale_type const scale = *scale_ptr;

// Must be performed using 64-bit math to avoid integer overflow.
out += token_idx * hidden_size;
mgoin marked this conversation as resolved.
Show resolved Hide resolved
input += token_idx * hidden_size;

for (int i = tid; i < hidden_size; i += blockDim.x) {
out[token_idx * hidden_size + i] = float_to_int8_rn(
static_cast<float>(input[token_idx * hidden_size + i]) / scale);
out[i] = float_to_int8_rn(static_cast<float>(input[i]) / scale);
}
}

Expand All @@ -111,14 +114,18 @@ __global__ void static_scaled_int8_azp_quant_kernel(
scale_type const* scale_ptr, azp_type const* azp_ptr,
const int hidden_size) {
int const tid = threadIdx.x;
int const token_idx = blockIdx.x;
int64_t const token_idx = blockIdx.x;
scale_type const scale = *scale_ptr;
azp_type const azp = *azp_ptr;

// Must be performed using 64-bit math to avoid integer overflow.
out += token_idx * hidden_size;
input += token_idx * hidden_size;

for (int i = tid; i < hidden_size; i += blockDim.x) {
auto const val = static_cast<float>(input[token_idx * hidden_size + i]);
auto const val = static_cast<float>(input[i]);
auto const quant_val = int32_to_int8(float_to_int32_rn(val / scale) + azp);
out[token_idx * hidden_size + i] = quant_val;
out[i] = quant_val;
}
}

Expand All @@ -127,12 +134,16 @@ __global__ void dynamic_scaled_int8_quant_kernel(
scalar_t const* __restrict__ input, int8_t* __restrict__ out,
scale_type* scale, const int hidden_size) {
int const tid = threadIdx.x;
int const token_idx = blockIdx.x;
int64_t const token_idx = blockIdx.x;
float absmax_val = 0.0f;
float const zero = 0.0f;

// Must be performed using 64-bit math to avoid integer overflow.
out += token_idx * hidden_size;
input += token_idx * hidden_size;

for (int i = tid; i < hidden_size; i += blockDim.x) {
float val = static_cast<float>(input[token_idx * hidden_size + i]);
float val = static_cast<float>(input[i]);
val = val > zero ? val : -val;
absmax_val = val > absmax_val ? val : absmax_val;
}
Expand All @@ -150,22 +161,25 @@ __global__ void dynamic_scaled_int8_quant_kernel(

float const tmp_scale = 127.0f / block_absmax_val;
for (int i = tid; i < hidden_size; i += blockDim.x) {
out[token_idx * hidden_size + i] = float_to_int8_rn(
static_cast<float>(input[token_idx * hidden_size + i]) * tmp_scale);
out[i] = float_to_int8_rn(static_cast<float>(input[i]) * tmp_scale);
}
}

template <typename scalar_t, typename scale_type, typename azp_type>
__global__ void dynamic_scaled_int8_azp_quant_kernel(
scalar_t const* __restrict__ input, int8_t* __restrict__ out,
scale_type* scale, azp_type* azp, const int hidden_size) {
int const token_idx = blockIdx.x;
int64_t const token_idx = blockIdx.x;

// Must be performed using 64-bit math to avoid integer overflow.
out += token_idx * hidden_size;
input += token_idx * hidden_size;

// Scan for the min and max value for this token
float max_val = std::numeric_limits<float>::min();
float min_val = std::numeric_limits<float>::max();
for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
auto val = static_cast<float>(input[token_idx * hidden_size + i]);
auto val = static_cast<float>(input[i]);
max_val = std::max(max_val, val);
min_val = std::min(min_val, val);
}
Expand Down Expand Up @@ -200,10 +214,10 @@ __global__ void dynamic_scaled_int8_azp_quant_kernel(

// Quantize the values
for (int i = threadIdx.x; i < hidden_size; i += blockDim.x) {
auto const val = static_cast<float>(input[token_idx * hidden_size + i]);
auto const val = static_cast<float>(input[i]);
auto const quant_val =
int32_to_int8(float_to_int32_rn(val / scale_val) + azp_val);
out[token_idx * hidden_size + i] = quant_val;
out[i] = quant_val;
}
}

Expand Down