Skip to content

Commit a0389c8

Browse files
jiqing-fengjiqing-feng
committed
fix lint
Signed-off-by: jiqing-feng <jiqing.feng@intel.com>
1 parent d8cbc68 commit a0389c8

4 files changed

Lines changed: 88 additions & 146 deletions

File tree

‎bitsandbytes/backends/cpu/ops.py‎

Lines changed: 0 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66

77
from bitsandbytes.functional import get_ptr
88

9-
from ..utils import CODE
109
from ..._ops import register_kernel
1110
from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib
1211

@@ -190,9 +189,6 @@ def _(
190189
ct.c_longlong(shape[0]),
191190
ct.c_longlong(shape[1]),
192191
)
193-
# out_2 = dequantize_nf4_test(_reverse_4bit_compress_format(A.reshape(-1)), absmax, blocksize, quant_type, shape, dtype)
194-
# if not torch.allclose(out, out_2, rtol=1e-2, atol=5e-2):
195-
# import pdb; pdb.set_trace()
196192
elif dtype == torch.float16:
197193
lib.cdequantize_blockwise_cpu_nf4_fp16(
198194
get_ptr(A),
@@ -206,40 +202,3 @@ def _(
206202
raise ValueError
207203

208204
return out
209-
210-
def dequantize_nf4_test(
211-
A: torch.Tensor,
212-
absmax: torch.Tensor,
213-
blocksize: int,
214-
quant_type: str,
215-
shape: Sequence[int],
216-
dtype: torch.dtype,
217-
):
218-
# Map nf4 to [-1, 1]
219-
out_dq = torch.empty(A.size(0) * 2, dtype=torch.int32, device=A.device)
220-
n = out_dq.numel()
221-
out_dq[1::2] = A & 0xF
222-
out_dq[::2] = A >> 4
223-
# code is fp32, cast to dtype to avoid the mismatch issue
224-
code = CODE[quant_type].to(dtype).to(A.device)
225-
out_dq = code[out_dq]
226-
227-
# Apply scales
228-
if out_dq.numel() != n:
229-
assert out_dq.numel() == n + 1
230-
out_dq = torch.narrow(out_dq, 0, 0, n)
231-
blocks = n // blocksize
232-
blocks += 1 if n % blocksize > 0 else 0
233-
rem = n % blocksize
234-
has_rem = rem > 0
235-
236-
if has_rem:
237-
out[: n - rem] = (out_dq[: n - rem].view(-1, blocksize) * absmax[: blocks - has_rem].view(-1, 1)).reshape(-1)
238-
out[n - rem :] = out_dq[n - rem :] * absmax[-1]
239-
else:
240-
out = out_dq.view(-1, blocksize) * absmax.view(-1, 1)
241-
242-
out = out.reshape(-1, *shape[1:]).to(dtype)
243-
244-
return out
245-

‎csrc/cpu_ops.cpp‎

Lines changed: 69 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -12,89 +12,61 @@
1212

1313
using namespace BinSearch;
1414

15-
1615
#if defined(__AVX512F__)
1716
#include <immintrin.h>
1817

1918
inline __m256i cvt_fp32_to_fp16(const __m512 src) {
2019
return _mm512_cvtps_ph(src, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
21-
}
20+
}
2221

2322
inline __m256i cvt_fp32_to_bf16(const __m512 src) {
24-
#if defined(__AVX512BF16__)
25-
return reinterpret_cast<__m256i>(_mm512_cvtneps_pbh(src));
26-
#else
27-
__m512i value = _mm512_castps_si512(src);
28-
__m512i nan = _mm512_set1_epi32(0xffff);
29-
auto mask_value = _mm512_cmp_ps_mask(src, src, _CMP_ORD_Q);
30-
__m512i ones = _mm512_set1_epi32(0x1);
31-
__m512i vec_bias = _mm512_set1_epi32(0x7fff);
32-
// uint32_t lsb = (input >> 16) & 1;
33-
auto t_value = _mm512_and_si512(_mm512_srli_epi32(value, 16), ones);
34-
// uint32_t rounding_bias = 0x7fff + lsb;
35-
t_value = _mm512_add_epi32(t_value, vec_bias);
36-
// input += rounding_bias;
37-
t_value = _mm512_add_epi32(t_value, value);
38-
// input = input >> 16;
39-
t_value = _mm512_srli_epi32(t_value, 16);
40-
// Check NaN before converting back to bf16
41-
t_value = _mm512_mask_blend_epi32(mask_value, nan, t_value);
42-
return _mm512_cvtusepi32_epi16(t_value);
43-
#endif
23+
#if defined(__AVX512BF16__)
24+
return reinterpret_cast<__m256i>(_mm512_cvtneps_pbh(src));
25+
#else
26+
__m512i value = _mm512_castps_si512(src);
27+
__m512i nan = _mm512_set1_epi32(0xffff);
28+
auto mask_value = _mm512_cmp_ps_mask(src, src, _CMP_ORD_Q);
29+
__m512i ones = _mm512_set1_epi32(0x1);
30+
__m512i vec_bias = _mm512_set1_epi32(0x7fff);
31+
// uint32_t lsb = (input >> 16) & 1;
32+
auto t_value = _mm512_and_si512(_mm512_srli_epi32(value, 16), ones);
33+
// uint32_t rounding_bias = 0x7fff + lsb;
34+
t_value = _mm512_add_epi32(t_value, vec_bias);
35+
// input += rounding_bias;
36+
t_value = _mm512_add_epi32(t_value, value);
37+
// input = input >> 16;
38+
t_value = _mm512_srli_epi32(t_value, 16);
39+
// Check NaN before converting back to bf16
40+
t_value = _mm512_mask_blend_epi32(mask_value, nan, t_value);
41+
return _mm512_cvtusepi32_epi16(t_value);
42+
#endif
4443
}
4544

4645
static inline __m512 set_nf4_lut() {
4746
return _mm512_set_ps(
48-
1.0f,
49-
0.7229568362236023,
50-
0.5626170039176941,
51-
0.44070982933044434,
52-
0.33791524171829224,
53-
0.24611230194568634,
54-
0.16093020141124725,
55-
0.07958029955625534,
56-
0.0f,
57-
-0.09105003625154495,
58-
-0.18477343022823334,
59-
-0.28444138169288635,
60-
-0.39491748809814453,
61-
-0.5250730514526367,
62-
-0.6961928009986877,
63-
-1.0f);
47+
1.0f, 0.7229568362236023, 0.5626170039176941, 0.44070982933044434, 0.33791524171829224, 0.24611230194568634,
48+
0.16093020141124725, 0.07958029955625534, 0.0f, -0.09105003625154495, -0.18477343022823334,
49+
-0.28444138169288635, -0.39491748809814453, -0.5250730514526367, -0.6961928009986877, -1.0f
50+
);
6451
}
52+
6553
static inline __m512 set_fp4_lut() {
6654
return _mm512_set_ps(
67-
-0.2500f,
68-
-0.16666667f,
69-
-0.5000f,
70-
-0.33333333f,
71-
-1.0000f,
72-
-0.66666667f,
73-
-5.208333333e-03f,
74-
0.0000f,
75-
0.2500f,
76-
0.16666667f,
77-
0.5000f,
78-
0.33333333f,
79-
1.0000f,
80-
0.66666667f,
81-
5.208333333e-03f,
82-
0.0000f);
55+
-0.2500f, -0.16666667f, -0.5000f, -0.33333333f, -1.0000f, -0.66666667f, -5.208333333e-03f, 0.0000f, 0.2500f,
56+
0.16666667f, 0.5000f, 0.33333333f, 1.0000f, 0.66666667f, 5.208333333e-03f, 0.0000f
57+
);
8358
}
8459
#endif
8560

8661
// 4-bit (FP4 / NF4) dequantization helper extracted from the original else branch.
87-
// DATA_TYPE: 1 = FP4, 0 = NF4
62+
// DATA_TYPE: 1 = FP4, 2 = NF4
8863
template <typename T, int DATA_TYPE>
89-
void dequantizeBlockwise4bitCpu(unsigned char* A,
90-
const float* absmax,
91-
T* out,
92-
long long blocksize,
93-
long long m,
94-
long long n) {
95-
static_assert(DATA_TYPE == 0 || DATA_TYPE == 1,
96-
"dequantizeBlockwise4bitCpu called with non 4-bit DATA_TYPE");
97-
if (blocksize <= 0 || m < 0 || n <= 0) return;
64+
void dequantizeBlockwise4bitCpu(
65+
unsigned char* A, const float* absmax, T* out, long long blocksize, long long m, long long n
66+
) {
67+
static_assert(DATA_TYPE == 1 || DATA_TYPE == 2, "dequantizeBlockwise4bitCpu called with non 4-bit DATA_TYPE");
68+
if (blocksize <= 0 || m < 0 || n <= 0)
69+
return;
9870

9971
#if defined(__AVX512F__)
10072
long long dim_0 = m;
@@ -119,10 +91,10 @@ void dequantizeBlockwise4bitCpu(unsigned char* A,
11991
uint64_t high = 0;
12092
uint64_t low = 0;
12193
for (int i = 0; i < 4; ++i) {
122-
low |= ((packed >> (2*i * 4)) & 0xf) << ((2*i+1) * 8);
123-
low |= ((packed >> ((2*i+1) * 4)) & 0xf) << (2*i * 8);
124-
high |= ((packed >> (2*i * 4 + 32)) & 0xf) << ((2*i+1) * 8);
125-
high |= ((packed >> ((2*i+1) * 4 + 32)) & 0xf) << (2*i * 8);
94+
low |= ((packed >> (2 * i * 4)) & 0xf) << ((2 * i + 1) * 8);
95+
low |= ((packed >> ((2 * i + 1) * 4)) & 0xf) << (2 * i * 8);
96+
high |= ((packed >> (2 * i * 4 + 32)) & 0xf) << ((2 * i + 1) * 8);
97+
high |= ((packed >> ((2 * i + 1) * 4 + 32)) & 0xf) << (2 * i * 8);
12698
}
12799
__m128i packed_128 = _mm_set_epi64x(high, low);
128100
__m512i vint32 = _mm512_cvtepu8_epi32(packed_128);
@@ -133,13 +105,11 @@ void dequantizeBlockwise4bitCpu(unsigned char* A,
133105
// Store results
134106
T* pout = &out[block_idx * dim_1 + k * 2];
135107
if constexpr (std::is_same<T, float>()) {
136-
_mm512_storeu_ps(pout, vout);
108+
_mm512_storeu_ps(pout, vout);
137109
} else if constexpr (std::is_same<T, bf16_t>()) {
138-
_mm256_storeu_si256(
139-
(__m256i*)pout, cvt_fp32_to_bf16(vout));
110+
_mm256_storeu_si256((__m256i*)pout, cvt_fp32_to_bf16(vout));
140111
} else if constexpr (std::is_same<T, fp16_t>()) {
141-
_mm256_storeu_si256(
142-
(__m256i*)pout, cvt_fp32_to_fp16(vout));
112+
_mm256_storeu_si256((__m256i*)pout, cvt_fp32_to_fp16(vout));
143113
}
144114
}
145115
}
@@ -157,11 +127,9 @@ void dequantizeBlockwise4bitCpu(unsigned char* A,
157127
unsigned char byte = A[byte_index];
158128

159129
// High nibble first (matches previous code logic)
160-
float v0 = (DATA_TYPE == 1 ? dDequantizeFP4(byte >> 4)
161-
: dDequantizeNF4(byte >> 4)) * scale;
130+
float v0 = (DATA_TYPE == 1 ? dDequantizeFP4(byte >> 4) : dDequantizeNF4(byte >> 4)) * scale;
162131
// Low nibble second
163-
float v1 = (DATA_TYPE == 1 ? dDequantizeFP4(byte & 0x0F)
164-
: dDequantizeNF4(byte & 0x0F)) * scale;
132+
float v1 = (DATA_TYPE == 1 ? dDequantizeFP4(byte & 0x0F) : dDequantizeNF4(byte & 0x0F)) * scale;
165133

166134
if constexpr (std::is_same<T, bf16_t>::value) {
167135
out[block_idx + i] = float_to_bf16(v0);
@@ -184,20 +152,17 @@ void dequantizeBlockwise4bitCpu(unsigned char* A,
184152
}
185153
}
186154

187-
188155
template <typename T>
189-
void dequantizeBlockwise8bitCpu(float* code,
190-
unsigned char* A,
191-
const float* absmax,
192-
T* out,
193-
long long blocksize,
194-
long long n) {
195-
if (blocksize <= 0 || n <= 0) return;
156+
void dequantizeBlockwise8bitCpu(
157+
float* code, unsigned char* A, const float* absmax, T* out, long long blocksize, long long n
158+
) {
159+
if (blocksize <= 0 || n <= 0)
160+
return;
196161
// 8-bit path
197162
BNB_OMP_PARALLEL_FOR
198163
for (long long block_idx = 0; block_idx < n; block_idx += blocksize) {
199164
long long valid_items = (n - block_idx >= blocksize ? blocksize : n - block_idx);
200-
long long block_end = block_idx + valid_items;
165+
long long block_end = block_idx + valid_items;
201166
float scale = absmax[block_idx / blocksize];
202167
for (long long i = block_idx; i < block_end; ++i) {
203168
float v = code[A[i]] * scale;
@@ -212,7 +177,6 @@ void dequantizeBlockwise8bitCpu(float* code,
212177
}
213178
}
214179

215-
216180
void quantize_cpu(float* code, float* A, float* absmax, unsigned char* out, long long blocksize, long long n) {
217181

218182
// the default code is has range [-0.993, 1.0] which can cause an error in the binary search algorithm used below
@@ -265,26 +229,35 @@ void quantize_cpu(float* code, float* A, float* absmax, unsigned char* out, long
265229
//==============================================================
266230

267231
template void dequantizeBlockwise8bitCpu<float>(
268-
float* code, unsigned char* A, const float* absmax, float* out, long long blocksize, long long n);
232+
float* code, unsigned char* A, const float* absmax, float* out, long long blocksize, long long n
233+
);
269234
template void dequantizeBlockwise8bitCpu<fp16_t>(
270-
float* code, unsigned char* A, const float* absmax, fp16_t* out, long long blocksize, long long n);
235+
float* code, unsigned char* A, const float* absmax, fp16_t* out, long long blocksize, long long n
236+
);
271237
template void dequantizeBlockwise8bitCpu<bf16_t>(
272-
float* code, unsigned char* A, const float* absmax, bf16_t* out, long long blocksize, long long n);
238+
float* code, unsigned char* A, const float* absmax, bf16_t* out, long long blocksize, long long n
239+
);
273240

274241
template void dequantizeBlockwise4bitCpu<float, FP4>(
275-
unsigned char* A, const float* absmax, float* out, long long blocksize, long long m, long long n);
242+
unsigned char* A, const float* absmax, float* out, long long blocksize, long long m, long long n
243+
);
276244
template void dequantizeBlockwise4bitCpu<float, NF4>(
277-
unsigned char* A, const float* absmax, float* out, long long blocksize, long long m, long long n);
245+
unsigned char* A, const float* absmax, float* out, long long blocksize, long long m, long long n
246+
);
278247

279248
template void dequantizeBlockwise4bitCpu<fp16_t, FP4>(
280-
unsigned char* A, const float* absmax, fp16_t* out, long long blocksize, long long m, long long n);
249+
unsigned char* A, const float* absmax, fp16_t* out, long long blocksize, long long m, long long n
250+
);
281251
template void dequantizeBlockwise4bitCpu<fp16_t, NF4>(
282-
unsigned char* A, const float* absmax, fp16_t* out, long long blocksize, long long m, long long n);
252+
unsigned char* A, const float* absmax, fp16_t* out, long long blocksize, long long m, long long n
253+
);
283254

284255
template void dequantizeBlockwise4bitCpu<bf16_t, FP4>(
285-
unsigned char* A, const float* absmax, bf16_t* out, long long blocksize, long long m, long long n);
256+
unsigned char* A, const float* absmax, bf16_t* out, long long blocksize, long long m, long long n
257+
);
286258
template void dequantizeBlockwise4bitCpu<bf16_t, NF4>(
287-
unsigned char* A, const float* absmax, bf16_t* out, long long blocksize, long long m, long long n);
259+
unsigned char* A, const float* absmax, bf16_t* out, long long blocksize, long long m, long long n
260+
);
288261

289262
// template void gemv_4bit_inference<fp16_t, 16>(
290263
// int m, int n, int k, fp16_t* A, unsigned char* B, float* absmax, float* datatype, fp16_t* out,

‎csrc/cpu_ops.h‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,9 @@
77
void quantize_cpu(float* code, float* A, float* absmax, unsigned char* out, long long blocksize, long long n);
88

99
typedef enum DataType_t {
10-
NF4 = 0,
10+
General8bit = 0,
1111
FP4 = 1,
12+
NF4 = 2,
1213
} DataType_t;
1314

1415
struct fp16_t {
@@ -30,17 +31,17 @@ static inline fp16_t float_to_fp16(float x) {
3031
uint32_t bits;
3132
std::memcpy(&bits, &x, 4);
3233
uint32_t sign = (bits >> 31) & 0x1;
33-
uint32_t exp = (bits >> 23) & 0xFF;
34+
uint32_t exp = (bits >> 23) & 0xFF;
3435
uint32_t mant = bits & 0x7FFFFF;
3536

3637
uint16_t h;
37-
if (exp == 0xFF) { // Inf / NaN
38+
if (exp == 0xFF) { // Inf / NaN
3839
uint16_t mant16 = mant ? 0x200 : 0; // quiet NaN: set MSB of mantissa
3940
h = (sign << 15) | (0x1F << 10) | mant16;
40-
} else if (exp > 0x70 + 0x1E) { // overflow: exp_f -127 +15 > 30 (exp_f > 142)
41+
} else if (exp > 0x70 + 0x1E) { // overflow: exp_f -127 +15 > 30 (exp_f > 142)
4142
h = (sign << 15) | (0x1F << 10); // Inf
42-
} else if (exp < 0x71) { // subnormal or zero (exp_f < 113)
43-
if (exp < 0x67) { // too small -> zero (exp_f < 103)
43+
} else if (exp < 0x71) { // subnormal or zero (exp_f < 113)
44+
if (exp < 0x67) { // too small -> zero (exp_f < 103)
4445
h = (sign << 15);
4546
} else {
4647
// subnormal: implicit leading 1
@@ -156,11 +157,14 @@ inline float dDequantizeNF4(unsigned char val) {
156157
return -1.0f; //*0000
157158
}
158159

159-
160160
template <typename T>
161-
void dequantizeBlockwise8bitCpu(float* code, unsigned char* A, const float* absmax, T* out, long long blocksize, long long n);
161+
void dequantizeBlockwise8bitCpu(
162+
float* code, unsigned char* A, const float* absmax, T* out, long long blocksize, long long n
163+
);
162164

163165
template <typename T, int DATA_TYPE>
164-
void dequantizeBlockwise4bitCpu(unsigned char* A, const float* absmax, T* out, long long blocksize, long long m, long long n);
166+
void dequantizeBlockwise4bitCpu(
167+
unsigned char* A, const float* absmax, T* out, long long blocksize, long long m, long long n
168+
);
165169

166170
#endif

0 commit comments

Comments
 (0)