From 8920cafbe7ce552437521084b4756a0810f075cb Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Thu, 1 Feb 2024 19:46:38 +1100
Subject: [PATCH] inference

---
 unsloth/kernels/utils.py | 6 +++---
 unsloth/models/llama.py  | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py
index 4f4ce41006..de97fea729 100644
--- a/unsloth/kernels/utils.py
+++ b/unsloth/kernels/utils.py
@@ -119,7 +119,7 @@ def fast_gemv(X, W, quant_state, out = None):
     # For fast X @ W where seq_len == 1
     # From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469
     bsz, q_len, hd = X.shape
-    assert(q_len == 1)
+    # assert(q_len == 1)
 
     if type(quant_state) is not list:
         # https://github.com/TimDettmers/bitsandbytes/pull/763/files
@@ -138,7 +138,7 @@ def fast_gemv(X, W, quant_state, out = None):
         offset, state2 = compressed_stats
         absmax2, code2, blocksize2, _, _, _, _ = state2
     pass
-    assert(dtype == X.dtype)
+    # assert(dtype == X.dtype)
     bout = shape[0]
 
     if out is None:
@@ -152,7 +152,7 @@ def fast_gemv(X, W, quant_state, out = None):
     k = shape[1]
     lda = shape[0]
     ldc = shape[0]
-    ldb = (X.shape[-1]+1)//2
+    ldb = (hd+1)//2
     m = ctypes.c_int32(m)
     n = ctypes.c_int32(n)
     k = ctypes.c_int32(k)
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index b0a10312d9..7095976b6d 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -286,7 +286,7 @@ def fast_mlp_inference(self, X):
     # up   = self.up_proj(X)
     bsz, _, hd = X.shape
     mlp_size = self.config.intermediate_size
-    temp = torch.empty((3, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda")
+    temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda")
 
     gate = fast_linear_forward(self.gate_proj, X, out = temp[0])
     up   = fast_linear_forward(self.  up_proj, X, out = temp[1])
@@ -294,7 +294,7 @@ def fast_mlp_inference(self, X):
     gate *= up
 
     # X = self.down_proj(gate)
-    down = fast_linear_forward(self.down_proj, gate, out = temp[2])
+    down = fast_linear_forward(self.down_proj, gate)
     return down
 pass