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

[Logprob] Serve input-logprob temporaries from CUDA-graph-pool dead space (#40038)

Browse Source 
Co-authored-by: cctry <csycfl@gmail.com>
This commit is contained in:
metamergebot
2026-09-19 12:03:52 -07:00
committed by GitHub
co-authored by cctry
parent c5326d28a3
commit 9e5a62a767
10 changed files with 674 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/logits_processor.py
+23
View File
@@ -243,6 +243,9 @@ class LogitsProcessorOutput:
)
input_token_ids_logprobs_idx: Optional[List] = None
# Completion of input-logprob copies from borrowed graph storage.
input_logprobs_copy_done: Optional[torch.cuda.Event] = None
## Part 4: Diffusion LLM only.
full_logits: Optional[torch.Tensor] = None
@@ -262,6 +265,22 @@ class LogitsProcessorOutput:
# Scheduler-local output copied alongside the ordinary generation result.
auxiliary_device_output: Optional[DeviceAuxiliaryOutput] = None
def finalize_input_logprobs(self) -> None:
if self.input_logprobs_copy_done is None:
return
self.input_logprobs_copy_done.synchronize()
self.input_logprobs_copy_done = None
# Only borrowed results contain spans within each sequence. Other
# producers (including multi-item scoring) keep their existing layout.
for sequences in (
self.input_top_logprobs_val,
self.input_top_logprobs_idx,
self.input_token_ids_logprobs_val,
):
if sequences is not None:
for i, spans in enumerate(sequences):
sequences[i] = [row for span in spans for row in span.tolist()]
@dataclasses.dataclass
class LogitsMetadata:
@@ -278,6 +297,8 @@ class LogitsMetadata:
extend_logprob_pruned_lens_cpu: Optional[List[int]] = None
top_logprobs_nums: Optional[List[int]] = None
extend_input_logprob_token_ids_gpu: Optional[torch.Tensor] = None
sample_indices_cpu: Optional[List[int]] = None
input_logprob_indices_cpu: Optional[List[int]] = None
token_ids_logprobs: Optional[List[List[int]]] = None
# logits and logprobs post processing
@@ -737,6 +758,8 @@ class LogitsProcessor(nn.Module):
else [torch.cat(lst) for lst in aux_pruned_states_lists]
)
logits_metadata.sample_indices_cpu = sample_indices
logits_metadata.input_logprob_indices_cpu = input_logprob_indices
# Build the index tensors via pinned host memory + non-blocking H2D
# so the small copy doesn't drain the stream.
sample_indices = torch.tensor(
python/sglang/srt/layers/logprob_processor.py
+207 -137
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import bisect
import dataclasses
import logging
from contextlib import nullcontext
@@ -17,6 +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
if TYPE_CHECKING:
from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput
@@ -45,6 +47,7 @@ class LogprobResult:
top_logprobs_idx: Optional[List] = None
token_ids_logprobs_val: Optional[List] = None
token_ids_logprobs_idx: Optional[List] = None
input_copy_done: Optional[torch.cuda.Event] = None
def write_input_to(self, logits_output: LogitsProcessorOutput) -> None:
if self.token_logprobs is not None:
@@ -55,6 +58,7 @@ class LogprobResult:
if self.token_ids_logprobs_val is not None:
logits_output.input_token_ids_logprobs_val = self.token_ids_logprobs_val
logits_output.input_token_ids_logprobs_idx = self.token_ids_logprobs_idx
logits_output.input_logprobs_copy_done = self.input_copy_done
def write_output_to(self, logits_output: LogitsProcessorOutput) -> None:
if self.token_logprobs is not None:
@@ -202,6 +206,7 @@ def get_top_logprobs_chunk(
split_pruned_len: int,
log_normalizer: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
precomputed_topk: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
copy_to_pinned_cpu: bool = False,
) -> int:
"""Get top-k logprobs for each sequence in the chunk.
@@ -217,6 +222,8 @@ def get_top_logprobs_chunk(
log_normalizer: Per-row (max, logsumexp - max) of the logits
precomputed_topk: (raw fp32 top values, indices) from the fused
logsumexp+top-k kernel, sorted with lowest-index tie-breaking
copy_to_pinned_cpu: Copy gathered results to pinned CPU spans;
finalize_input_logprobs converts them to rows after the copy completes.
Returns:
int: Number of remaining tokens to process in next chunk
@@ -234,8 +241,14 @@ def get_top_logprobs_chunk(
]
else:
values_tensor, indices_tensor = logprobs.topk(max_k, dim=1)
values = values_tensor.tolist()
indices = indices_tensor.tolist()
if copy_to_pinned_cpu:
values = async_d2h(values_tensor)
indices = async_d2h(indices_tensor)
rows_avail = values.shape[0]
else:
values = values_tensor.tolist()
indices = indices_tensor.tolist()
rows_avail = len(values)
pt = 0
next_split_pruned_len = 0
@@ -254,17 +267,19 @@ def get_top_logprobs_chunk(
top_logprobs_idx.append([])
continue
# Handle remaining tokens in next chunk if any
available_len = min(pruned_len, max(rows_avail - pt, 0))
if available_len < pruned_len:
next_split_pruned_len = split_pruned_len + available_len
# Get the top-k logprobs
val = []
idx = []
for j in range(pruned_len):
# Handle remaining tokens in next chunk if any
if pt + j >= len(values):
next_split_pruned_len = split_pruned_len + j
break
# Append the top-k logprobs
val.append(values[pt + j][:k])
idx.append(indices[pt + j][:k])
if copy_to_pinned_cpu:
# Keep one span per sequence to avoid per-row tensor slicing.
val = [values[pt : pt + available_len, :k]] if available_len else []
idx = [indices[pt : pt + available_len, :k]] if available_len else []
else:
val = [values[pt + j][:k] for j in range(available_len)]
idx = [indices[pt + j][:k] for j in range(available_len)]
# Append or extend based on whether the sequence was split across chunks
# Split-sequence continuations extend; everyone else owns a fresh
@@ -288,6 +303,7 @@ def get_token_ids_logprobs_chunk(
token_ids_logprobs_idx: List,
split_pruned_len: int = 0,
log_normalizer: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
copy_to_pinned_cpu: bool = False,
):
"""Get token_ids logprobs for each sequence in the chunk.
@@ -301,6 +317,8 @@ def get_token_ids_logprobs_chunk(
token_ids_logprobs_idx: List to store token indices
split_pruned_len: Length of pruned tokens from previous chunk
log_normalizer: Per-row (max, logsumexp - max) of the logits
copy_to_pinned_cpu: Copy gathered results to pinned CPU spans;
finalize_input_logprobs converts them to rows after the copy completes.
Returns:
int: Number of remaining tokens to process in next chunk
@@ -329,20 +347,25 @@ def get_token_ids_logprobs_chunk(
token_ids_logprobs_idx.append([])
continue
# Handle remaining tokens in next chunk if any
available_len = min(pruned_len, max(logprobs.shape[0] - pt, 0))
if available_len < pruned_len:
next_split_pruned_len = split_pruned_len + available_len
# Get the token ids logprobs
val = []
idx = []
for j in range(pruned_len):
# Handle remaining tokens in next chunk if any
if pt + j >= logprobs.shape[0]:
next_split_pruned_len = split_pruned_len + j
break
if token_ids is not None:
row = logprobs[pt + j, token_ids]
if log_normalizer is not None:
row = (row.float() - row_max[pt + j]) - row_log_sum[pt + j]
val.append(row.tolist())
idx.append(token_ids)
if token_ids is not None and available_len > 0:
rows = logprobs[pt : pt + available_len, token_ids]
if log_normalizer is not None:
rows = (
rows.float() - row_max[pt : pt + available_len, None]
) - row_log_sum[pt : pt + available_len, None]
if copy_to_pinned_cpu:
val.append(async_d2h(rows))
else:
val.extend(rows.tolist())
idx.extend([token_ids] * available_len)
# Split-sequence continuations extend; everyone else owns a fresh
# (possibly empty) entry.
@@ -497,9 +520,12 @@ class InputLogprobProcessor:
else:
chunk_size = self.logprobs_chunk_size
borrow_logits_memory = False
if pruned_states.is_cuda and not skip_chunking_for_dp_attn:
borrow_logits_memory = self._can_borrow_logits_memory(chunk_size)
borrow_logprob_memory = False
if pruned_states.is_cuda:
borrow_logprob_memory = self._can_borrow_logprob_memory(
chunk_size=chunk_size,
borrow_logits=not skip_chunking_for_dp_attn,
)
return self._forward_by_chunk(
pruned_states,
@@ -510,23 +536,36 @@ class InputLogprobProcessor:
get_logits_fn,
logits_metadata,
chunk_size,
borrow_logits_memory=borrow_logits_memory,
borrow_logprob_memory=borrow_logprob_memory,
borrow_logits_memory=(
borrow_logprob_memory and not skip_chunking_for_dp_attn
),
)
def _can_borrow_logits_memory(self, chunk_size: int) -> bool:
def _can_borrow_logprob_memory(
self,
chunk_size: int,
borrow_logits: bool,
) -> bool:
"""Borrow only when the planned chunk fits on every TP rank.
Resizing chunks to graph-pool capacity changes LM-head GEMM shapes
and their rounding. Keep chunk boundaries independent of the graph
memory layout, including when borrowing is disabled on another runner.
"""
# TP gathering can hold the local projection, gathered tensor, and
# contiguous reshape together; FP32 bounds their possible dtypes.
bytes_per_row = 3 * self.vocab_size * 4
# NCCL's symmetric allocator owns its collective buffers, so those
# allocations cannot be counted as borrowed storage.
# A TP gather can hold the local projection, gathered tensor, and its
# contiguous reshape together. Scoring alone needs at most two matrices.
matrices_per_chunk = (
3 if borrow_logits else (1 if self.enable_fast_input_logprobs else 2)
)
bytes_per_row = matrices_per_chunk * self.vocab_size * 4
# NCCL's symmetric allocator owns its collective buffers, so they
# cannot be counted as borrowed storage.
free_run = (
0 if is_symmetric_memory_enabled() else graph_pool_borrow_largest_run()
0
if borrow_logits and is_symmetric_memory_enabled()
else graph_pool_borrow_largest_run()
)
fit_rows = max(0, free_run - _GRAPH_POOL_BORROW_SLACK_BYTES) // bytes_per_row
if self.chunking_group is not None:
@@ -551,13 +590,16 @@ class InputLogprobProcessor:
get_logits_fn: Callable,
logits_metadata: LogitsMetadata,
chunk_size: int,
borrow_logprob_memory: bool = False,
borrow_logits_memory: bool = False,
) -> Tuple[LogprobResult, torch.Tensor]:
"""Compute input logprobs chunk by chunk to cap peak memory."""
total_size = pruned_states.shape[0]
num_chunks = (total_size + chunk_size - 1) // chunk_size
token_logprobs = []
sample_indices_cpu = logits_metadata.sample_indices_cpu
input_logprob_indices_cpu = logits_metadata.input_logprob_indices_cpu
token_logprobs = None if borrow_logprob_memory else []
if logits_metadata.extend_return_top_logprob:
top_logprobs_val = []
top_logprobs_idx = []
@@ -595,15 +637,10 @@ class InputLogprobProcessor:
if num_chunks > 1 and hasattr(lm_head, "set_lm_head_pass"):
lm_head.set_lm_head_pass(i)
# Get indices for this chunk
chunk_mask = (input_logprob_indices >= start_idx) & (
input_logprob_indices < end_idx
)
global_indices = input_logprob_indices[chunk_mask]
chunk_indices = global_indices - start_idx
# Get the positions in the original array where chunk_mask is True
# This is needed to correctly index into extend_input_logprob_token_ids_gpu
mask_indices = torch.nonzero(chunk_mask, as_tuple=True)[0]
# The sorted host indices avoid per-chunk device synchronization.
lp_lo = bisect.bisect_left(input_logprob_indices_cpu, start_idx)
lp_hi = bisect.bisect_left(input_logprob_indices_cpu, end_idx)
chunk_indices = input_logprob_indices[lp_lo:lp_hi] - start_idx
# Get the logits for this chunk. Each chunk must own its output:
# writing through the shared graph logits buffer would alias
@@ -621,107 +658,133 @@ class InputLogprobProcessor:
use_logits_buffer=num_chunks == 1,
)
# Sampled outputs must survive graph replay, so they are allocated
# outside borrowing. The transient logits are released below.
# These outputs must survive graph replay. Allocate them outside
# borrowing, then consume and release chunk_logits in the next scope.
if i == 0:
sampled_logits = torch.empty(
(sample_indices.shape[0], chunk_logits.shape[1]),
dtype=chunk_logits.dtype,
device=chunk_logits.device,
)
# Handle sampled logits for the chunk if needed
# This must be done before the continue statement to ensure all sampled_logits are filled
chunk_sample_mask = (sample_indices >= start_idx) & (
sample_indices < end_idx
)
if chunk_sample_mask.any():
chunk_sample_indices = sample_indices[chunk_sample_mask] - start_idx
sampled_logits[chunk_sample_mask] = chunk_logits[chunk_sample_indices]
# Zero-logprob-row chunks still need the per-sequence bookkeeping below.
chunk_logprobs = chunk_logits[chunk_indices]
del chunk_logits
# End at the last row inside the chunk; token_to_seq_idx[end_idx]
# belongs to the next chunk and would emit its sequence twice.
chunk_slice = slice(
token_to_seq_idx[start_idx], token_to_seq_idx[end_idx - 1] + 1
)
chunk_precomputed_topk = None
if self.enable_fast_input_logprobs:
# Every consumer below needs only small gathers / top-k plus a
# per-row normalizer, so keep the raw logits and skip the
# full-vocab log-softmax materialization entirely. When top-k
# is requested, the fused kernel produces the normalizer and
# the top-k in the same single read of the logits.
max_k = (
max(logits_metadata.top_logprobs_nums[chunk_slice])
if logits_metadata.extend_return_top_logprob
else 0
)
if 0 < max_k <= fused_max_k:
row_max, row_log_sum, top_vals, top_idx = fused_kernel(
chunk_logprobs, max_k
if borrow_logprob_memory:
token_logprobs = torch.empty(
input_logprob_indices.shape,
dtype=(
torch.float32
if self.enable_fast_input_logprobs
else chunk_logits.dtype
),
device="cpu",
pin_memory=True,
)
chunk_log_normalizer = (row_max, row_log_sum)
chunk_precomputed_topk = (top_vals, top_idx)
# Fill the sampled logits whose rows fall in this chunk.
s_lo = bisect.bisect_left(sample_indices_cpu, start_idx)
s_hi = bisect.bisect_left(sample_indices_cpu, end_idx)
if s_hi > s_lo:
sampled_logits[s_lo:s_hi] = chunk_logits[
sample_indices[s_lo:s_hi] - start_idx
]
borrow_scope = (
borrow_graph_pool(user="input logprob processing")
if borrow_logprob_memory
else nullcontext()
)
with borrow_scope:
chunk_logprobs = chunk_logits[chunk_indices]
del chunk_logits
# End at the last row inside the chunk; token_to_seq_idx[end_idx]
# belongs to the next chunk and would emit its sequence twice.
chunk_slice = slice(
token_to_seq_idx[start_idx], token_to_seq_idx[end_idx - 1] + 1
)
chunk_precomputed_topk = None
if self.enable_fast_input_logprobs:
# Every consumer below needs only small gathers / top-k plus a
# per-row normalizer, so keep the raw logits and skip the
# full-vocab log-softmax materialization entirely. When top-k
# is requested, the fused kernel produces the normalizer and
# the top-k in the same single read of the logits.
max_k = (
max(logits_metadata.top_logprobs_nums[chunk_slice])
if logits_metadata.extend_return_top_logprob
else 0
)
if 0 < max_k <= fused_max_k:
row_max, row_log_sum, top_vals, top_idx = fused_kernel(
chunk_logprobs, max_k
)
chunk_log_normalizer = (row_max, row_log_sum)
chunk_precomputed_topk = (top_vals, top_idx)
else:
chunk_log_normalizer = compute_row_log_normalizer(
chunk_logprobs
)
else:
chunk_log_normalizer = compute_row_log_normalizer(chunk_logprobs)
else:
# Free the raw logits before the out-of-place log_softmax:
# keeping all three alive is a 3x peak, which OOMs when the
# single chunk covers a large batch.
chunk_log_normalizer = None
chunk_logprobs = torch.nn.functional.log_softmax(chunk_logprobs, dim=-1)
# Free the raw logits before the out-of-place log_softmax:
# keeping all three alive is a 3x peak, which OOMs when the
# single chunk covers a large batch.
chunk_log_normalizer = None
chunk_logprobs = torch.nn.functional.log_softmax(
chunk_logprobs, dim=-1
)
# Get the logprob of top-k tokens
if logits_metadata.extend_return_top_logprob:
top_k_nums = logits_metadata.top_logprobs_nums[chunk_slice]
pruned_lens = logits_metadata.extend_logprob_pruned_lens_cpu[
chunk_slice
# Get the logprob of top-k tokens
if logits_metadata.extend_return_top_logprob:
top_k_nums = logits_metadata.top_logprobs_nums[chunk_slice]
pruned_lens = logits_metadata.extend_logprob_pruned_lens_cpu[
chunk_slice
]
split_len_topk = get_top_logprobs_chunk(
chunk_logprobs,
top_k_nums,
pruned_lens,
top_logprobs_val,
top_logprobs_idx,
split_len_topk,
log_normalizer=chunk_log_normalizer,
precomputed_topk=chunk_precomputed_topk,
copy_to_pinned_cpu=borrow_logprob_memory,
)
# Get the logprob of given token id
if logits_metadata.extend_token_ids_logprob:
token_ids_logprobs = logits_metadata.token_ids_logprobs[chunk_slice]
pruned_lens = logits_metadata.extend_logprob_pruned_lens_cpu[
chunk_slice
]
split_len_token_ids = get_token_ids_logprobs_chunk(
chunk_logprobs,
token_ids_logprobs,
pruned_lens,
token_ids_logprobs_val,
token_ids_logprobs_idx,
split_len_token_ids,
log_normalizer=chunk_log_normalizer,
copy_to_pinned_cpu=borrow_logprob_memory,
)
# Get the logprob of the requested token ids
chunk_token_logprobs = chunk_logprobs[
torch.arange(chunk_logprobs.shape[0], device=chunk_logprobs.device),
logits_metadata.extend_input_logprob_token_ids_gpu[lp_lo:lp_hi],
]
split_len_topk = get_top_logprobs_chunk(
chunk_logprobs,
top_k_nums,
pruned_lens,
top_logprobs_val,
top_logprobs_idx,
split_len_topk,
log_normalizer=chunk_log_normalizer,
precomputed_topk=chunk_precomputed_topk,
)
# Get the logprob of given token id
if logits_metadata.extend_token_ids_logprob:
token_ids_logprobs = logits_metadata.token_ids_logprobs[chunk_slice]
pruned_lens = logits_metadata.extend_logprob_pruned_lens_cpu[
chunk_slice
]
split_len_token_ids = get_token_ids_logprobs_chunk(
chunk_logprobs,
token_ids_logprobs,
pruned_lens,
token_ids_logprobs_val,
token_ids_logprobs_idx,
split_len_token_ids,
log_normalizer=chunk_log_normalizer,
)
# Get the logprob of the requested token ids
chunk_token_logprobs = chunk_logprobs[
torch.arange(chunk_logprobs.shape[0], device=chunk_logprobs.device),
logits_metadata.extend_input_logprob_token_ids_gpu[mask_indices],
]
if chunk_log_normalizer is not None:
row_max, row_log_sum = chunk_log_normalizer
chunk_token_logprobs = (
chunk_token_logprobs.float() - row_max
) - row_log_sum
token_logprobs.append(chunk_token_logprobs)
# Free before the next chunk's logits (bf16 + fp32) materialize.
del chunk_logprobs
if chunk_log_normalizer is not None:
row_max, row_log_sum = chunk_log_normalizer
chunk_token_logprobs = (
chunk_token_logprobs.float() - row_max
) - row_log_sum
if borrow_logprob_memory:
token_logprobs[lp_lo:lp_hi].copy_(
chunk_token_logprobs, non_blocking=True
)
else:
token_logprobs.append(chunk_token_logprobs)
# Free before the next chunk's logits (bf16 + fp32) materialize.
del chunk_logprobs
# Restore the full-pruned lm_head batch_info after chunk iteration.
if num_chunks > 1 and hasattr(lm_head, "reset_lm_head_pass"):
@@ -730,8 +793,14 @@ class InputLogprobProcessor:
)
lm_head.reset_lm_head_pass()
# Concatenate the results
token_logprobs = torch.cat(token_logprobs, dim=0)
input_copy_done = None
if borrow_logprob_memory:
# The copies and the next replay share the forward stream.
input_copy_done = torch.cuda.Event()
input_copy_done.record(torch.cuda.current_stream(pruned_states.device))
else:
# Concatenate the results
token_logprobs = torch.cat(token_logprobs, dim=0)
return (
LogprobResult(
@@ -740,6 +809,7 @@ class InputLogprobProcessor:
top_logprobs_idx=top_logprobs_idx,
token_ids_logprobs_val=token_ids_logprobs_val,
token_ids_logprobs_idx=token_ids_logprobs_idx,
input_copy_done=input_copy_done,
),
sampled_logits,
)
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+1
View File
@@ -519,6 +519,7 @@ class SchedulerBatchResultProcessor:
logits_output: LogitsProcessorOutput,
) -> None:
if batch.return_logprob:
logits_output.finalize_input_logprobs()
if logits_output.next_token_logprobs is not None:
logits_output.next_token_logprobs = (
logits_output.next_token_logprobs.tolist()
python/sglang/srt/managers/utils.py
+4 -13
View File
@@ -21,6 +21,7 @@ from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
from sglang.srt.runtime_context import get_spec, max_speculative_num_draft_tokens
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.state_capturer.base import TopkCaptureOutput
from sglang.srt.utils.common import async_d2h as _async_d2h
if TYPE_CHECKING:
from sglang.srt.managers.auxiliary_output import HostAuxiliaryOutput
@@ -45,19 +46,6 @@ def allocate_distinct_stream(device_module, avoid_streams):
raise RuntimeError("Unable to allocate a distinct stream")
def _async_d2h(t: torch.Tensor) -> torch.Tensor:
"""Async D2H copy for overlap scheduling. On CUDA the dest is pinned (a D2H
to pageable host memory blocks the caller until done) and record_stream keeps
the source alive until the copy stream drains, so the caching allocator can't
recycle it early. Non-CUDA falls back to a plain copy."""
if not t.is_cuda:
return t.to("cpu", non_blocking=True)
cpu_t = torch.empty(t.shape, dtype=t.dtype, pin_memory=True)
cpu_t.copy_(t, non_blocking=True)
t.record_stream(torch.cuda.current_stream(t.device))
return cpu_t
@dataclasses.dataclass
class GenerationBatchResult:
logits_output: Optional[LogitsProcessorOutput] = None
@@ -292,6 +280,9 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict:
logits_output = result.logits_output
assert logits_output is not None
# Nested pinned CPU tensors are serialized as Python metadata by PP, so
# their forward-stream copies must be complete before pickling starts.
logits_output.finalize_input_logprobs()
sampling_mask_output = logits_output.sampling_mask_output
return {
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+1
View File
@@ -2052,6 +2052,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
input_top_logprobs_idx=output.input_top_logprobs_idx,
input_token_ids_logprobs_val=output.input_token_ids_logprobs_val,
input_token_ids_logprobs_idx=output.input_token_ids_logprobs_idx,
input_logprobs_copy_done=output.input_logprobs_copy_done,
mm_input_embeds=mm_input_embeds,
)
python/sglang/srt/utils/common.py
+10
View File
@@ -558,6 +558,16 @@ def is_pin_memory_available(device=None) -> bool:
return current_platform.is_pin_memory_available(device)
def async_d2h(tensor: torch.Tensor) -> torch.Tensor:
"""Enqueue a CUDA-to-pinned-host copy on the current stream."""
if not tensor.is_cuda:
return tensor.to("cpu", non_blocking=True)
host = torch.empty(tensor.shape, dtype=tensor.dtype, pin_memory=True)
host.copy_(tensor, non_blocking=True)
tensor.record_stream(torch.cuda.current_stream(tensor.device))
return host
def get_dispatch_device_backend():
if is_cuda_alike():
dispatch_key = "CUDA"