Skip to content

Add NPU support. #2200

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
brian-dellabetta merged 3 commits into from
Jan 13, 2026
+20 −5
Merged

Add NPU support. #2200

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
318c4dd
Add NPU support.
menogrey Jan 8, 2026
7bffc2c
Merge branch 'main' into add_npu_support
kylesayrs Jan 13, 2026
58f95a5
Merge branch 'main' into add_npu_support
brian-dellabetta Jan 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion src/llmcompressor/entrypoints/model_free/helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,13 @@ def gpu_if_available(device: torch.device | str | None) -> torch.device:
elif hasattr(torch, "xpu") and torch.xpu.is_available():
return torch.device("xpu:0")

elif hasattr(torch, "npu") and torch.npu.is_available():
return torch.device("npu:0")

else:
logger.warning("CUDA/XPU is not available! Compressing model on CPU instead")
logger.warning(
"CUDA/XPU/NPU is not available! Compressing model on CPU instead"
)
return torch.device("cpu")


Expand Down
6 changes: 5 additions & 1 deletion src/llmcompressor/pipelines/sequential/helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -543,8 +543,12 @@ def dispatch_for_sequential(model: PreTrainedModel) -> PreTrainedModel:
offloaded_dispatch(model, execution_device=torch.device("cuda:0"))
elif hasattr(torch, "xpu") and torch.xpu.is_available():
offloaded_dispatch(model, execution_device=torch.device("xpu:0"))
elif hasattr(torch, "npu") and torch.npu.is_available():
offloaded_dispatch(model, execution_device=torch.device("npu:0"))
else:
logger.warning("CUDA/XPU is not available! Compressing model on CPU instead")
logger.warning(
"CUDA/XPU/NPU is not available! Compressing model on CPU instead"
)

return model

Expand Down
12 changes: 9 additions & 3 deletions tools/collect_env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,20 +21,25 @@ def get_torch_hardware_info():

cuda_devices = []
amd_devices = []
npu_devices = []
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
name = torch.cuda.get_device_name(i)
if "AMD" in name.upper():
amd_devices.append(name)
else:
cuda_devices.append(name)
return cuda_devices, amd_devices
if hasattr(torch, "npu") and torch.npu.is_available():
for i in range(torch.npu.device_count()):
name = torch.npu.get_device_name(i)
npu_devices.append(name)
return cuda_devices, amd_devices, npu_devices
except ImportError:
return [], []
return [], [], []


def collect_environment_info():
cuda_devices, amd_devices = get_torch_hardware_info()
cuda_devices, amd_devices, npu_devices = get_torch_hardware_info()

info = {
"Operating System": platform.platform(),
Expand All @@ -45,6 +50,7 @@ def collect_environment_info():
"torch Version": get_version("torch"),
"CUDA Devices": cuda_devices if cuda_devices else "None",
"AMD Devices": amd_devices if amd_devices else "None",
"NPU Devices": npu_devices if npu_devices else "None",
}

print("### Environment Information ###")
Expand Down