diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f4d2a63eb..c294893b8 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2,7 +2,7 @@ from __future__ import annotations from sglang.srt.dllm.config import DllmConfig from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.utils.common import ceil_align +from sglang.srt.utils.common import ceil_align, is_pin_memory_available # Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); @@ -1347,6 +1347,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): return self.dllm_config is not None def prepare_encoder_info_extend(self, input_ids: List[int], seq_lens: List[int]): + _pin = is_pin_memory_available(self.device) self.encoder_lens_cpu = [] self.encoder_cached = [] @@ -1363,9 +1364,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): or len(req.prefix_indices) >= im.num_image_tokens ) - self.encoder_lens = torch.tensor(self.encoder_lens_cpu, dtype=torch.int64).to( - self.device, non_blocking=True - ) + self.encoder_lens = torch.tensor( + self.encoder_lens_cpu, dtype=torch.int64, pin_memory=_pin + ).to(self.device, non_blocking=True) # Strip encoder infos pt = 0 @@ -1394,10 +1395,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): pt += req.extend_input_len # Reassign - self.input_ids = torch.tensor(sum(input_ids, []), dtype=torch.int64).to( - self.device, non_blocking=True - ) - self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64).to( + self.input_ids = torch.tensor( + sum(input_ids, []), dtype=torch.int64, pin_memory=_pin + ).to(self.device, non_blocking=True) + self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to( self.device, non_blocking=True ) self.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64) @@ -1449,7 +1450,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): r.token_type_ids for r in reqs if r.token_type_ids is not None ] - _pin = self.device != "cpu" and torch.cuda.is_available() + _pin = is_pin_memory_available(self.device) input_ids_tensor = torch.tensor( list(chain.from_iterable(input_ids)), dtype=torch.int64, pin_memory=_pin ).to(self.device, non_blocking=True) @@ -2056,9 +2057,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # No need to filter return - keep_indices_device = torch.tensor(keep_indices, dtype=torch.int64).to( - self.device, non_blocking=True - ) + keep_indices_device = torch.tensor( + keep_indices, + dtype=torch.int64, + pin_memory=is_pin_memory_available(self.device), + ).to(self.device, non_blocking=True) if self.model_config.is_encoder_decoder: self.encoder_lens = self.encoder_lens[keep_indices_device] diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index ed48bd4f2..cb3204e3d 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -609,8 +609,12 @@ def get_available_gpu_memory( return free_gpu_memory / (1 << 30) -def is_pin_memory_available() -> bool: - return torch.cuda.is_available() +def is_pin_memory_available(device=None) -> bool: + if not torch.cuda.is_available(): + return False + if device is not None and str(device) == "cpu": + return False + return True class LayerFn(Protocol):