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

Update utils.py

Browse source
This commit is contained in:
Daniel Han 2025-02-02 15:17:50 -08:00
commit 64f1e817b3
1 changed files with 5 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/kernels/utils.py
View file

@ -227,7 +227,7 @@ if HAS_CUDA_STREAM:
if quant_state is None: return torch.matmul(X, W, out = out)
# For fast X @ W where seq_len == 1
# From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469
_, q_len, hd = X.shape
bsz, q_len, hd = X.shape
# assert(q_len == 1)
if type(quant_state) is not list:
@ -254,7 +254,7 @@ if HAS_CUDA_STREAM:
bout = shape[0]
if out is None:
out = torch.empty((1, 1, bout,), dtype = dtype, device = "cuda:0")
out = torch.empty((bsz, 1, bout,), dtype = dtype, device = "cuda:0")
# else:
# assert(out.shape == (1, 1, bout,))
# pass
@ -284,8 +284,9 @@ if HAS_CUDA_STREAM:
cgemm_4bit_inference_naive_bf16
blocksize = ctypes.c_int32(blocksize)
fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out),
lda, ldb, ldc, blocksize, CUDA_STREAM,)
for i in range(bsz):
fx(m, n, k, get_ptr(X[i]), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out[i]),
lda, ldb, ldc, blocksize, CUDA_STREAM,)
return out
pass