[MISC] fix pin_memory=torch.cuda.is_available(), use is_pin_memory_available (#37415)

Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
Kunshang Ji
2026-03-19 17:23:24 +08:00
committed by GitHub
parent da70c87e81
commit ca21483bf9
2 changed files with 4 additions and 2 deletions

View File

@@ -11,6 +11,7 @@ from transformers import PreTrainedTokenizerBase
from vllm.sampling_params import SamplingParams
from vllm.utils.import_utils import LazyLoader
from vllm.utils.platform_utils import is_pin_memory_available
from vllm.v1.structured_output.backend_types import (
StructuredOutputBackend,
StructuredOutputGrammar,
@@ -138,7 +139,7 @@ class LMFormatEnforcerBackend(StructuredOutputBackend):
(max_num_seqs, (self.vocab_size + 31) // 32),
-1,
dtype=torch.int32,
pin_memory=torch.cuda.is_available(),
pin_memory=is_pin_memory_available(),
)
def destroy(self):

View File

@@ -15,6 +15,7 @@ from regex import escape as regex_escape
from vllm.sampling_params import SamplingParams
from vllm.utils.import_utils import LazyLoader
from vllm.utils.platform_utils import is_pin_memory_available
from vllm.v1.structured_output.backend_types import (
StructuredOutputBackend,
StructuredOutputGrammar,
@@ -96,7 +97,7 @@ class OutlinesBackend(StructuredOutputBackend):
(max_num_seqs, (self.vocab_size + 31) // 32),
-1,
dtype=torch.int32,
pin_memory=torch.cuda.is_available(),
pin_memory=is_pin_memory_available(),
)
def destroy(self):