From f2e87251c721482d05bf8dd21452e4eb5c20ba02 Mon Sep 17 00:00:00 2001 From: Fizza-Mukhtar Date: Tue, 30 Dec 2025 07:56:01 -0800 Subject: [PATCH] Fix 3D tensor support for bitsandbytes 8-bit matmul in forward pass --- unsloth/kernels/fast_lora.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 16b679d005..fbb18c3a15 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -383,12 +383,12 @@ class LoRA_QKV(torch.autograd.Function): # 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 = X.view(-1, X.shape[-1]) - - 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) + 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: