[misc][cuda] use nvml to avoid accidentally cuda initialization (#6007)

2024-06-30 20:07:34 -07:00
parent af9ad46fca
commit 614aa51203
13 changed files with 86 additions and 68 deletions
--- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py
+++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py
@@ -12,8 +12,9 @@ from vllm.model_executor.layers.quantization.utils.marlin_perms import (
    marlin_perm, marlin_scale_perm, marlin_scale_perm_single)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
    get_pack_factor, quantize_weights, sort_weights)
+from vllm.utils import get_device_capability_stateless

-__cuda_arch = torch.cuda.get_device_capability()
+__cuda_arch = get_device_capability_stateless()

 MARLIN_TILE = 16