Skip to content

Commit 249a3cd

Browse files
authored
Enable XPU int matmul (#1547)
Signed-off-by: jiqing-feng <[email protected]>
1 parent 2354bdd commit 249a3cd

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

bitsandbytes/backends/cpu_xpu_common.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -194,8 +194,10 @@ def int8_linear_matmul_impl(
194194

195195
A_reshaped = A.reshape(m, k)
196196

197-
# torch._int_mm is available on CPU since torch 2.4
198-
if _torch_version_prereq(2, 4) and A.device.type == "cpu":
197+
# torch._int_mm is available on CPU since torch 2.4, XPU since torch 2.6
198+
if (A.device.type == "cpu" and _torch_version_prereq(2, 4)) or (
199+
A.device.type == "xpu" and _torch_version_prereq(2, 6)
200+
):
199201
C = torch._int_mm(A_reshaped, B.T).to(dtype)
200202
else:
201203
C = torch.matmul(A_reshaped.float(), B.t().float()).to(dtype)

0 commit comments

Comments
 (0)