minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Use pinned memory for asynchronous sampling metadata transfers (#39777)

Browse Source 
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
Yuxuan Zhang
2026-09-20 15:31:31 -07:00
committed by GitHub
co-authored by Xinyuan Tong hnyls2002
parent 95521da18d
commit f31a7bd45c
7 changed files with 245 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/logprob_processor.py
+8 -7
View File
@@ -18,7 +18,7 @@ from sglang.srt.model_executor.runner_utils.pool import (
graph_pool_borrow_largest_run,
)
from sglang.srt.runtime_context import get_exec
from sglang.srt.utils.common import async_d2h
from sglang.srt.utils.common import async_d2h, is_pin_memory_available
if TYPE_CHECKING:
from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput
@@ -151,15 +151,16 @@ def get_token_ids_logprobs_raw(
no_copy_to_cpu: bool = False,
):
vals, idxs = [], []
pin_memory = is_pin_memory_available(logprobs.device)
if stage == LogprobStage.DECODE:
for i, token_ids in enumerate(token_ids_logprobs_list):
if token_ids is None:
vals.append([])
idxs.append([])
else:
token_ids_tensor = torch.tensor(token_ids, dtype=torch.long).to(
logprobs.device, non_blocking=True
)
token_ids_tensor = torch.tensor(
token_ids, dtype=torch.long, pin_memory=pin_memory
).to(logprobs.device, non_blocking=True)
row = logprobs[i, token_ids_tensor]
vals.append(row if no_copy_to_cpu else row.tolist())
idxs.append(token_ids)
@@ -178,9 +179,9 @@ def get_token_ids_logprobs_raw(
idxs.append([])
pt += pruned_len
continue
token_ids_tensor = torch.tensor(token_ids, dtype=torch.long).to(
logprobs.device, non_blocking=True
)
token_ids_tensor = torch.tensor(
token_ids, dtype=torch.long, pin_memory=pin_memory
).to(logprobs.device, non_blocking=True)
pos_logprobs = logprobs[pt : pt + pruned_len, token_ids_tensor]
vals.append(pos_logprobs if no_copy_to_cpu else pos_logprobs.tolist())
idxs.append([token_ids for _ in range(pruned_len)])
python/sglang/srt/sampling/penaltylib/frequency_penalty.py
+6 -2
View File
@@ -1,6 +1,7 @@
import torch
from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer
from sglang.srt.utils.common import is_pin_memory_available
class BatchedFrequencyPenalizer(_BatchedPenalizer):
@@ -15,6 +16,7 @@ class BatchedFrequencyPenalizer(_BatchedPenalizer):
)
def _prepare(self):
pin_memory = is_pin_memory_available(self.orchestrator.device)
self.cumulated_frequency_penalties = torch.zeros(
(len(self.orchestrator.reqs()), self.orchestrator.vocab_size),
dtype=torch.float32,
@@ -28,9 +30,11 @@ class BatchedFrequencyPenalizer(_BatchedPenalizer):
for req in self.orchestrator.reqs()
],
dtype=torch.float32,
device=self.orchestrator.device,
pin_memory=pin_memory,
)
).unsqueeze_(1)
.to(self.orchestrator.device, non_blocking=True)
.unsqueeze_(1)
)
def _cumulate_output_tokens(self, output_ids: torch.Tensor):
self.cumulated_frequency_penalties.scatter_add_(
python/sglang/srt/sampling/penaltylib/min_new_tokens.py
+20 -9
View File
@@ -1,6 +1,7 @@
import torch
from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer
from sglang.srt.utils.common import is_pin_memory_available
class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
@@ -14,15 +15,21 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
)
def _prepare(self):
self.min_new_tokens = torch.tensor(
data=[
req.sampling_params.min_new_tokens for req in self.orchestrator.reqs()
],
dtype=torch.int32,
device=self.orchestrator.device,
).unsqueeze_(1)
pin_memory = is_pin_memory_available(self.orchestrator.device)
self.min_new_tokens = (
torch.tensor(
data=[
req.sampling_params.min_new_tokens
for req in self.orchestrator.reqs()
],
dtype=torch.int32,
pin_memory=pin_memory,
)
.to(self.orchestrator.device, non_blocking=True)
.unsqueeze_(1)
)
padded_stop_token_ids = torch.nn.utils.rnn.pad_sequence(
padded_stop_token_ids_cpu = torch.nn.utils.rnn.pad_sequence(
sequences=[
torch.tensor(
data=[
@@ -40,13 +47,17 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
if token_id is not None
],
dtype=torch.int64,
device=self.orchestrator.device,
)
for req in self.orchestrator.reqs()
],
batch_first=True,
padding_value=self.orchestrator.vocab_size,
)
if pin_memory:
padded_stop_token_ids_cpu = padded_stop_token_ids_cpu.pin_memory()
padded_stop_token_ids = padded_stop_token_ids_cpu.to(
self.orchestrator.device, non_blocking=True
)
self.stop_token_penalties = torch.zeros(
size=(len(self.orchestrator.reqs()), self.orchestrator.vocab_size + 1),
dtype=torch.float32,
python/sglang/srt/sampling/penaltylib/presence_penalty.py
+6 -2
View File
@@ -1,6 +1,7 @@
import torch
from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer
from sglang.srt.utils.common import is_pin_memory_available
class BatchedPresencePenalizer(_BatchedPenalizer):
@@ -15,6 +16,7 @@ class BatchedPresencePenalizer(_BatchedPenalizer):
)
def _prepare(self):
pin_memory = is_pin_memory_available(self.orchestrator.device)
self.cumulated_presence_penalties = torch.zeros(
(len(self.orchestrator.reqs()), self.orchestrator.vocab_size),
dtype=torch.float32,
@@ -28,9 +30,11 @@ class BatchedPresencePenalizer(_BatchedPenalizer):
for req in self.orchestrator.reqs()
],
dtype=torch.float32,
device=self.orchestrator.device,
pin_memory=pin_memory,
)
).unsqueeze_(1)
.to(self.orchestrator.device, non_blocking=True)
.unsqueeze_(1)
)
def _cumulate_output_tokens(self, output_ids: torch.Tensor):
self.cumulated_presence_penalties.scatter_(
python/sglang/srt/sampling/penaltylib/repetition_penalty.py
+6 -2
View File
@@ -2,6 +2,7 @@ import torch
from sglang.srt.sampling.penaltylib.orchestrator import _BatchedPenalizer
from sglang.srt.utils import get_compiler_backend, is_npu
from sglang.srt.utils.common import is_pin_memory_available
_is_npu = is_npu()
@@ -29,6 +30,7 @@ class BatchedRepetitionPenalizer(_BatchedPenalizer):
)
def _prepare(self):
pin_memory = is_pin_memory_available(self.orchestrator.device)
self.cumulated_repetition_penalties = torch.ones(
(len(self.orchestrator.reqs()), self.orchestrator.vocab_size),
dtype=torch.float32,
@@ -41,9 +43,11 @@ class BatchedRepetitionPenalizer(_BatchedPenalizer):
for req in self.orchestrator.reqs()
],
dtype=torch.float32,
device=self.orchestrator.device,
pin_memory=pin_memory,
)
).unsqueeze_(1)
.to(self.orchestrator.device, non_blocking=True)
.unsqueeze_(1)
)
def _cumulate_output_tokens(self, output_ids: torch.Tensor):
self.cumulated_repetition_penalties.scatter_(
python/sglang/srt/sampling/sampling_batch_info.py
+18 -2
View File
@@ -158,10 +158,26 @@ class SamplingBatchInfo:
logit_bias = None
if any(r.sampling_params.logit_bias is not None for r in reqs):
logit_bias = torch.zeros(len(reqs), vocab_size, device=device)
rows, cols, vals = [], [], []
for i, r in enumerate(reqs):
if r.sampling_params.logit_bias is not None:
for key, value in r.sampling_params.logit_bias.items():
logit_bias[i, int(key)] = value
# Dedup on int(key) first ("1" and "01" collide): duplicate
# indices make the index_put below nondeterministic on CUDA.
row_bias = {
int(key): value
for key, value in r.sampling_params.logit_bias.items()
}
for token_id, value in row_bias.items():
rows.append(i)
cols.append(token_id)
vals.append(value)
if rows:
logit_bias[
_rows_to_device_indices(rows, device, _pin),
_rows_to_device_indices(cols, device, _pin),
] = torch.tensor(vals, dtype=logit_bias.dtype, pin_memory=_pin).to(
device, non_blocking=True
)
# Check if any request has custom logit processor
has_custom_logit_processor = (