A high-throughput and memory-efficient inference and serving engine for LLMs
[CPU] Add missing scalar fallback for CPU W4A8 INT4 GEMM (#44523)
Signed-off-by: wcy <233313160abc@gmail.com> Co-authored-by: lyd1992 <liuyudong@iscas.ac.cn>
wcy committed
Jun 11, 2026 at 16:52 UTC
f06aefb4e3757f0fc76bc117a7aa5c41632ce72b
4 files changed
+61
-15
cmake/cpu_extension.cmake
+6
@@ -438,6 +438,12 @@ if(USE_ONEDNN)
438
${VLLM_EXT_SRC})
439
endif()
440
441
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
442
+ set(VLLM_EXT_SRC
443
+ "csrc/cpu/sgl-kernels/gemm_int4.cpp"
444
+ ${VLLM_EXT_SRC})
445
+endif()
446
+
447
if (ENABLE_X86_ISA)
448
set(VLLM_EXT_SRC_SGL
449
"csrc/cpu/sgl-kernels/conv.cpp"
csrc/cpu/sgl-kernels/gemm_int4.cpp
+37
-1
@@ -268,6 +268,23 @@ void _dequant_gemm_accum_small_M(
268
_dequant_gemm_accum_small_M<M, N, ldb, sym_quant_act>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
269
#endif
270
271
+template <int64_t N, int64_t ldb>
272
+inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t n) {
273
+ // B is packed as [_block_k / 4, N / 2, 4] for VNNI4. Each byte stores two
274
+ // columns from adjacent 8-column groups for one K lane.
275
+ constexpr int64_t n_group_size = 8;
276
+ constexpr int64_t vnni_size = 4;
277
+ static_assert(N % (2 * n_group_size) == 0);
278
+
279
+ int64_t n_group = n / n_group_size;
280
+ int64_t ni = n % n_group_size;
281
+ int64_t ki = k % vnni_size;
282
+ int64_t k_base = k - ki;
283
+ int64_t packed_n = (n_group / 2) * n_group_size + ni;
284
+ uint8_t packed = B[k_base * ldb + packed_n * vnni_size + ki];
285
+ return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f);
286
+}
287
+
288
template <int64_t N, int64_t ldb, bool sym_quant_act>
289
void _dequant_gemm_accum(
290
float* C,
@@ -321,7 +338,24 @@ void _dequant_gemm_accum(
338
} else
339
#endif
340
{
324
- TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!");
341
+ for (int64_t m = 0; m < M; ++m) {
342
+ for (int64_t n = 0; n < N; ++n) {
343
+ int32_t acc = 0;
344
+ for (int64_t k = 0; k < K; ++k) {
345
+ int32_t b = load_uint4_vnni<N, ldb>(B, k, n) - qzeros_b[n];
346
+ if constexpr (sym_quant_act) {
347
+ const int8_t* A_s8 = reinterpret_cast<const int8_t*>(A);
348
+ acc += static_cast<int32_t>(A_s8[m * lda + k]) * b;
349
+ } else {
350
+ acc += static_cast<int32_t>(A[m * lda + k]) * b;
351
+ }
352
+ }
353
+ if constexpr (!sym_quant_act) {
354
+ acc -= qzeros_a[m] * compensation[n];
355
+ }
356
+ C[m * ldc + n] += static_cast<float>(acc) * scales_a[m] * scales_b[n];
357
+ }
358
+ }
359
}
360
}
361
@@ -496,9 +530,11 @@ void _da8w4_linear_impl(
530
store_out<out_dtype, BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/);
531
}
532
}
533
+#if defined(CPU_CAPABILITY_AVX512)
534
if (use_brgemm) {
535
at::native::cpublas::brgemm_release();
536
}
537
+#endif
538
});
539
}
540
csrc/cpu/sgl-kernels/vec.h
+1
-1
@@ -245,7 +245,7 @@ quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restric
245
246
for (int64_t k = 0; k < K; ++k) {
247
const float val = static_cast<float>(A[k]) * inv_scale;
248
- Aq[k] = (uint8_t)(std::round(val)) + 128;
248
+ Aq[k] = static_cast<uint8_t>(static_cast<int32_t>(std::round(val)) + 128);
249
}
250
As = scale;
251
}
csrc/cpu/torch_bindings.cpp
+17
-13
@@ -429,19 +429,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
429
ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
430
&int8_scaled_mm_with_quant);
431
432
- // Adapted from sglang: INT4 W4A8 kernels
433
- ops.def(
434
- "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
435
- "scales, int quant_method_4bit) -> (Tensor, "
436
- "Tensor, Tensor)");
437
- ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
438
- &convert_weight_packed_scale_zp);
439
-
440
- ops.def(
441
- "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
442
- "Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
443
- ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
444
-
432
// Adapted from sglang: FP8 W8A16 kernel
433
ops.def(
434
"fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) "
@@ -468,6 +455,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
455
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
456
#endif
457
458
+#if (defined(__AVX512BF16__) && defined(__AVX512F__) && \
459
+ defined(__AVX512VNNI__)) || \
460
+ defined(__riscv)
461
+ // Adapted from sglang: INT4 W4A8 kernels
462
+ ops.def(
463
+ "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
464
+ "scales, int quant_method_4bit) -> (Tensor, "
465
+ "Tensor, Tensor)");
466
+ ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
467
+ &convert_weight_packed_scale_zp);
468
+
469
+ ops.def(
470
+ "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
471
+ "Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
472
+ ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
473
+#endif
474
+
475
// Adapted from sglang: GDN kernels
476
ops.def(
477
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "