From e4fae9989c8ad60736bf2421c71b88bb94c3758d Mon Sep 17 00:00:00 2001 From: jiqing-feng Date: Mon, 14 Jul 2025 12:36:27 +0000 Subject: [PATCH] register triton kernel for quantization Signed-off-by: jiqing-feng --- bitsandbytes/backends/xpu/ops.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/bitsandbytes/backends/xpu/ops.py b/bitsandbytes/backends/xpu/ops.py index 1c1422c35..6e877cff8 100644 --- a/bitsandbytes/backends/xpu/ops.py +++ b/bitsandbytes/backends/xpu/ops.py @@ -139,6 +139,13 @@ def _gemv_4bit_impl( if not isinstance(lib, ErrorHandlerMockBNBNativeLibrary): logger.info("Register sycl bitsandbytes kernels for XPU") + # TODO: Remove the triton register when quantization sycl kernel is ready. + if triton_available: + from ..triton import ops as triton_ops + + register_kernel("bitsandbytes::quantize_blockwise", "xpu")(triton_ops.quantize_blockwise) + register_kernel("bitsandbytes::quantize_4bit", "xpu")(triton_ops.quantize_4bit) + @register_kernel("bitsandbytes::dequantize_4bit", "xpu") def _( A: torch.Tensor,