ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #3806 from Fizza-Mukhtar/fix/3d-tensor-matmul

Browse source 
Fix 3D tensor support for bitsandbytes 8-bit matmul in forward pass
This commit is contained in:
Daniel Han 2026-01-01 04:07:43 -08:00 • committed by GitHub
commit fbf0745eb0
1 changed files with 16 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

19
unsloth/kernels/fast_lora.py
View file

@ -379,9 +379,22 @@ class LoRA_QKV(torch.autograd.Function):
):
dtype = X.dtype
Q = matmul_lora(X, QW, QW_quant, QA, QB, QS)
K = matmul_lora(X, KW, KW_quant, KA, KB, KS)
V = matmul_lora(X, VW, VW_quant, VA, VB, VS)
# bitsandbytes 8-bit matmul expects 2D inputs.
# TorchInductor/AOTAutograd fails on 3D tensors during backward,
# so we explicitly flatten the sequence dimension.
orig_shape = X.shape
X_for_matmul = X
if X.dim() == 3:
X_for_matmul = X.view(-1, X.shape[-1])
Q = matmul_lora(X_for_matmul, QW, QW_quant, QA, QB, QS)
K = matmul_lora(X_for_matmul, KW, KW_quant, KA, KB, KS)
V = matmul_lora(X_for_matmul, VW, VW_quant, VA, VB, VS)
# Restore original shape after matmul
if len(orig_shape) == 3:
Q = Q.view(orig_shape[0], orig_shape[1], -1)
K = K.view(orig_shape[0], orig_shape[1], -1)
V = V.view(orig_shape[0], orig_shape[1], -1)
ctx.custom_saved_tensors = (
QW,