A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 02:15 UTC
Enabled BnB NF4 inference on Gaudi (#20172)
Signed-off-by: Ruheena Suhani Shaik <rsshaik@habana.ai>
Ruheena Suhani Shaik committed
Jul 15, 2025 at 08:56 UTC
016b8d1b7f1d51a46306ebf0d62f0499bce2a36d
2 files changed
+18
-8
vllm/model_executor/layers/quantization/bitsandbytes.py
+6
-6
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
13
from vllm.model_executor.layers.quantization import QuantizationMethods
14
from vllm.model_executor.layers.quantization.base_config import (
15
QuantizationConfig)
16
+from vllm.platforms import current_platform
17
from vllm.utils import direct_register_custom_op
18
19
@@ -390,12 +391,11 @@ def _apply_bnb_4bit_fake(
391
392
393
try:
393
- direct_register_custom_op(
394
- op_name="apply_bnb_4bit",
395
- op_func=_apply_bnb_4bit,
396
- mutates_args=["out"],
397
- fake_impl=_apply_bnb_4bit_fake,
398
- )
394
+ direct_register_custom_op(op_name="apply_bnb_4bit",
395
+ op_func=_apply_bnb_4bit,
396
+ mutates_args=["out"],
397
+ fake_impl=_apply_bnb_4bit_fake,
398
+ dispatch_key=current_platform.dispatch_key)
399
apply_bnb_4bit = torch.ops.vllm.apply_bnb_4bit
400
401
except AttributeError as error:
vllm/model_executor/model_loader/bitsandbytes_loader.py
+12
-2
@@ -199,6 +199,10 @@ class BitsAndBytesModelLoader(BaseModelLoader):
199
200
if self.pre_quant:
201
if self.load_8bit:
202
+ if current_platform.is_hpu():
203
+ raise ValueError(
204
+ "currently hpu supports 4bit quantization only")
205
+
206
return self._quantized_8bit_generator(
207
hf_weights_files, use_safetensors,
208
quant_state_dict), quant_state_dict
@@ -302,6 +306,10 @@ class BitsAndBytesModelLoader(BaseModelLoader):
306
in temp_state_dict):
307
quant_state = _parse_quant_state(mapped_weight_name,
308
temp_state_dict)
309
+ if current_platform.is_hpu():
310
+ assert quant_state.quant_type == "nf4", (
311
+ "currently hpu supports nf4 quant_type only")
312
+
313
quant_state_dict[mapped_weight_name] = quant_state
314
yield org_weight_name, weight_tensor
315
else:
@@ -372,10 +380,12 @@ class BitsAndBytesModelLoader(BaseModelLoader):
380
...]
381
382
# bitsandbytes requires data in GPU
375
- if weight_sub_tensor.is_cuda:
383
+ if (weight_sub_tensor.is_cuda
384
+ or weight_sub_tensor.device.type == "hpu"):
385
loaded_weight = weight_sub_tensor
386
else:
378
- loaded_weight = weight_sub_tensor.cuda()
387
+ loaded_weight = weight_sub_tensor.to(
388
+ device=current_platform.device_type)
389
390
# remove the following after the issue is fixed:
391
# https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1342