Skip to content

Commit 69b2146

Browse files
committed
update cuda op to support fp16/bf16 absmax
Signed-off-by: jiqing-feng <[email protected]>
1 parent 50ee994 commit 69b2146

File tree

1 file changed

+2
-2
lines changed
  • bitsandbytes/backends/cuda

1 file changed

+2
-2
lines changed

bitsandbytes/backends/cuda/ops.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor
237237
else:
238238
raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {A.dtype}")
239239

240-
return out, absmax
240+
return out, absmax.to(A.dtype)
241241

242242

243243
@register_kernel("bitsandbytes::dequantize_blockwise", "cuda")
@@ -332,7 +332,7 @@ def _(
332332
else:
333333
lib.cquantize_blockwise_fp32_nf4(*args)
334334

335-
return out, absmax
335+
return out, absmax.to(A.dtype)
336336

337337

338338
@register_kernel("bitsandbytes::dequantize_4bit", "cuda")

0 commit comments

Comments
 (0)