Skip to content

Commit 1be1dc8

Browse files
Xia-WeiwenXia-Weiwen
authored
WOQ: enable lowp-mode=INT8 for NF4 weight (#3395)
* WOQ: enable lowp-mode=INT8 for NF4 weight * Remove -mavx512vbmi for avx512 * Remove -mavx512vbmi for avx512_vnni and avx512_bf16
1 parent 6b4d80b commit 1be1dc8

File tree

13 files changed

+217
-55
lines changed

13 files changed

+217
-55
lines changed

‎cmake/Modules/FindAVX.cmake

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -186,14 +186,14 @@ CHECK_SSE(CXX "AVX512_BF16" " ;-mavx512f -mavx512dq -mavx512vl -mavx512bw -mavx5
186186

187187
# gcc start to support amx from version 11.2
188188
# https://gcc.gnu.org/onlinedocs/gcc-11.2.0/gcc/x86-Options.html#x86-Options
189-
CHECK_SSE(C "AMX" " ;-mavx512f -mavx512dq -mavx512vl -mavx512bw -mavx512bf16 -mfma\
189+
CHECK_SSE(C "AMX" " ;-mavx512f -mavx512dq -mavx512vbmi -mavx512vl -mavx512bw -mavx512bf16 -mfma\
190190
-mamx-tile -mamx-int8 -mamx-bf16;/arch:AVX512")
191-
CHECK_SSE(CXX "AMX" " ;-mavx512f -mavx512dq -mavx512vl -mavx512bw -mavx512bf16 -mfma\
191+
CHECK_SSE(CXX "AMX" " ;-mavx512f -mavx512dq -mavx512vbmi -mavx512vl -mavx512bw -mavx512bf16 -mfma\
192192
-mamx-tile -mamx-int8 -mamx-bf16;/arch:AVX512")
193193

194194
# gcc starts to support avx512fp16 from version 12.1
195195
# https://gcc.gnu.org/onlinedocs/gcc-12.1.0/gcc/x86-Options.html#x86-Options
196-
CHECK_SSE(C "AVX512_FP16" " ;-mavx512f -mavx512dq -mavx512vl -mavx512bw -mavx512bf16 -mfma\
196+
CHECK_SSE(C "AVX512_FP16" " ;-mavx512f -mavx512dq -mavx512vbmi -mavx512vl -mavx512bw -mavx512bf16 -mfma\
197197
-mamx-tile -mamx-int8 -mamx-bf16 -mavx512fp16;/arch:AVX512")
198-
CHECK_SSE(CXX "AVX512_FP16" " ;-mavx512f -mavx512dq -mavx512vl -mavx512bw -mavx512bf16 -mfma\
198+
CHECK_SSE(CXX "AVX512_FP16" " ;-mavx512f -mavx512dq -mavx512vbmi -mavx512vl -mavx512bw -mavx512bf16 -mfma\
199199
-mamx-tile -mamx-int8 -mamx-bf16 -mavx512fp16;/arch:AVX512")

‎cmake/cpu/IsaCodegen.cmake

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ if(CXX_AVX512_FP16_FOUND)
4444
else(MSVC)
4545
list(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG} -D__AVX512F__ ${AVX512_OPTIMIZE_FLAGS} -DCPU_CAPABILITY_AVX512 \
4646
-DCPU_CAPABILITY_AVX512_VNNI -DCPU_CAPABILITY_AVX512_BF16 -DCPU_CAPABILITY_AMX \
47-
-mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vnni \
47+
-mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vbmi -mavx512vnni \
4848
-mavx512bf16 -mfma -mamx-tile -mamx-int8 -mamx-bf16 -mavx512fp16")
4949
endif(MSVC)
5050
else(CXX_AVX512_FP16_FOUND)
@@ -60,7 +60,7 @@ if(CXX_AMX_FOUND)
6060
list(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG}/arch:AVX512") # TODO: CHECK HERE
6161
else(MSVC)
6262
list(APPEND CPU_CAPABILITY_FLAGS "${OPT_FLAG} -D__AVX512F__ ${AVX512_OPTIMIZE_FLAGS} -DCPU_CAPABILITY_AVX512 \
63-
-DCPU_CAPABILITY_AVX512_VNNI -DCPU_CAPABILITY_AVX512_BF16 -mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vnni -mavx512bf16 -mfma \
63+
-DCPU_CAPABILITY_AVX512_VNNI -DCPU_CAPABILITY_AVX512_BF16 -mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vbmi -mavx512vnni -mavx512bf16 -mfma \
6464
-mamx-tile -mamx-int8 -mamx-bf16")
6565
endif(MSVC)
6666
else(CXX_AMX_FOUND)

‎csrc/cpu/aten/Linear.cpp

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -363,7 +363,7 @@ static size_t get_block_k(
363363
int64_t lowp_mode,
364364
int64_t group_size,
365365
int64_t K) {
366-
size_t default_block_k = lowp_mode == 3 ? 128 : 64;
366+
size_t default_block_k = lowp_mode == LOWP_MODE_INT8 ? 128 : 64;
367367
size_t block_k = group_size > 0
368368
? std::min((size_t)group_size, default_block_k)
369369
: default_block_k;
@@ -389,7 +389,7 @@ at::Tensor woq_linear_pack_weight(
389389
size_t block_n = WOQ_N_BLOCK_SIZE;
390390
size_t block_k = get_block_k(weight_dtype, lowp_mode, group_size, K);
391391
if (weight_dtype == WOQ_DTYPE_INT4 || weight_dtype == WOQ_DTYPE_NF4) {
392-
if (block_k % 4 && lowp_mode == 3) {
392+
if (block_k % 4 && lowp_mode == LOWP_MODE_INT8) {
393393
// This case is not supported by kernel
394394
return weight;
395395
}
@@ -428,7 +428,7 @@ at::Tensor woq_linear_compute_compensation(
428428
TORCH_CHECK(weight.dim() == 2);
429429
auto N = weight.size(0), K = weight.size(1);
430430
if (N % WOQ_N_BLOCK_SIZE == 0 && weight_dtype == WOQ_DTYPE_INT8 &&
431-
lowp_mode == 3) {
431+
lowp_mode == LOWP_MODE_INT8) {
432432
size_t block_k = get_block_k(weight_dtype, lowp_mode, group_size, K);
433433
int64_t Nc = N / WOQ_N_BLOCK_SIZE, Kc = K / block_k;
434434
auto weight_reshaped = weight.reshape({Nc, WOQ_N_BLOCK_SIZE, Kc, block_k});
@@ -445,7 +445,7 @@ at::Tensor woq_linear_unpack_weight(
445445
const at::Tensor& weight,
446446
int64_t weight_dtype,
447447
int64_t lowp_mode) {
448-
if (weight_dtype == WOQ_DTYPE_INT8 && lowp_mode == 3) {
448+
if (weight_dtype == WOQ_DTYPE_INT8 && lowp_mode == LOWP_MODE_INT8) {
449449
// Unpack weight for INT8 GEMM.
450450
// weight is packed in 5d (Nc, Kc, block_k / 4, block_n, 4)
451451
// but viewd as 4d (Nc, Kc, block_k, block_n)
@@ -553,7 +553,8 @@ at::Tensor woq_linear_forward_v2(
553553
WOQ_DTYPE_MAP.find(weight_dtype) != WOQ_DTYPE_MAP.end(),
554554
"Unsupported weight dtype: ",
555555
weight_dtype);
556-
if (WOQ_DTYPE_MAP.at(weight_dtype) == WOQ_DTYPE_INT8 && lowp_mode == 3) {
556+
if (WOQ_DTYPE_MAP.at(weight_dtype) == WOQ_DTYPE_INT8 &&
557+
lowp_mode == LOWP_MODE_INT8) {
557558
TORCH_CHECK(compensation.has_value() && compensation.value().defined());
558559
}
559560
static const at::Tensor empty_tensor = at::Tensor();

‎csrc/cpu/aten/kernels/WoqInt8GemmAPerKBlockKrnl.cpp

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,9 +98,6 @@ at::Tensor woq_gemm_int8(
9898
auto quant_w_mode_ = std::get<1>(tuple);
9999
using act_type =
100100
typename c10::impl::ScalarTypeToCPPType<act_dtype>::type;
101-
TLA_ASSERT(
102-
qw_type == WOQ_DTYPE_INT4 || qw_type == WOQ_DTYPE_INT8,
103-
"LOWP_MODE_INT8 only support qw_type = INT4 or INT8");
104101
auto block_k = w_sizes[2];
105102
if (quant_block_k <= 0)
106103
quant_block_k = block_k;

‎csrc/cpu/aten/kernels/WoqInt8GemmAPerMKBlockKrnl.cpp

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,9 +98,6 @@ at::Tensor woq_gemm_int8(
9898
auto quant_w_mode_ = std::get<1>(tuple);
9999
using act_type =
100100
typename c10::impl::ScalarTypeToCPPType<act_dtype>::type;
101-
TLA_ASSERT(
102-
qw_type == WOQ_DTYPE_INT4 || qw_type == WOQ_DTYPE_INT8,
103-
"LOWP_MODE_INT8 only support qw_type = INT4 or INT8");
104101
auto block_k = w_sizes[2];
105102
if (quant_block_k <= 0)
106103
quant_block_k = block_k;

‎csrc/cpu/aten/kernels/WoqInt8GemmAPerMKrnl.cpp

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,9 +98,6 @@ at::Tensor woq_gemm_int8(
9898
auto quant_w_mode_ = std::get<1>(tuple);
9999
using act_type =
100100
typename c10::impl::ScalarTypeToCPPType<act_dtype>::type;
101-
TLA_ASSERT(
102-
qw_type == WOQ_DTYPE_INT4 || qw_type == WOQ_DTYPE_INT8,
103-
"LOWP_MODE_INT8 only support qw_type = INT4 or INT8");
104101
auto block_k = w_sizes[2];
105102
if (quant_block_k <= 0)
106103
quant_block_k = block_k;

‎csrc/cpu/aten/kernels/WoqInt8GemmAPerTensorKrnl.cpp

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,9 +98,6 @@ at::Tensor woq_gemm_int8(
9898
auto quant_w_mode_ = std::get<1>(tuple);
9999
using act_type =
100100
typename c10::impl::ScalarTypeToCPPType<act_dtype>::type;
101-
TLA_ASSERT(
102-
qw_type == WOQ_DTYPE_INT4 || qw_type == WOQ_DTYPE_INT8,
103-
"LOWP_MODE_INT8 only support qw_type = INT4 or INT8");
104101
float scale_a;
105102
int32_t zp_a;
106103
bool is_sym_quant = !is_asymmetric_quant_a(quant_a_mode);

‎csrc/cpu/aten/utils/woq.h

Lines changed: 75 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -525,6 +525,51 @@ inline std::array<__m256i, 2> load_sint4_as_int8(uint8_t* qB) {
525525
return {low, high};
526526
}
527527

528+
// load nf4
529+
inline std::array<__m256i, 2> load_nf4_as_int8(uint8_t* qB) {
530+
__m256i packed = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(qB));
531+
const __m256i low_mask = _mm256_set1_epi8(0x0f);
532+
__m256i high = _mm256_srli_epi16(packed, 4);
533+
high = _mm256_and_si256(high, low_mask);
534+
__m256i low = _mm256_and_si256(packed, low_mask);
535+
const __m256i lut = _mm256_set_epi8(
536+
0,
537+
0,
538+
0,
539+
0,
540+
0,
541+
0,
542+
0,
543+
0,
544+
0,
545+
0,
546+
0,
547+
0,
548+
0,
549+
0,
550+
0,
551+
0,
552+
127,
553+
92,
554+
71,
555+
56,
556+
43,
557+
31,
558+
20,
559+
10,
560+
0,
561+
-12,
562+
-23,
563+
-36,
564+
-50,
565+
-67,
566+
-88,
567+
-127);
568+
low = _mm256_permutexvar_epi8(low, lut);
569+
high = _mm256_permutexvar_epi8(high, lut);
570+
return {low, high};
571+
}
572+
528573
#else
529574
inline std::array<__m256i, 2> load_zps_4vnni(int8_t* zps) {
530575
TLA_ASSERT(false, "not implemented");
@@ -541,6 +586,11 @@ inline std::array<__m256i, 2> load_sint4_as_int8(uint8_t* qB) {
541586
return std::array<__m256i, 2>();
542587
}
543588

589+
inline std::array<__m256i, 2> load_nf4_as_int8(uint8_t* qB) {
590+
TLA_ASSERT(false, "not implemented");
591+
return std::array<__m256i, 2>();
592+
}
593+
544594
#endif
545595

546596
template <long N, bool sym_quant, typename T>
@@ -831,6 +881,10 @@ struct GemmMicroKernel<
831881
// Load scales and zps
832882
compile_time_for<COLS>::op([&](auto i) {
833883
vscales[i] = _mm512_loadu_ps(scales + i * 16);
884+
if constexpr (qw_type == WOQ_DTYPE_NF4) {
885+
const __m512 factor = _mm512_set1_ps(1.0f / 127.0f);
886+
vscales[i] = _mm512_mul_ps(vscales[i], factor);
887+
}
834888
// TODO(jgong5): should we use 512 or two 256 here?
835889
if constexpr (!sym_quant_w) {
836890
vzps[i] = combine_m256i(load_zps_4vnni(zps + i * 16));
@@ -859,8 +913,10 @@ struct GemmMicroKernel<
859913
if constexpr (!sym_quant_w) {
860914
vb[col] = combine_m256i(load_uint4_as_int8(pqB[k / 4][col * 16]));
861915
vb[col] = _mm512_sub_epi8(vb[col], vzps[col]);
862-
} else {
916+
} else if constexpr (qw_type == WOQ_DTYPE_INT4) {
863917
vb[col] = combine_m256i(load_sint4_as_int8(pqB[k / 4][col * 16]));
918+
} else {
919+
vb[col] = combine_m256i(load_nf4_as_int8(pqB[k / 4][col * 16]));
864920
}
865921
if constexpr (is_asymmetric_quant_a(quant_a_mode)) {
866922
vcompensate[col] =
@@ -1290,12 +1346,12 @@ struct Dequantize<half, ldb, N_GROUP_SIZE, qw_type, sym_quant_w> {
12901346
}
12911347
};
12921348

1293-
template <long ldb, bool sym_quant_w>
1349+
template <long ldb, int qw_type, bool sym_quant_w>
12941350
struct Dequantize<
12951351
int8_t,
12961352
ldb,
12971353
/*N_GROUP_SIZE*/ 16,
1298-
/*qw_type*/ WOQ_DTYPE_INT4,
1354+
qw_type,
12991355
sym_quant_w> {
13001356
template <int quant_a_mode>
13011357
static inline void call(
@@ -1330,10 +1386,14 @@ struct Dequantize<
13301386
auto [low, high] = load_uint4_as_int8(pqB[k][n]);
13311387
vb_high = _mm256_sub_epi8(high, vzps_high);
13321388
vb_low = _mm256_sub_epi8(low, vzps_low);
1333-
} else {
1389+
} else if constexpr (qw_type == WOQ_DTYPE_INT4) {
13341390
auto [low, high] = load_sint4_as_int8(pqB[k][n]);
13351391
vb_low = low;
13361392
vb_high = high;
1393+
} else {
1394+
auto [low, high] = load_nf4_as_int8(pqB[k][n]);
1395+
vb_low = low;
1396+
vb_high = high;
13371397
}
13381398
if constexpr (is_asymmetric_quant_a(quant_a_mode)) {
13391399
vcompensate[0] = _mm256_dpbusd_epi32(vcompensate[0], ones, vb_low);
@@ -1585,6 +1645,7 @@ template <
15851645
long ldb,
15861646
bool transA,
15871647
bool ACC,
1648+
int qw_type,
15881649
int quant_a_mode,
15891650
int quant_w_mode,
15901651
long PREFETCH_K_DIST>
@@ -1598,7 +1659,7 @@ class DequantGemmTPP<
15981659
ldb,
15991660
transA,
16001661
ACC,
1601-
/*qw_type*/ WOQ_DTYPE_INT4,
1662+
qw_type,
16021663
quant_a_mode,
16031664
quant_w_mode,
16041665
PREFETCH_K_DIST> {
@@ -1696,7 +1757,7 @@ class DequantGemmTPP<
16961757
ACC,
16971758
quant_a_mode,
16981759
PREFETCH_K_DIST>::
1699-
template call<WOQ_DTYPE_INT4, sym_quant_w>(
1760+
template call<qw_type, sym_quant_w>(
17001761
K,
17011762
qA[m],
17021763
lda,
@@ -1725,7 +1786,7 @@ class DequantGemmTPP<
17251786
ACC,
17261787
quant_a_mode,
17271788
PREFETCH_K_DIST>::
1728-
template call<WOQ_DTYPE_INT4, sym_quant_w>(
1789+
template call<qw_type, sym_quant_w>(
17291790
K,
17301791
qA[m],
17311792
lda,
@@ -1748,12 +1809,7 @@ class DequantGemmTPP<
17481809
int8_t B[K / 4][N][4];
17491810
int32_t qC[M][N];
17501811
int32_t compensation[N];
1751-
Dequantize<
1752-
int8_t,
1753-
ldb,
1754-
N_GROUP_SIZE,
1755-
/*qw_type*/ WOQ_DTYPE_INT4,
1756-
sym_quant_w>::
1812+
Dequantize<int8_t, ldb, N_GROUP_SIZE, qw_type, sym_quant_w>::
17571813
template call<quant_a_mode>(qB, K, N, zps, B[0][0], compensation);
17581814
(*pgemm)((int8_t*)qA[0], B[0][0], qC[0], 1, no_tile_cfg);
17591815
if constexpr (PREFETCH_K_DIST > 0) {
@@ -1782,11 +1838,14 @@ class DequantGemmTPP<
17821838
}
17831839
}
17841840
float c = 0;
1841+
auto scale = scales[n];
1842+
if constexpr (qw_type == WOQ_DTYPE_NF4) {
1843+
scale *= (1.0f / 127.0f);
1844+
}
17851845
if constexpr (is_asymmetric_quant_a(quant_a_mode)) {
1786-
c = (qC[m][n] - compensation[n] * (*zp_a_m)) * (*scale_a_m) *
1787-
scales[n];
1846+
c = (qC[m][n] - compensation[n] * (*zp_a_m)) * (*scale_a_m) * scale;
17881847
} else {
1789-
c = (qC[m][n]) * (*scale_a_m) * scales[n];
1848+
c = (qC[m][n]) * (*scale_a_m) * scale;
17901849
}
17911850
if constexpr (ACC) {
17921851
C[m * ldc + n] += c;

‎csrc/cpu/jit/cpu/kernels/LinearWoqPacked.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -362,7 +362,7 @@ ContextLinearWoq create(
362362
lowp_mode,
363363
act_quant_mode,
364364
cache_weight_for_large_batch);
365-
if (weight_dtype == WOQ_DTYPE_INT8 && lowp_mode == 3) {
365+
if (weight_dtype == WOQ_DTYPE_INT8 && lowp_mode == LOWP_MODE_INT8) {
366366
auto compensation = woq_linear_compute_compensation(
367367
weight, weight_dtype, group_size, lowp_mode);
368368
context.cached_compensation_ =

‎docs/tutorials/features/isa_dynamic_dispatch.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,9 @@ The CodeGen will copy each cpp files from **Kernel implementation**, and then ad
7474
>
7575
> AVX512_BF16: `build/Release/csrc/isa_codegen/cpu/aten/kernels/AdaptiveAveragePoolingKrnl.cpp.AVX512_BF16.cpp -O3 -D__AVX512F__ -DCPU_CAPABILITY_AVX512 -DCPU_CAPABILITY_AVX512_VNNI -mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vnni -mavx512bf16 -mfma -DCPU_CAPABILITY=AVX512_BF16 -DCPU_CAPABILITY_AVX512_BF16`
7676
>
77-
> AMX: `build/Release/csrc/isa_codegen/cpu/aten/kernels/AdaptiveAveragePoolingKrnl.cpp.AMX.cpp -O3 -D__AVX512F__ -DCPU_CAPABILITY_AVX512 -DCPU_CAPABILITY_AVX512_VNNI -DCPU_CAPABILITY_AVX512_BF16 -mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vnni -mavx512bf16 -mfma -mamx-tile -mamx-int8 -mamx-bf16 -DCPU_CAPABILITY=AMX -DCPU_CAPABILITY_AMX`
77+
> AMX: `build/Release/csrc/isa_codegen/cpu/aten/kernels/AdaptiveAveragePoolingKrnl.cpp.AMX.cpp -O3 -D__AVX512F__ -DCPU_CAPABILITY_AVX512 -DCPU_CAPABILITY_AVX512_VNNI -DCPU_CAPABILITY_AVX512_BF16 -mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vbmi -mavx512vnni -mavx512bf16 -mfma -mamx-tile -mamx-int8 -mamx-bf16 -DCPU_CAPABILITY=AMX -DCPU_CAPABILITY_AMX`
7878
>
79-
> AVX512_FP16: `build/Release/csrc/isa_codegen/cpu/aten/kernels/AdaptiveAveragePoolingKrnl.cpp.AVX512_FP16.cpp -O3 -D__AVX512F__ -DCPU_CAPABILITY_AVX512 -DCPU_CAPABILITY_AVX512_VNNI -DCPU_CAPABILITY_AVX512_BF16 -mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vnni -mavx512bf16 -mfma -mamx-tile -mamx-int8 -mamx-bf16 -mavx512fp16 -DCPU_CAPABILITY_AMX -DCPU_CAPABILITY=AVX512_FP16 -DCPU_CAPABILITY_AVX512_FP16`
79+
> AVX512_FP16: `build/Release/csrc/isa_codegen/cpu/aten/kernels/AdaptiveAveragePoolingKrnl.cpp.AVX512_FP16.cpp -O3 -D__AVX512F__ -DCPU_CAPABILITY_AVX512 -DCPU_CAPABILITY_AVX512_VNNI -DCPU_CAPABILITY_AVX512_BF16 -mavx512f -mavx512bw -mavx512vl -mavx512dq -mavx512vbmi -mavx512vnni -mavx512bf16 -mfma -mamx-tile -mamx-int8 -mamx-bf16 -mavx512fp16 -DCPU_CAPABILITY_AMX -DCPU_CAPABILITY=AVX512_FP16 -DCPU_CAPABILITY_AVX512_FP16`
8080
---
8181

8282
>**Note:**

0 commit comments

Comments
 (0)